4 files changed, 27 insertions, 14 deletions

diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/viewupdate/ModelUpdateListener.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/viewupdate/ModelUpdateListener.java
index 25919888..5a0da315 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/viewupdate/ModelUpdateListener.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/viewupdate/ModelUpdateListener.java
@@ -1,5 +1,6 @@
 package tools.refinery.store.query.viatra.internal.viewupdate;
 
+import org.eclipse.viatra.query.runtime.matchers.context.IInputKey;
 import org.eclipse.viatra.query.runtime.matchers.context.IQueryRuntimeContextListener;
 import org.eclipse.viatra.query.runtime.matchers.tuple.ITuple;
 import tools.refinery.store.model.Tuple;
@@ -46,20 +47,23 @@ public class ModelUpdateListener {
                 return view2Buffers.containsKey(relationalKey);
         }
 
-        public <D> void addListener(RelationView<D> relationView, ITuple seed, IQueryRuntimeContextListener listener) {
+        public <D> void addListener(IInputKey key, RelationView<D> relationView, ITuple seed,
+                                                                IQueryRuntimeContextListener listener) {
                 if (view2Buffers.containsKey(relationView)) {
-                        ViewUpdateTranslator<D> updateListener = new ViewUpdateTranslator<>(relationView, seed, listener);
+                        ViewUpdateTranslator<D> updateListener = new ViewUpdateTranslator<>(key, relationView, seed, listener);
                         ViewUpdateBuffer<D> updateBuffer = new ViewUpdateBuffer<>(updateListener);
                         view2Buffers.get(relationView).add(updateBuffer);
-                } else
+                } else {
                         throw new IllegalArgumentException();
+                }
         }
 
-        public void removeListener(RelationView<?> relationView, ITuple seed, IQueryRuntimeContextListener listener) {
+        public void removeListener(IInputKey key, RelationView<?> relationView, ITuple seed,
+                                                           IQueryRuntimeContextListener listener) {
                 if (view2Buffers.containsKey(relationView)) {
                         Set<ViewUpdateBuffer<?>> buffers = this.view2Buffers.get(relationView);
                         for (var buffer : buffers) {
-                                if (buffer.getUpdateListener().equals(relationView, seed, listener)) {
+                                if (buffer.getUpdateListener().equals(key, relationView, seed, listener)) {
                                         // remove buffer and terminate immediately, or it will break iterator.
                                         buffers.remove(buffer);
                                         return;
diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/viewupdate/ViewUpdate.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/viewupdate/ViewUpdate.java
index c727f046..b9406018 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/viewupdate/ViewUpdate.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/viewupdate/ViewUpdate.java
@@ -3,8 +3,7 @@ package tools.refinery.store.query.viatra.internal.viewupdate;
 import java.util.Arrays;
 import java.util.Objects;
 
-record ViewUpdate (Object[] tuple, boolean isInsertion) {
-
+record ViewUpdate(Object[] tuple, boolean isInsertion) {
         @Override
         public int hashCode() {
                 final int prime = 31;
@@ -28,7 +27,6 @@ record ViewUpdate (Object[] tuple, boolean isInsertion) {
 
         @Override
         public String toString() {
-                return "ViewUpdate [" + Arrays.toString(tuple) + "insertion= "+this.isInsertion+"]";
+                return "ViewUpdate [" + Arrays.toString(tuple) + "insertion= " + this.isInsertion + "]";
         }
-
 }
diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/viewupdate/ViewUpdateBuffer.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/viewupdate/ViewUpdateBuffer.java
index 5a4243f2..e13a9cb8 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/viewupdate/ViewUpdateBuffer.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/viewupdate/ViewUpdateBuffer.java
@@ -8,6 +8,7 @@ import java.util.List;
 
 public class ViewUpdateBuffer<D> {
         protected final ViewUpdateTranslator<D> updateListener;
+
         protected final List<ViewUpdate> buffer = new ArrayList<>();
 
         public ViewUpdateBuffer(ViewUpdateTranslator<D> updateListener) {
diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/viewupdate/ViewUpdateTranslator.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/viewupdate/ViewUpdateTranslator.java
index 2f7f9a9c..62a10e45 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/viewupdate/ViewUpdateTranslator.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/viewupdate/ViewUpdateTranslator.java
@@ -1,5 +1,6 @@
 package tools.refinery.store.query.viatra.internal.viewupdate;
 
+import org.eclipse.viatra.query.runtime.matchers.context.IInputKey;
 import org.eclipse.viatra.query.runtime.matchers.context.IQueryRuntimeContextListener;
 import org.eclipse.viatra.query.runtime.matchers.tuple.ITuple;
 import org.eclipse.viatra.query.runtime.matchers.tuple.Tuples;
@@ -9,29 +10,38 @@ import tools.refinery.store.query.view.RelationView;
 import java.util.Objects;
 
 public class ViewUpdateTranslator<D> {
+        private final IInputKey wrappedKey;
+
         private final RelationView<D> key;
 
         private final ITuple filter;
 
         private final IQueryRuntimeContextListener listener;
 
-        public ViewUpdateTranslator(RelationView<D> key, ITuple filter, IQueryRuntimeContextListener listener) {
+        public ViewUpdateTranslator(IInputKey wrappedKey, RelationView<D> key, ITuple filter,
+                                                                IQueryRuntimeContextListener listener) {
                 super();
+                this.wrappedKey = wrappedKey;
                 this.key = key;
                 this.filter = filter;
                 this.listener = listener;
         }
 
-        public boolean equals(RelationView<?> relationView, ITuple seed, IQueryRuntimeContextListener listener) {
-                return key == relationView && filter.equals(seed) && this.listener == listener;
+        public boolean equals(IInputKey wrappedKey, RelationView<?> relationView, ITuple seed,
+                                                  IQueryRuntimeContextListener listener) {
+                return this.wrappedKey == wrappedKey && key == relationView && filter.equals(seed) && this.listener == listener;
         }
 
         public void processChange(ViewUpdate change) {
-                listener.update(key, Tuples.flatTupleOf(change.tuple()), change.isInsertion());
+                listener.update(wrappedKey, Tuples.flatTupleOf(change.tuple()), change.isInsertion());
         }
 
+        @SuppressWarnings("squid:S1168")
         public Object[] isMatching(Tuple tuple, D value) {
-                return isMatching(key.getWrappedKey().transform(tuple, value), filter);
+                if (!key.filter(tuple, value)) {
+                        return null;
+                }
+                return isMatching(key.forwardMap(tuple, value), filter);
         }
 
         @SuppressWarnings("squid:S1168")