Map.put returns old value, ugly solution

6 files changed, 35 insertions, 19 deletions

diff --git a/model-data/src/main/java/org/eclipse/viatra/solver/data/map/VersionedMap.java b/model-data/src/main/java/org/eclipse/viatra/solver/data/map/VersionedMap.java
index 2f2cb7a2..3a35b9f0 100644
--- a/model-data/src/main/java/org/eclipse/viatra/solver/data/map/VersionedMap.java
+++ b/model-data/src/main/java/org/eclipse/viatra/solver/data/map/VersionedMap.java
@@ -4,7 +4,7 @@ public interface VersionedMap<K,V> extends Versioned{
         public V get(K key);
         public Cursor<K,V> getAll();
         
-        public void put(K key, V value);
+        public V put(K key, V value);
         public void putAll(Cursor<K,V> cursor);
         
         public long getSize();
diff --git a/model-data/src/main/java/org/eclipse/viatra/solver/data/map/VersionedMapStoreImpl.java b/model-data/src/main/java/org/eclipse/viatra/solver/data/map/VersionedMapStoreImpl.java
index 4f3d8b10..83d0e8cd 100644
--- a/model-data/src/main/java/org/eclipse/viatra/solver/data/map/VersionedMapStoreImpl.java
+++ b/model-data/src/main/java/org/eclipse/viatra/solver/data/map/VersionedMapStoreImpl.java
@@ -4,6 +4,7 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -79,7 +80,7 @@ public class VersionedMapStoreImpl<K, V> implements VersionedMapStore<K, V> {
         
         @Override
         public synchronized Set<Long> getStates() {
-                return states.keySet();
+                return new HashSet<>(states.keySet());
         }
 
         @Override
diff --git a/model-data/src/main/java/org/eclipse/viatra/solver/data/map/internal/ImmutableNode.java b/model-data/src/main/java/org/eclipse/viatra/solver/data/map/internal/ImmutableNode.java
index 99df4d48..04a9b19a 100644
--- a/model-data/src/main/java/org/eclipse/viatra/solver/data/map/internal/ImmutableNode.java
+++ b/model-data/src/main/java/org/eclipse/viatra/solver/data/map/internal/ImmutableNode.java
@@ -129,7 +129,7 @@ public class ImmutableNode<K, V> extends Node<K, V> {
         
         @SuppressWarnings("unchecked")
         @Override
-        public Node<K,V> putValue(K key, V value, ContinousHashProvider<? super K> hashProvider, V defaultValue, int hash, int depth) {
+        public Node<K,V> putValue(K key, V value, OldValueBox<V> oldValue, ContinousHashProvider<? super K> hashProvider, V defaultValue, int hash, int depth) {
                 int selectedHashFragment = hashFragment(hash,shiftDepth(depth));
                 int bitposition = 1 << selectedHashFragment;
                 if((dataMap & bitposition) != 0) {
@@ -139,23 +139,25 @@ public class ImmutableNode<K, V> extends Node<K, V> {
                                 if(value == defaultValue) {
                                         // delete
                                         MutableNode<K, V> mutable = this.toMutable();
-                                        return mutable.removeEntry(selectedHashFragment);
+                                        return mutable.removeEntry(selectedHashFragment,oldValue);
                                 } else if(value == content[keyIndex+1]) {
                                         // dont change
+                                        oldValue.setOldValue(value);
                                         return this;
                                 } else {
                                         // update existing value
                                         MutableNode<K, V> mutable = this.toMutable();
-                                        return mutable.updateValue(value, selectedHashFragment);
+                                        return mutable.updateValue(value, oldValue, selectedHashFragment);
                                 }
                         } else {
                                 if(value == defaultValue) {
                                         // dont change
+                                        oldValue.setOldValue(defaultValue);
                                         return this;
                                 } else {
                                         // add new key + value
                                         MutableNode<K, V> mutable = this.toMutable();
-                                        return mutable.putValue(key, value, hashProvider, defaultValue, hash, depth);
+                                        return mutable.putValue(key, value, oldValue, hashProvider, defaultValue, hash, depth);
                                 }
                         }
                 } else if((nodeMap & bitposition)!=0) {
@@ -164,7 +166,7 @@ public class ImmutableNode<K, V> extends Node<K, V> {
                         ImmutableNode<K,V> subNode = (ImmutableNode<K,V>) content[keyIndex];
                         int newDepth = depth+1;
                         int newHash = newHash(hashProvider, key, hash, newDepth);
-                        Node<K,V> newsubNode = subNode.putValue(key, value, hashProvider, defaultValue, newHash, newDepth);
+                        Node<K,V> newsubNode = subNode.putValue(key, value, oldValue, hashProvider, defaultValue, newHash, newDepth);
                         
                         if(subNode == newsubNode) {
                                 // nothing changed
@@ -176,7 +178,7 @@ public class ImmutableNode<K, V> extends Node<K, V> {
                 } else {
                         // add new key + value
                         MutableNode<K, V> mutable = this.toMutable();
-                        return mutable.putValue(key, value, hashProvider, defaultValue, hash, depth);
+                        return mutable.putValue(key, value, oldValue, hashProvider, defaultValue, hash, depth);
                 }
         }
         
diff --git a/model-data/src/main/java/org/eclipse/viatra/solver/data/map/internal/MutableNode.java b/model-data/src/main/java/org/eclipse/viatra/solver/data/map/internal/MutableNode.java
index 2f794a7b..9e63b941 100644
--- a/model-data/src/main/java/org/eclipse/viatra/solver/data/map/internal/MutableNode.java
+++ b/model-data/src/main/java/org/eclipse/viatra/solver/data/map/internal/MutableNode.java
@@ -77,7 +77,7 @@ public class MutableNode<K, V> extends Node<K, V> {
 
         @SuppressWarnings("unchecked")
         @Override
-        public Node<K, V> putValue(K key, V value, ContinousHashProvider<? super K> hashProvider, V defaultValue, int hash,
+        public Node<K, V> putValue(K key, V value, OldValueBox<V> oldValue, ContinousHashProvider<? super K> hashProvider, V defaultValue, int hash,
                         int depth) {
                 int selectedHashFragment = hashFragment(hash, shiftDepth(depth));
                 K keyCandidate = (K) content[2 * selectedHashFragment];
@@ -86,18 +86,20 @@ public class MutableNode<K, V> extends Node<K, V> {
                         if (keyCandidate.equals(key)) {
                                 // The key is equals to an existing key -> update entry
                                 if (value == defaultValue) {
-                                        return removeEntry(selectedHashFragment);
+                                        return removeEntry(selectedHashFragment, oldValue);
                                 } else {
-                                        return updateValue(value, selectedHashFragment);
+                                        return updateValue(value, oldValue, selectedHashFragment);
                                 }
                         } else {
                                 // The key is not equivalent to an existing key on the same hash bin
                                 // -> split entry if it is necessary
                                 if (value == defaultValue) {
                                         // Value is default -> do not need to add new node
+                                        oldValue.setOldValue(defaultValue);
                                         return this;
                                 } else {
                                         // Value is not default -> Split entry data to a new node
+                                        oldValue.setOldValue(defaultValue);
                                         return moveDownAndSplit(hashProvider, key, value, keyCandidate, hash, depth, selectedHashFragment);
                                 }
                         }
@@ -106,24 +108,27 @@ public class MutableNode<K, V> extends Node<K, V> {
                         Node<K, V> nodeCandidate = (Node<K, V>) content[2 * selectedHashFragment + 1];
                         if (nodeCandidate != null) {
                                 // If it has value, it is a subnode -> upate that
-                                Node<K, V> newNode = nodeCandidate.putValue(key, value, hashProvider, defaultValue,
+                                Node<K, V> newNode = nodeCandidate.putValue(key, value, oldValue, hashProvider, defaultValue,
                                                 newHash(hashProvider, key, hash, depth + 1), depth + 1);
                                 return updateWithSubNode(selectedHashFragment, newNode, value.equals(defaultValue));
                         } else {
                                 // If it does not have value, put it in the empty place
                                 if (value == defaultValue) {
                                         // dont need to add new key-value pair
+                                        oldValue.setOldValue(defaultValue);
                                         return this;
                                 } else {
-                                        return addEntry(key, value, selectedHashFragment);
+                                        return addEntry(key, value, oldValue, selectedHashFragment);
                                 }
 
                         }
                 }
         }
 
-        private Node<K, V> addEntry(K key, V value, int selectedHashFragment) {
+        @SuppressWarnings("unchecked")
+        private Node<K, V> addEntry(K key, V value, OldValueBox<V> oldValue, int selectedHashFragment) {
                 content[2 * selectedHashFragment] = key;
+                oldValue.setOldValue((V) content[2 * selectedHashFragment + 1]);
                 content[2 * selectedHashFragment + 1] = value;
                 updateHash();
                 return this;
@@ -136,7 +141,9 @@ public class MutableNode<K, V> extends Node<K, V> {
          * @param selectedHashFragment
          * @return
          */
-        Node<K, V> updateValue(V value, int selectedHashFragment) {
+        @SuppressWarnings("unchecked")
+        Node<K, V> updateValue(V value, OldValueBox<V> oldValue, int selectedHashFragment) {
+                oldValue.setOldValue((V) content[2 * selectedHashFragment + 1]);
                 content[2 * selectedHashFragment + 1] = value;
                 updateHash();
                 return this;
@@ -247,8 +254,10 @@ public class MutableNode<K, V> extends Node<K, V> {
                 return subNode;
         }
 
-        Node<K, V> removeEntry(int selectedHashFragment) {
+        @SuppressWarnings("unchecked")
+        Node<K, V> removeEntry(int selectedHashFragment, OldValueBox<V> oldValue) {
                 content[2 * selectedHashFragment] = null;
+                oldValue.setOldValue((V) content[2 * selectedHashFragment + 1]);
                 content[2 * selectedHashFragment + 1] = null;
                 if (hasContent()) {
                         updateHash();
diff --git a/model-data/src/main/java/org/eclipse/viatra/solver/data/map/internal/Node.java b/model-data/src/main/java/org/eclipse/viatra/solver/data/map/internal/Node.java
index acf7a463..d40f980a 100644
--- a/model-data/src/main/java/org/eclipse/viatra/solver/data/map/internal/Node.java
+++ b/model-data/src/main/java/org/eclipse/viatra/solver/data/map/internal/Node.java
@@ -58,7 +58,7 @@ public abstract class Node<K,V>{
         
         
         public abstract V getValue(K key, ContinousHashProvider<? super K> hashProvider, V defaultValue,  int hash, int depth);
-        public abstract Node<K,V> putValue(K key, V value, ContinousHashProvider<? super K> hashProvider, V defaultValue, int hash, int depth);
+        public abstract Node<K,V> putValue(K key, V value, OldValueBox<V> old, ContinousHashProvider<? super K> hashProvider, V defaultValue, int hash, int depth);
         public abstract long getSize();
         
         abstract MutableNode<K, V> toMutable();
diff --git a/model-data/src/main/java/org/eclipse/viatra/solver/data/map/internal/VersionedMapImpl.java b/model-data/src/main/java/org/eclipse/viatra/solver/data/map/internal/VersionedMapImpl.java
index ef962f13..de41e602 100644
--- a/model-data/src/main/java/org/eclipse/viatra/solver/data/map/internal/VersionedMapImpl.java
+++ b/model-data/src/main/java/org/eclipse/viatra/solver/data/map/internal/VersionedMapImpl.java
@@ -24,6 +24,8 @@ public class VersionedMapImpl<K,V> implements VersionedMap<K,V>{
         protected final V defaultValue;
         protected Node<K,V> root;
         
+        private OldValueBox<V> oldValueBox = new OldValueBox<>();
+        
         public VersionedMapImpl(
                         VersionedMapStoreImpl<K,V> store,
                         ContinousHashProvider<K> hashProvider,
@@ -52,11 +54,13 @@ public class VersionedMapImpl<K,V> implements VersionedMap<K,V>{
                 return hashProvider;
         }
         @Override
-        public void put(K key, V value) {
+        public V put(K key, V value) {
                 if(root!=null) {
-                        root = root.putValue(key, value, hashProvider, defaultValue, hashProvider.getHash(key, 0), 0);
+                        root = root.putValue(key, value, oldValueBox, hashProvider, defaultValue, hashProvider.getHash(key, 0), 0);
+                        return oldValueBox.getOldValue();
                 } else {
                         root = MutableNode.initialize(key, value, hashProvider, defaultValue);
+                        return defaultValue;
                 }
         }