feat: query rewriters for Viatra

8 files changed, 113 insertions, 64 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 931a07aa..66279c94 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
@@ -9,9 +9,10 @@ import org.eclipse.viatra.query.runtime.api.ViatraQueryEngineOptions;
 import org.eclipse.viatra.query.runtime.matchers.backend.IQueryBackendFactory;
 import org.eclipse.viatra.query.runtime.matchers.backend.QueryEvaluationHint;
 import tools.refinery.store.model.ModelStore;
+import tools.refinery.store.query.ModelQueryBuilder;
 import tools.refinery.store.query.dnf.AnyQuery;
 import tools.refinery.store.query.dnf.Dnf;
-import tools.refinery.store.query.ModelQueryBuilder;
+import tools.refinery.store.query.rewriter.DnfRewriter;
 
 import java.util.Collection;
 import java.util.function.Function;
@@ -43,12 +44,11 @@ public interface ViatraModelQueryBuilder extends ModelQueryBuilder {
         @Override
         ViatraModelQueryBuilder query(AnyQuery query);
 
-        ViatraModelQueryBuilder query(AnyQuery query, QueryEvaluationHint queryEvaluationHint);
+        @Override
+        ViatraModelQueryBuilder rewriter(DnfRewriter rewriter);
 
         ViatraModelQueryBuilder computeHint(Function<Dnf, QueryEvaluationHint> computeHint);
 
-        ViatraModelQueryBuilder hint(Dnf dnf, QueryEvaluationHint queryEvaluationHint);
-
         @Override
         ViatraModelQueryStoreAdapter build(ModelStore store);
 }
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 5f3e86b4..e17386e1 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
@@ -131,7 +131,8 @@ public class ViatraModelQueryAdapterImpl implements ViatraModelQueryAdapter, Mod
 
         @Override
         public <T> ResultSet<T> getResultSet(Query<T> query) {
-                var resultSet = resultSets.get(query);
+                var canonicalQuery = storeAdapter.getCanonicalQuery(query);
+                var resultSet = resultSets.get(canonicalQuery);
                 if (resultSet == null) {
                         throw new IllegalArgumentException("No matcher for query %s in model".formatted(query.name()));
                 }
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 ce2467b4..5ee8bd74 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
@@ -13,9 +13,12 @@ import org.eclipse.viatra.query.runtime.matchers.backend.QueryEvaluationHint;
 import org.eclipse.viatra.query.runtime.rete.matcher.ReteBackendFactory;
 import tools.refinery.store.adapter.AbstractModelAdapterBuilder;
 import tools.refinery.store.model.ModelStore;
-import tools.refinery.store.model.ModelStoreBuilder;
 import tools.refinery.store.query.dnf.AnyQuery;
 import tools.refinery.store.query.dnf.Dnf;
+import tools.refinery.store.query.rewriter.CompositeRewriter;
+import tools.refinery.store.query.rewriter.DnfRewriter;
+import tools.refinery.store.query.rewriter.DuplicateDnfRemover;
+import tools.refinery.store.query.rewriter.InputParameterResolver;
 import tools.refinery.store.query.viatra.ViatraModelQueryBuilder;
 import tools.refinery.store.query.viatra.internal.localsearch.FlatCostFunction;
 import tools.refinery.store.query.viatra.internal.localsearch.RelationalLocalSearchBackendFactory;
@@ -32,15 +35,18 @@ public class ViatraModelQueryBuilderImpl extends AbstractModelAdapterBuilder<Via
                         // Use a cost function that ignores the initial (empty) model but allows higher arity input keys.
                         LocalSearchHintOptions.PLANNER_COST_FUNCTION, new FlatCostFunction()
         ), (IQueryBackendFactory) null);
+        private final CompositeRewriter rewriter;
         private final Dnf2PQuery dnf2PQuery = new Dnf2PQuery();
-        private final Set<AnyQuery> vacuousQueries = new LinkedHashSet<>();
-        private final Map<AnyQuery, IQuerySpecification<RawPatternMatcher>> querySpecifications = new LinkedHashMap<>();
+        private final Set<AnyQuery> queries = new LinkedHashSet<>();
 
         public ViatraModelQueryBuilderImpl() {
                 engineOptionsBuilder = new ViatraQueryEngineOptions.Builder()
                                 .withDefaultBackend(ReteBackendFactory.INSTANCE)
                                 .withDefaultCachingBackend(ReteBackendFactory.INSTANCE)
                                 .withDefaultSearchBackend(RelationalLocalSearchBackendFactory.INSTANCE);
+                rewriter = new CompositeRewriter();
+                rewriter.addFirst(new DuplicateDnfRemover());
+                rewriter.addFirst(new InputParameterResolver());
         }
 
         @Override
@@ -79,58 +85,58 @@ public class ViatraModelQueryBuilderImpl extends AbstractModelAdapterBuilder<Via
         }
 
         @Override
-        public ViatraModelQueryBuilder query(AnyQuery query) {
+        public ViatraModelQueryBuilder queries(Collection<? extends AnyQuery> queries) {
                 checkNotConfigured();
-                if (querySpecifications.containsKey(query) || vacuousQueries.contains(query)) {
-                        // Ignore duplicate queries.
-                        return this;
-                }
-                var dnf = query.getDnf();
-                var reduction = dnf.getReduction();
-                switch (reduction) {
-                case NOT_REDUCIBLE -> {
-                        var pQuery = dnf2PQuery.translate(dnf);
-                        querySpecifications.put(query, pQuery.build());
-                }
-                case ALWAYS_FALSE -> vacuousQueries.add(query);
-                case ALWAYS_TRUE -> throw new IllegalArgumentException(
-                                "Query %s is relationally unsafe (it matches every tuple)".formatted(query.name()));
-                default -> throw new IllegalArgumentException("Unknown reduction: " + reduction);
-                }
+                this.queries.addAll(queries);
                 return this;
         }
 
         @Override
-        public ViatraModelQueryBuilder query(AnyQuery query, QueryEvaluationHint queryEvaluationHint) {
-                hint(query.getDnf(), queryEvaluationHint);
-                query(query);
+        public ViatraModelQueryBuilder query(AnyQuery query) {
+                checkNotConfigured();
+                queries.add(query);
                 return this;
         }
 
         @Override
-        public ViatraModelQueryBuilder computeHint(Function<Dnf, QueryEvaluationHint> computeHint) {
-                checkNotConfigured();
-                dnf2PQuery.setComputeHint(computeHint);
+        public ViatraModelQueryBuilder rewriter(DnfRewriter rewriter) {
+                this.rewriter.addFirst(rewriter);
                 return this;
         }
 
         @Override
-        public ViatraModelQueryBuilder hint(Dnf dnf, QueryEvaluationHint queryEvaluationHint) {
+        public ViatraModelQueryBuilder computeHint(Function<Dnf, QueryEvaluationHint> computeHint) {
                 checkNotConfigured();
-                dnf2PQuery.hint(dnf, queryEvaluationHint);
+                dnf2PQuery.setComputeHint(computeHint);
                 return this;
         }
 
         @Override
-        public void doConfigure(ModelStoreBuilder storeBuilder) {
-                dnf2PQuery.assertNoUnusedHints();
-        }
-
-        @Override
         public ViatraModelQueryStoreAdapterImpl doBuild(ModelStore store) {
+                var canonicalQueryMap = new HashMap<AnyQuery, AnyQuery>();
+                var querySpecifications = new LinkedHashMap<AnyQuery, IQuerySpecification<RawPatternMatcher>>();
+                var vacuousQueries = new LinkedHashSet<AnyQuery>();
+                for (var query : queries) {
+                        var canonicalQuery = rewriter.rewrite(query);
+                        canonicalQueryMap.put(query, canonicalQuery);
+                        var dnf = canonicalQuery.getDnf();
+                        var reduction = dnf.getReduction();
+                        switch (reduction) {
+                        case NOT_REDUCIBLE -> {
+                                var pQuery = dnf2PQuery.translate(dnf);
+                                querySpecifications.put(canonicalQuery, pQuery.build());
+                        }
+                        case ALWAYS_FALSE -> vacuousQueries.add(canonicalQuery);
+                        case ALWAYS_TRUE -> throw new IllegalArgumentException(
+                                        "Query %s is relationally unsafe (it matches every tuple)".formatted(query.name()));
+                        default -> throw new IllegalArgumentException("Unknown reduction: " + reduction);
+                        }
+                }
+
                 validateSymbols(store);
                 return new ViatraModelQueryStoreAdapterImpl(store, buildEngineOptions(), dnf2PQuery.getSymbolViews(),
-                                Collections.unmodifiableMap(querySpecifications), Collections.unmodifiableSet(vacuousQueries));
+                                Collections.unmodifiableMap(canonicalQueryMap), Collections.unmodifiableMap(querySpecifications),
+                                Collections.unmodifiableSet(vacuousQueries));
         }
 
         private ViatraQueryEngineOptions buildEngineOptions() {
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 11a3c7fd..25f1cd02 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
@@ -11,6 +11,7 @@ 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.AnyQuery;
+import tools.refinery.store.query.dnf.Query;
 import tools.refinery.store.query.viatra.ViatraModelQueryStoreAdapter;
 import tools.refinery.store.query.viatra.internal.matcher.RawPatternMatcher;
 import tools.refinery.store.query.view.AnySymbolView;
@@ -21,17 +22,20 @@ public class ViatraModelQueryStoreAdapterImpl implements ViatraModelQueryStoreAd
         private final ModelStore store;
         private final ViatraQueryEngineOptions engineOptions;
         private final Map<AnySymbolView, IInputKey> inputKeys;
+        private final Map<AnyQuery, AnyQuery> canonicalQueryMap;
         private final Map<AnyQuery, IQuerySpecification<RawPatternMatcher>> querySpecifications;
         private final Set<AnyQuery> vacuousQueries;
         private final Set<AnyQuery> allQueries;
 
         ViatraModelQueryStoreAdapterImpl(ModelStore store, ViatraQueryEngineOptions engineOptions,
                                                                          Map<AnySymbolView, IInputKey> inputKeys,
+                                                                         Map<AnyQuery, AnyQuery> canonicalQueryMap,
                                                                          Map<AnyQuery, IQuerySpecification<RawPatternMatcher>> querySpecifications,
                                                                          Set<AnyQuery> vacuousQueries) {
                 this.store = store;
                 this.engineOptions = engineOptions;
                 this.inputKeys = inputKeys;
+                this.canonicalQueryMap = canonicalQueryMap;
                 this.querySpecifications = querySpecifications;
                 this.vacuousQueries = vacuousQueries;
                 var mutableAllQueries = new LinkedHashSet<AnyQuery>(querySpecifications.size() + vacuousQueries.size());
@@ -58,6 +62,17 @@ public class ViatraModelQueryStoreAdapterImpl implements ViatraModelQueryStoreAd
                 return allQueries;
         }
 
+        @Override
+        public <T> Query<T> getCanonicalQuery(Query<T> query) {
+                // We know that canonical forms of queries do not change output types.
+                @SuppressWarnings("unchecked")
+                var canonicalQuery = (Query<T>) canonicalQueryMap.get(query);
+                if (canonicalQuery == null) {
+                        throw new IllegalArgumentException("Unknown query: " + query);
+                }
+                return canonicalQuery;
+        }
+
         Map<AnyQuery, IQuerySpecification<RawPatternMatcher>> getQuerySpecifications() {
                 return querySpecifications;
         }
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 d51bc9fc..8110a98f 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
@@ -36,14 +36,12 @@ import tools.refinery.store.util.CycleDetectingMapper;
 
 import java.util.*;
 import java.util.function.Function;
-import java.util.stream.Collectors;
 
 public class Dnf2PQuery {
         private static final Object P_CONSTRAINT_LOCK = new Object();
         private final CycleDetectingMapper<Dnf, RawPQuery> mapper = new CycleDetectingMapper<>(Dnf::name,
                         this::doTranslate);
         private final QueryWrapperFactory wrapperFactory = new QueryWrapperFactory(this);
-        private final Map<Dnf, QueryEvaluationHint> hintOverrides = new LinkedHashMap<>();
         private Function<Dnf, QueryEvaluationHint> computeHint = dnf -> new QueryEvaluationHint(null,
                         (IQueryBackendFactory) null);
 
@@ -59,30 +57,9 @@ public class Dnf2PQuery {
                 return wrapperFactory.getSymbolViews();
         }
 
-        public void hint(Dnf dnf, QueryEvaluationHint hint) {
-                hintOverrides.compute(dnf, (ignoredKey, existingHint) ->
-                                existingHint == null ? hint : existingHint.overrideBy(hint));
-        }
-
-        private QueryEvaluationHint consumeHint(Dnf dnf) {
-                var defaultHint = computeHint.apply(dnf);
-                var existingHint = hintOverrides.remove(dnf);
-                return defaultHint.overrideBy(existingHint);
-        }
-
-        public void assertNoUnusedHints() {
-                if (hintOverrides.isEmpty()) {
-                        return;
-                }
-                var unusedHints = hintOverrides.keySet().stream().map(Dnf::name).collect(Collectors.joining(", "));
-                throw new IllegalStateException(
-                                "Unused query evaluation hints for %s. Hints must be set before a query is added to the engine"
-                                                .formatted(unusedHints));
-        }
-
         private RawPQuery doTranslate(Dnf dnfQuery) {
                 var pQuery = new RawPQuery(dnfQuery.getUniqueName());
-                pQuery.setEvaluationHints(consumeHint(dnfQuery));
+                pQuery.setEvaluationHints(computeHint.apply(dnfQuery));
 
                 Map<SymbolicParameter, PParameter> parameters = new HashMap<>();
                 List<PParameter> parameterList = new ArrayList<>();
diff --git a/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/viatra/DiagonalQueryTest.java b/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/viatra/DiagonalQueryTest.java
index 6aae2ebe..3d2d5f83 100644
--- a/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/viatra/DiagonalQueryTest.java
+++ b/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/viatra/DiagonalQueryTest.java
@@ -98,6 +98,7 @@ class DiagonalQueryTest {
                                 .build();
 
                 var model = store.createEmptyModel();
+
                 var personInterpretation = model.getInterpretation(person);
                 var symbolInterpretation = model.getInterpretation(symbol);
                 var queryEngine = model.getAdapter(ModelQueryAdapter.class);
diff --git a/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/viatra/QueryTest.java b/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/viatra/QueryTest.java
index be05b5e6..8ecbe11a 100644
--- a/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/viatra/QueryTest.java
+++ b/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/viatra/QueryTest.java
@@ -9,7 +9,9 @@ import org.eclipse.viatra.query.runtime.matchers.backend.QueryEvaluationHint;
 import org.junit.jupiter.api.Test;
 import tools.refinery.store.model.ModelStore;
 import tools.refinery.store.query.ModelQueryAdapter;
+import tools.refinery.store.query.dnf.Dnf;
 import tools.refinery.store.query.dnf.Query;
+import tools.refinery.store.query.term.ParameterDirection;
 import tools.refinery.store.query.term.Variable;
 import tools.refinery.store.query.viatra.tests.QueryEngineTest;
 import tools.refinery.store.query.view.AnySymbolView;
@@ -312,6 +314,53 @@ class QueryTest {
         }
 
         @QueryEngineTest
+        void patternCallInputArgumentTest(QueryEvaluationHint hint) {
+                var friendPredicate = Dnf.of("Friend", builder -> {
+                        var p1 = builder.parameter("p1", ParameterDirection.IN);
+                        var p2 = builder.parameter("p2", ParameterDirection.IN);
+                        builder.clause(
+                                        personView.call(p1),
+                                        personView.call(p2),
+                                        friendMustView.call(p1, p2)
+                        );
+                });
+                var predicate = Query.of("PositivePatternCall", (builder, p3, p4) -> builder.clause(
+                                personView.call(p3),
+                                personView.call(p4),
+                                friendPredicate.call(p3, p4)
+                ));
+
+                var store = ModelStore.builder()
+                                .symbols(person, friend)
+                                .with(ViatraModelQueryAdapter.builder()
+                                                .defaultHint(hint)
+                                                .queries(predicate))
+                                .build();
+
+                var model = store.createEmptyModel();
+                var personInterpretation = model.getInterpretation(person);
+                var friendInterpretation = model.getInterpretation(friend);
+                var queryEngine = model.getAdapter(ModelQueryAdapter.class);
+                var predicateResultSet = queryEngine.getResultSet(predicate);
+
+                personInterpretation.put(Tuple.of(0), true);
+                personInterpretation.put(Tuple.of(1), true);
+                personInterpretation.put(Tuple.of(2), true);
+
+                friendInterpretation.put(Tuple.of(0, 1), TruthValue.TRUE);
+                friendInterpretation.put(Tuple.of(1, 0), TruthValue.TRUE);
+                friendInterpretation.put(Tuple.of(1, 2), TruthValue.TRUE);
+
+                queryEngine.flushChanges();
+                assertResults(Map.of(
+                                Tuple.of(0, 1), true,
+                                Tuple.of(1, 0), true,
+                                Tuple.of(1, 2), true,
+                                Tuple.of(2, 1), false
+                ), predicateResultSet);
+        }
+
+        @QueryEngineTest
         void negativeRelationViewTest(QueryEvaluationHint hint) {
                 var predicate = Query.of("NegativePatternCall", (builder, p1, p2) -> builder.clause(
                                 personView.call(p1),
diff --git a/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/viatra/tests/QueryAssertions.java b/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/viatra/tests/QueryAssertions.java
index ca089a9d..5c1c4fc1 100644
--- a/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/viatra/tests/QueryAssertions.java
+++ b/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/viatra/tests/QueryAssertions.java
@@ -30,7 +30,7 @@ public final class QueryAssertions {
         }
 
         public static <T> void assertResults(Map<Tuple, T> expected, ResultSet<T> resultSet) {
-                var defaultValue = resultSet.getQuery().defaultValue();
+                var defaultValue = resultSet.getCanonicalQuery().defaultValue();
                 var filteredExpected = new LinkedHashMap<Tuple, T>();
                 var executables = new ArrayList<Executable>();
                 for (var entry : expected.entrySet()) {