1 files changed, 126 insertions, 75 deletions

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 201e0ed0..7afeb977 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,45 +1,44 @@
 package tools.refinery.store.query.viatra.internal.pquery;
 
+import org.eclipse.viatra.query.runtime.matchers.backend.IQueryBackendFactory;
 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.aggregations.BoundAggregator;
+import org.eclipse.viatra.query.runtime.matchers.psystem.aggregations.IMultisetAggregationOperator;
 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;
-import org.eclipse.viatra.query.runtime.matchers.psystem.basicdeferred.NegativePatternCall;
+import org.eclipse.viatra.query.runtime.matchers.psystem.basicdeferred.*;
 import org.eclipse.viatra.query.runtime.matchers.psystem.basicenumerables.BinaryTransitiveClosure;
 import org.eclipse.viatra.query.runtime.matchers.psystem.basicenumerables.ConstantValue;
 import org.eclipse.viatra.query.runtime.matchers.psystem.basicenumerables.PositivePatternCall;
 import org.eclipse.viatra.query.runtime.matchers.psystem.basicenumerables.TypeConstraint;
 import org.eclipse.viatra.query.runtime.matchers.psystem.queries.PParameter;
-import org.eclipse.viatra.query.runtime.matchers.psystem.queries.PVisibility;
+import org.eclipse.viatra.query.runtime.matchers.psystem.queries.PQuery;
 import org.eclipse.viatra.query.runtime.matchers.tuple.Tuple;
 import org.eclipse.viatra.query.runtime.matchers.tuple.Tuples;
-import tools.refinery.store.query.Dnf;
-import tools.refinery.store.query.DnfClause;
-import tools.refinery.store.query.DnfUtils;
-import tools.refinery.store.query.Variable;
+import tools.refinery.store.query.dnf.Dnf;
+import tools.refinery.store.query.dnf.DnfClause;
 import tools.refinery.store.query.literal.*;
+import tools.refinery.store.query.term.ConstantTerm;
+import tools.refinery.store.query.term.StatefulAggregator;
+import tools.refinery.store.query.term.StatelessAggregator;
+import tools.refinery.store.query.term.Variable;
 import tools.refinery.store.query.view.AnyRelationView;
 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 Map<AnyRelationView, RelationViewWrapper> view2WrapperMap = new LinkedHashMap<>();
-
-        private final Map<AnyRelationView, RawPQuery> view2EmbeddedMap = new HashMap<>();
-
+        private final QueryWrapperFactory wrapperFactory = new QueryWrapperFactory(this);
+        private final Map<Dnf, QueryEvaluationHint> hintOverrides = new LinkedHashMap<>();
         private Function<Dnf, QueryEvaluationHint> computeHint = dnf -> new QueryEvaluationHint(null,
-                        QueryEvaluationHint.BackendRequirement.UNSPECIFIED);
+                        (IQueryBackendFactory) null);
 
         public void setComputeHint(Function<Dnf, QueryEvaluationHint> computeHint) {
                 this.computeHint = computeHint;
@@ -50,16 +49,33 @@ public class Dnf2PQuery {
         }
 
         public Map<AnyRelationView, IInputKey> getRelationViews() {
-                return Collections.unmodifiableMap(view2WrapperMap);
+                return wrapperFactory.getRelationViews();
         }
 
-        public RawPQuery getAlreadyTranslated(Dnf dnfQuery) {
-                return mapper.getAlreadyMapped(dnfQuery);
+        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(computeHint.apply(dnfQuery));
+                pQuery.setEvaluationHints(consumeHint(dnfQuery));
 
                 Map<Variable, PParameter> parameters = new HashMap<>();
                 for (Variable variable : dnfQuery.getParameters()) {
@@ -97,7 +113,7 @@ public class Dnf2PQuery {
                                 body.setSymbolicParameters(symbolicParameters);
                                 pQuery.addBody(body);
                                 for (Literal literal : clause.literals()) {
-                                        translateLiteral(literal, body);
+                                        translateLiteral(literal, clause, body);
                                 }
                         }
                 }
@@ -105,15 +121,21 @@ public class Dnf2PQuery {
                 return pQuery;
         }
 
-        private void translateLiteral(Literal literal, PBody body) {
+        private void translateLiteral(Literal literal, DnfClause clause, PBody body) {
                 if (literal instanceof EquivalenceLiteral equivalenceLiteral) {
                         translateEquivalenceLiteral(equivalenceLiteral, body);
-                } else if (literal instanceof RelationViewLiteral relationViewLiteral) {
-                        translateRelationViewLiteral(relationViewLiteral, body);
-                } else if (literal instanceof DnfCallLiteral dnfCallLiteral) {
-                        translateDnfCallLiteral(dnfCallLiteral, body);
+                } else if (literal instanceof CallLiteral callLiteral) {
+                        translateCallLiteral(callLiteral, clause, body);
                 } else if (literal instanceof ConstantLiteral constantLiteral) {
                         translateConstantLiteral(constantLiteral, body);
+                } else if (literal instanceof AssignLiteral<?> assignLiteral) {
+                        translateAssignLiteral(assignLiteral, body);
+                } else if (literal instanceof AssumeLiteral assumeLiteral) {
+                        translateAssumeLiteral(assumeLiteral, body);
+                } else if (literal instanceof CountLiteral countLiteral) {
+                        translateCountLiteral(countLiteral, clause, body);
+                } else if (literal instanceof AggregationLiteral<?, ?> aggregationLiteral) {
+                        translateAggregationLiteral(aggregationLiteral, clause, body);
                 } else {
                         throw new IllegalArgumentException("Unknown literal: " + literal.toString());
                 }
@@ -129,20 +151,43 @@ public class Dnf2PQuery {
                 }
         }
 
-        private void translateRelationViewLiteral(RelationViewLiteral relationViewLiteral, PBody body) {
-                var substitution = translateSubstitution(relationViewLiteral.getArguments(), body);
-                var polarity = relationViewLiteral.getPolarity();
-                var relationView = relationViewLiteral.getTarget();
-                if (polarity == CallPolarity.POSITIVE) {
-                        new TypeConstraint(body, substitution, wrapView(relationView));
-                } else {
-                        var embeddedPQuery = translateEmbeddedRelationViewPQuery(relationView);
-                        switch (polarity) {
-                        case TRANSITIVE -> new BinaryTransitiveClosure(body, substitution, embeddedPQuery);
-                        case NEGATIVE -> new NegativePatternCall(body, substitution, embeddedPQuery);
-                        default -> throw new IllegalArgumentException("Unknown polarity: " + polarity);
+        private void translateCallLiteral(CallLiteral callLiteral, DnfClause clause, PBody body) {
+                var polarity = callLiteral.getPolarity();
+                switch (polarity) {
+                case POSITIVE -> {
+                        var substitution = translateSubstitution(callLiteral.getArguments(), body);
+                        var constraint = callLiteral.getTarget();
+                        if (constraint instanceof Dnf dnf) {
+                                var pattern = translate(dnf);
+                                new PositivePatternCall(body, substitution, pattern);
+                        } else if (constraint instanceof AnyRelationView relationView) {
+                                var inputKey = wrapperFactory.getInputKey(relationView);
+                                new TypeConstraint(body, substitution, inputKey);
+                        } else {
+                                throw new IllegalArgumentException("Unknown Constraint: " + constraint);
                         }
                 }
+                case TRANSITIVE -> {
+                        var substitution = translateSubstitution(callLiteral.getArguments(), body);
+                        var constraint = callLiteral.getTarget();
+                        PQuery pattern;
+                        if (constraint instanceof Dnf dnf) {
+                                pattern = translate(dnf);
+                        } else if (constraint instanceof AnyRelationView relationView) {
+                                pattern = wrapperFactory.wrapRelationViewIdentityArguments(relationView);
+                        } else {
+                                throw new IllegalArgumentException("Unknown Constraint: " + constraint);
+                        }
+                        new BinaryTransitiveClosure(body, substitution, pattern);
+                }
+                case NEGATIVE -> {
+                        var wrappedCall = wrapperFactory.maybeWrapConstraint(callLiteral, clause);
+                        var substitution = translateSubstitution(wrappedCall.remappedArguments(), body);
+                        var pattern = wrappedCall.pattern();
+                        new NegativePatternCall(body, substitution, pattern);
+                }
+                default -> throw new IllegalArgumentException("Unknown polarity: " + polarity);
+                }
         }
 
         private static Tuple translateSubstitution(List<Variable> substitution, PBody body) {
@@ -155,51 +200,57 @@ public class Dnf2PQuery {
                 return Tuples.flatTupleOf(variables);
         }
 
-        private RawPQuery translateEmbeddedRelationViewPQuery(AnyRelationView relationView) {
-                return view2EmbeddedMap.computeIfAbsent(relationView, this::doTranslateEmbeddedRelationViewPQuery);
+        private void translateConstantLiteral(ConstantLiteral constantLiteral, PBody body) {
+                var variable = body.getOrCreateVariableByName(constantLiteral.variable().getUniqueName());
+                new ConstantValue(body, variable, constantLiteral.nodeId());
         }
 
-        private RawPQuery doTranslateEmbeddedRelationViewPQuery(AnyRelationView relationView) {
-                var embeddedPQuery = new RawPQuery(DnfUtils.generateUniqueName(relationView.name()), PVisibility.EMBEDDED);
-                var body = new PBody(embeddedPQuery);
-                int arity = relationView.arity();
-                var parameters = new ArrayList<PParameter>(arity);
-                var arguments = new Object[arity];
-                var symbolicParameters = new ArrayList<ExportedParameter>(arity);
-                for (int i = 0; i < arity; i++) {
-                        var parameterName = "p" + i;
-                        var parameter = new PParameter(parameterName);
-                        parameters.add(parameter);
-                        var variable = body.getOrCreateVariableByName(parameterName);
-                        arguments[i] = variable;
-                        symbolicParameters.add(new ExportedParameter(body, variable, parameter));
+        private <T> void translateAssignLiteral(AssignLiteral<T> assignLiteral, PBody body) {
+                var variable = body.getOrCreateVariableByName(assignLiteral.variable().getUniqueName());
+                var term = assignLiteral.term();
+                if (term instanceof ConstantTerm<T> constantTerm) {
+                        new ConstantValue(body, variable, constantTerm.getValue());
+                } else {
+                        var evaluator = new TermEvaluator<>(term);
+                        new ExpressionEvaluation(body, evaluator, variable);
                 }
-                embeddedPQuery.setParameters(parameters);
-                body.setSymbolicParameters(symbolicParameters);
-                var argumentTuple = Tuples.flatTupleOf(arguments);
-                new TypeConstraint(body, argumentTuple, wrapView(relationView));
-                embeddedPQuery.addBody(body);
-                return embeddedPQuery;
         }
 
-        private RelationViewWrapper wrapView(AnyRelationView relationView) {
-                return view2WrapperMap.computeIfAbsent(relationView, RelationViewWrapper::new);
+        private void translateAssumeLiteral(AssumeLiteral assumeLiteral, PBody body) {
+                var evaluator = new AssumptionEvaluator(assumeLiteral.term());
+                new ExpressionEvaluation(body, evaluator, null);
         }
 
-        private void translateDnfCallLiteral(DnfCallLiteral dnfCallLiteral, PBody body) {
-                var variablesTuple = translateSubstitution(dnfCallLiteral.getArguments(), body);
-                var translatedReferred = translate(dnfCallLiteral.getTarget());
-                var polarity = dnfCallLiteral.getPolarity();
-                switch (polarity) {
-                case POSITIVE -> new PositivePatternCall(body, variablesTuple, translatedReferred);
-                case TRANSITIVE -> new BinaryTransitiveClosure(body, variablesTuple, translatedReferred);
-                case NEGATIVE -> new NegativePatternCall(body, variablesTuple, translatedReferred);
-                default -> throw new IllegalArgumentException("Unknown polarity: " + polarity);
-                }
+        private void translateCountLiteral(CountLiteral countLiteral, DnfClause clause, PBody body) {
+                var wrappedCall = wrapperFactory.maybeWrapConstraint(countLiteral, clause);
+                var substitution = translateSubstitution(wrappedCall.remappedArguments(), body);
+                var resultVariable = body.getOrCreateVariableByName(countLiteral.getResultVariable().getUniqueName());
+                new PatternMatchCounter(body, substitution, wrappedCall.pattern(), resultVariable);
         }
 
-        private void translateConstantLiteral(ConstantLiteral constantLiteral, PBody body) {
-                var variable = body.getOrCreateVariableByName(constantLiteral.variable().getUniqueName());
-                new ConstantValue(body, variable, constantLiteral.nodeId());
+        private <R, T> void translateAggregationLiteral(AggregationLiteral<R, T> aggregationLiteral, DnfClause clause,
+                                                                                                        PBody body) {
+                var aggregator = aggregationLiteral.getAggregator();
+                IMultisetAggregationOperator<T, ?, R> aggregationOperator;
+                if (aggregator instanceof StatelessAggregator<R, T> statelessAggregator) {
+                        aggregationOperator = new StatelessMultisetAggregator<>(statelessAggregator);
+                } else if (aggregator instanceof StatefulAggregator<R, T> statefulAggregator) {
+                        aggregationOperator = new StatefulMultisetAggregator<>(statefulAggregator);
+                } else {
+                        throw new IllegalArgumentException("Unknown aggregator: " + aggregator);
+                }
+                var wrappedCall = wrapperFactory.maybeWrapConstraint(aggregationLiteral, clause);
+                var substitution = translateSubstitution(wrappedCall.remappedArguments(), body);
+                var inputVariable = body.getOrCreateVariableByName(aggregationLiteral.getInputVariable().getUniqueName());
+                var aggregatedColumn = substitution.invertIndex().get(inputVariable);
+                if (aggregatedColumn == null) {
+                        throw new IllegalStateException("Input variable %s not found in substitution %s".formatted(inputVariable,
+                                        substitution));
+                }
+                var boundAggregator = new BoundAggregator(aggregationOperator, aggregator.getInputType(),
+                                aggregator.getResultType());
+                var resultVariable = body.getOrCreateVariableByName(aggregationLiteral.getResultVariable().getUniqueName());
+                new AggregatorConstraint(boundAggregator, body, substitution, wrappedCall.pattern(), resultVariable,
+                                aggregatedColumn);
         }
 }