feat: model query functional dependencies

11 files changed, 115 insertions, 36 deletions

diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/ViatraModelQueryBuilder.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/ViatraModelQueryBuilder.java
index ee445a79..efc6146c 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/ViatraModelQueryBuilder.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/ViatraModelQueryBuilder.java
@@ -29,7 +29,7 @@ public interface ViatraModelQueryBuilder extends ModelQueryBuilder {
         }
 
         @Override
-        default ViatraModelQueryBuilder queries(Collection<? extends DNF> queries) {
+        default ViatraModelQueryBuilder queries(Collection<DNF> queries) {
                 ModelQueryBuilder.super.queries(queries);
                 return this;
         }
diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/RelationalScope.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/RelationalScope.java
index 21dcaf15..b4d43ef8 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/RelationalScope.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/RelationalScope.java
@@ -5,26 +5,25 @@ import org.eclipse.viatra.query.runtime.api.ViatraQueryEngine;
 import org.eclipse.viatra.query.runtime.api.scope.IEngineContext;
 import org.eclipse.viatra.query.runtime.api.scope.IIndexingErrorListener;
 import org.eclipse.viatra.query.runtime.api.scope.QueryScope;
+import org.eclipse.viatra.query.runtime.matchers.context.IInputKey;
 import tools.refinery.store.model.Model;
 import tools.refinery.store.query.viatra.internal.context.RelationalEngineContext;
-import tools.refinery.store.query.viatra.internal.update.ModelUpdateListener;
 import tools.refinery.store.query.view.AnyRelationView;
 
-import java.util.Collection;
+import java.util.Map;
 
 public class RelationalScope extends QueryScope {
         private final Model model;
+        private final Map<AnyRelationView, IInputKey> relationViews;
 
-        private final ModelUpdateListener updateListener;
-
-        public RelationalScope(Model model, Collection<AnyRelationView> relationViews) {
+        public RelationalScope(Model model, Map<AnyRelationView, IInputKey> relationViews) {
                 this.model = model;
-                updateListener = new ModelUpdateListener(model, relationViews);
+                this.relationViews = relationViews;
         }
 
         @Override
         protected IEngineContext createEngineContext(ViatraQueryEngine engine, IIndexingErrorListener errorListener,
                                                                                                  Logger logger) {
-                return new RelationalEngineContext(model, updateListener);
+                return new RelationalEngineContext(model, relationViews);
         }
 }
diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryAdapterImpl.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryAdapterImpl.java
index 3c276935..810d2c32 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryAdapterImpl.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryAdapterImpl.java
@@ -30,7 +30,7 @@ public class ViatraModelQueryAdapterImpl implements ModelQueryAdapter {
         ViatraModelQueryAdapterImpl(Model model, ViatraModelQueryStoreAdapterImpl storeAdapter) {
                 this.model = model;
                 this.storeAdapter = storeAdapter;
-                var scope = new RelationalScope(model, storeAdapter.getRelationViews());
+                var scope = new RelationalScope(model, storeAdapter.getInputKeys());
                 queryEngine = (ViatraQueryEngineImpl) AdvancedViatraQueryEngine.createUnmanagedEngine(scope);
 
                 try {
diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryBuilderImpl.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryBuilderImpl.java
index 5105c9a7..9f1e55b1 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryBuilderImpl.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryBuilderImpl.java
@@ -92,7 +92,7 @@ public class ViatraModelQueryBuilderImpl extends AbstractModelAdapterBuilder imp
                         throw new IllegalArgumentException(
                                         "Cannot specify hint for %s, because it was not added to the query engine".formatted(dnf.name()));
                 }
-                pQuery.setEvaluationHints(queryEvaluationHint);
+                pQuery.setEvaluationHints(pQuery.getEvaluationHints().overrideBy(queryEvaluationHint));
                 return this;
         }
 
@@ -105,7 +105,7 @@ public class ViatraModelQueryBuilderImpl extends AbstractModelAdapterBuilder imp
 
         private void validateSymbols(ModelStore store) {
                 var symbols = store.getSymbols();
-                for (var relationView : dnf2PQuery.getRelationViews()) {
+                for (var relationView : dnf2PQuery.getRelationViews().keySet()) {
                         var symbol = relationView.getSymbol();
                         if (!symbols.contains(symbol)) {
                                 throw new IllegalArgumentException("Cannot query relation view %s: symbol %s is not in the model"
diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryStoreAdapterImpl.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryStoreAdapterImpl.java
index d77b7f4b..69f1f146 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryStoreAdapterImpl.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/ViatraModelQueryStoreAdapterImpl.java
@@ -2,6 +2,7 @@ package tools.refinery.store.query.viatra.internal;
 
 import org.eclipse.viatra.query.runtime.api.IQuerySpecification;
 import org.eclipse.viatra.query.runtime.api.ViatraQueryEngineOptions;
+import org.eclipse.viatra.query.runtime.matchers.context.IInputKey;
 import tools.refinery.store.model.Model;
 import tools.refinery.store.model.ModelStore;
 import tools.refinery.store.query.DNF;
@@ -15,15 +16,15 @@ import java.util.Map;
 public class ViatraModelQueryStoreAdapterImpl implements ViatraModelQueryStoreAdapter {
         private final ModelStore store;
         private final ViatraQueryEngineOptions engineOptions;
-        private final Collection<AnyRelationView> relationViews;
+        private final Map<AnyRelationView, IInputKey> inputKeys;
         private final Map<DNF, IQuerySpecification<RawPatternMatcher>> querySpecifications;
 
         ViatraModelQueryStoreAdapterImpl(ModelStore store, ViatraQueryEngineOptions engineOptions,
-                                                                         Collection<AnyRelationView> relationViews,
+                                                                         Map<AnyRelationView, IInputKey> inputKeys,
                                                                          Map<DNF, IQuerySpecification<RawPatternMatcher>> querySpecifications) {
                 this.store = store;
                 this.engineOptions = engineOptions;
-                this.relationViews = relationViews;
+                this.inputKeys = inputKeys;
                 this.querySpecifications = querySpecifications;
         }
 
@@ -32,9 +33,12 @@ public class ViatraModelQueryStoreAdapterImpl implements ViatraModelQueryStoreAd
                 return store;
         }
 
-        @Override
         public Collection<AnyRelationView> getRelationViews() {
-                return relationViews;
+                return inputKeys.keySet();
+        }
+
+        Map<AnyRelationView, IInputKey> getInputKeys() {
+                return inputKeys;
         }
 
         @Override
diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/context/DummyBaseIndexer.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/context/DummyBaseIndexer.java
index 29f0536c..2a24b67c 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/context/DummyBaseIndexer.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/context/DummyBaseIndexer.java
@@ -12,6 +12,9 @@ import java.util.concurrent.Callable;
  * Copied from <code>org.eclipse.viatra.query.runtime.tabular.TabularEngineContext</code>
  */
 public class DummyBaseIndexer implements IBaseIndex {
+        DummyBaseIndexer() {
+        }
+
         @Override
         public <V> V coalesceTraversals(Callable<V> callable) throws InvocationTargetException {
                 try {
diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/context/RelationalEngineContext.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/context/RelationalEngineContext.java
index 4eb8898b..3bad01b9 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/context/RelationalEngineContext.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/context/RelationalEngineContext.java
@@ -2,17 +2,19 @@ package tools.refinery.store.query.viatra.internal.context;
 
 import org.eclipse.viatra.query.runtime.api.scope.IBaseIndex;
 import org.eclipse.viatra.query.runtime.api.scope.IEngineContext;
+import org.eclipse.viatra.query.runtime.matchers.context.IInputKey;
 import org.eclipse.viatra.query.runtime.matchers.context.IQueryRuntimeContext;
-
 import tools.refinery.store.model.Model;
-import tools.refinery.store.query.viatra.internal.update.ModelUpdateListener;
+import tools.refinery.store.query.view.AnyRelationView;
+
+import java.util.Map;
 
 public class RelationalEngineContext implements IEngineContext {
         private final IBaseIndex baseIndex = new DummyBaseIndexer();
         private final RelationalRuntimeContext runtimeContext;
 
-        public RelationalEngineContext(Model model, ModelUpdateListener updateListener) {
-                runtimeContext = new RelationalRuntimeContext(model, updateListener);
+        public RelationalEngineContext(Model model, Map< AnyRelationView, IInputKey> inputKeys) {
+                runtimeContext = new RelationalRuntimeContext(model, inputKeys);
         }
 
         @Override
diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/context/RelationalQueryMetaContext.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/context/RelationalQueryMetaContext.java
index 81046b2a..cba3fa08 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/context/RelationalQueryMetaContext.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/context/RelationalQueryMetaContext.java
@@ -4,24 +4,29 @@ import org.eclipse.viatra.query.runtime.matchers.context.AbstractQueryMetaContex
 import org.eclipse.viatra.query.runtime.matchers.context.IInputKey;
 import org.eclipse.viatra.query.runtime.matchers.context.InputKeyImplication;
 import tools.refinery.store.query.viatra.internal.pquery.RelationViewWrapper;
+import tools.refinery.store.query.view.AnyRelationView;
 
-import java.util.Collection;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 
 /**
  * The meta context information for String scopes.
  */
 public class RelationalQueryMetaContext extends AbstractQueryMetaContext {
+        private final Map<AnyRelationView, IInputKey> inputKeys;
+
+        RelationalQueryMetaContext(Map<AnyRelationView, IInputKey> inputKeys) {
+                this.inputKeys = inputKeys;
+        }
+
         @Override
         public boolean isEnumerable(IInputKey key) {
-                ensureValidKey(key);
+                checkKey(key);
                 return key.isEnumerable();
         }
 
         @Override
         public boolean isStateless(IInputKey key) {
-                ensureValidKey(key);
+                checkKey(key);
                 return true;
         }
 
@@ -32,19 +37,58 @@ public class RelationalQueryMetaContext extends AbstractQueryMetaContext {
 
         @Override
         public Collection<InputKeyImplication> getImplications(IInputKey implyingKey) {
-                ensureValidKey(implyingKey);
-                return Set.of();
+                var relationView = checkKey(implyingKey);
+                var relationViewImplications = relationView.getImpliedRelationViews();
+                var inputKeyImplications = new HashSet<InputKeyImplication>(relationViewImplications.size());
+                for (var relationViewImplication : relationViewImplications) {
+                        if (!relationView.equals(relationViewImplication.implyingRelationView())) {
+                                throw new IllegalArgumentException("Relation view %s returned unrelated implication %s".formatted(
+                                                relationView, relationViewImplication));
+                        }
+                        var impliedInputKey = inputKeys.get(relationViewImplication.impliedRelationView());
+                        // Ignore implications not relevant for any queries included in the model.
+                        if (impliedInputKey != null) {
+                                inputKeyImplications.add(new InputKeyImplication(implyingKey, impliedInputKey,
+                                                relationViewImplication.impliedIndices()));
+                        }
+                }
+                return inputKeyImplications;
         }
 
         @Override
         public Map<Set<Integer>, Set<Integer>> getFunctionalDependencies(IInputKey key) {
-                ensureValidKey(key);
-                return Map.of();
+                var relationView = checkKey(key);
+                var functionalDependencies = relationView.getFunctionalDependencies();
+                var flattened = new HashMap<Set<Integer>, Set<Integer>>(functionalDependencies.size());
+                for (var functionalDependency : functionalDependencies) {
+                        var forEach = functionalDependency.forEach();
+                        checkValidIndices(relationView, forEach);
+                        var unique = functionalDependency.unique();
+                        checkValidIndices(relationView, unique);
+                        var existing = flattened.get(forEach);
+                        if (existing == null) {
+                                flattened.put(forEach, new HashSet<>(unique));
+                        } else {
+                                existing.addAll(unique);
+                        }
+                }
+                return flattened;
+        }
+
+        private static void checkValidIndices(AnyRelationView relationView, Collection<Integer> indices) {
+                indices.stream().filter(relationView::invalidIndex).findAny().ifPresent(i -> {
+                        throw new IllegalArgumentException("Index %d is invalid for %s".formatted(i, relationView));
+                });
         }
 
-        public void ensureValidKey(IInputKey key) {
-                if (!(key instanceof RelationViewWrapper)) {
+        public AnyRelationView checkKey(IInputKey key) {
+                if (!(key instanceof RelationViewWrapper wrapper)) {
                         throw new IllegalArgumentException("The input key %s is not a valid input key".formatted(key));
                 }
+                var relationView = wrapper.getWrappedKey();
+                if (!inputKeys.containsKey(relationView)) {
+                        throw new IllegalArgumentException("The input key %s is not present in the model".formatted(key));
+                }
+                return relationView;
         }
 }
diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/context/RelationalRuntimeContext.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/context/RelationalRuntimeContext.java
index 9c1d966c..71ab5cb4 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/context/RelationalRuntimeContext.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/context/RelationalRuntimeContext.java
@@ -13,6 +13,7 @@ import tools.refinery.store.query.view.AnyRelationView;
 
 import java.lang.reflect.InvocationTargetException;
 import java.util.Iterator;
+import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.Callable;
 
@@ -20,15 +21,16 @@ import static tools.refinery.store.util.CollectionsUtil.filter;
 import static tools.refinery.store.util.CollectionsUtil.map;
 
 public class RelationalRuntimeContext implements IQueryRuntimeContext {
-        private final RelationalQueryMetaContext metaContext = new RelationalQueryMetaContext();
+        private final RelationalQueryMetaContext metaContext;
 
         private final ModelUpdateListener modelUpdateListener;
 
         private final Model model;
 
-        public RelationalRuntimeContext(Model model, ModelUpdateListener relationUpdateListener) {
+        RelationalRuntimeContext(Model model, Map<AnyRelationView, IInputKey> inputKeys) {
                 this.model = model;
-                this.modelUpdateListener = relationUpdateListener;
+                metaContext = new RelationalQueryMetaContext(inputKeys);
+                modelUpdateListener = new ModelUpdateListener(model, inputKeys.keySet());
         }
 
         @Override
diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/pquery/DNF2PQuery.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/pquery/DNF2PQuery.java
index f2cb8d8d..60f1bcae 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/pquery/DNF2PQuery.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/pquery/DNF2PQuery.java
@@ -1,8 +1,10 @@
 package tools.refinery.store.query.viatra.internal.pquery;
 
 import org.eclipse.viatra.query.runtime.matchers.backend.QueryEvaluationHint;
+import org.eclipse.viatra.query.runtime.matchers.context.IInputKey;
 import org.eclipse.viatra.query.runtime.matchers.psystem.PBody;
 import org.eclipse.viatra.query.runtime.matchers.psystem.PVariable;
+import org.eclipse.viatra.query.runtime.matchers.psystem.annotations.PAnnotation;
 import org.eclipse.viatra.query.runtime.matchers.psystem.basicdeferred.Equality;
 import org.eclipse.viatra.query.runtime.matchers.psystem.basicdeferred.ExportedParameter;
 import org.eclipse.viatra.query.runtime.matchers.psystem.basicdeferred.Inequality;
@@ -35,6 +37,8 @@ public class DNF2PQuery {
 
         private final Map<AnyRelationView, RelationViewWrapper> view2WrapperMap = new LinkedHashMap<>();
 
+        private final Map<AnyRelationView, RawPQuery> view2EmbeddedMap = new HashMap<>();
+
         private Function<DNF, QueryEvaluationHint> computeHint = dnf -> new QueryEvaluationHint(null,
                         QueryEvaluationHint.BackendRequirement.UNSPECIFIED);
 
@@ -63,8 +67,8 @@ public class DNF2PQuery {
                 return pQuery;
         }
 
-        public Collection<AnyRelationView> getRelationViews() {
-                return Collections.unmodifiableCollection(view2WrapperMap.keySet());
+        public Map<AnyRelationView, IInputKey> getRelationViews() {
+                return Collections.unmodifiableMap(view2WrapperMap);
         }
 
         public RawPQuery getAlreadyTranslated(DNF dnfQuery) {
@@ -86,6 +90,17 @@ public class DNF2PQuery {
                 }
                 pQuery.setParameters(parameterList);
 
+                for (var functionalDependency : dnfQuery.getFunctionalDependencies()) {
+                        var functionalDependencyAnnotation = new PAnnotation("FunctionalDependency");
+                        for (var forEachVariable : functionalDependency.forEach()) {
+                                functionalDependencyAnnotation.addAttribute("forEach", forEachVariable.getUniqueName());
+                        }
+                        for (var uniqueVariable : functionalDependency.unique()) {
+                                functionalDependencyAnnotation.addAttribute("unique", uniqueVariable.getUniqueName());
+                        }
+                        pQuery.addAnnotation(functionalDependencyAnnotation);
+                }
+
                 // The constructor of {@link org.eclipse.viatra.query.runtime.matchers.psystem.BasePConstraint} mutates
                 // global static state (<code>nextID</code>) without locking. Therefore, we need to synchronize before creating
                 // any query constraints to avoid a data race.
@@ -159,6 +174,10 @@ public class DNF2PQuery {
         }
 
         private RawPQuery translateEmbeddedRelationViewPQuery(AnyRelationView relationView) {
+                return view2EmbeddedMap.computeIfAbsent(relationView, this::doTranslateEmbeddedRelationViewPQuery);
+        }
+
+        private RawPQuery doTranslateEmbeddedRelationViewPQuery(AnyRelationView relationView) {
                 var embeddedPQuery = new RawPQuery(DNFUtils.generateUniqueName(relationView.name()), PVisibility.EMBEDDED);
                 var body = new PBody(embeddedPQuery);
                 int arity = relationView.arity();
diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/pquery/RawPQuery.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/pquery/RawPQuery.java
index 830495b2..71b74396 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/pquery/RawPQuery.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/viatra/internal/pquery/RawPQuery.java
@@ -4,6 +4,7 @@ import org.eclipse.viatra.query.runtime.api.GenericQuerySpecification;
 import org.eclipse.viatra.query.runtime.api.ViatraQueryEngine;
 import org.eclipse.viatra.query.runtime.api.scope.QueryScope;
 import org.eclipse.viatra.query.runtime.matchers.psystem.PBody;
+import org.eclipse.viatra.query.runtime.matchers.psystem.annotations.PAnnotation;
 import org.eclipse.viatra.query.runtime.matchers.psystem.queries.BasePQuery;
 import org.eclipse.viatra.query.runtime.matchers.psystem.queries.PParameter;
 import org.eclipse.viatra.query.runtime.matchers.psystem.queries.PVisibility;
@@ -37,6 +38,11 @@ public class RawPQuery extends BasePQuery {
         }
 
         @Override
+        public void addAnnotation(PAnnotation annotation) {
+                super.addAnnotation(annotation);
+        }
+
+        @Override
         public List<PParameter> getParameters() {
                 return parameters;
         }