1 files changed, 130 insertions, 0 deletions

diff --git a/subprojects/language/src/main/java/tools/refinery/language/resource/state/DerivedVariableComputer.java b/subprojects/language/src/main/java/tools/refinery/language/resource/state/DerivedVariableComputer.java
new file mode 100644
index 00000000..f0baf35f
--- /dev/null
+++ b/subprojects/language/src/main/java/tools/refinery/language/resource/state/DerivedVariableComputer.java
@@ -0,0 +1,130 @@
+/*
+ * SPDX-FileCopyrightText: 2021-2023 The Refinery Authors <https://refinery.tools/>
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ */
+package tools.refinery.language.resource.state;
+
+import com.google.inject.Inject;
+import com.google.inject.Singleton;
+import com.google.inject.name.Named;
+import org.eclipse.xtext.linking.impl.LinkingHelper;
+import org.eclipse.xtext.naming.IQualifiedNameConverter;
+import org.eclipse.xtext.scoping.IScopeProvider;
+import org.eclipse.xtext.scoping.impl.AbstractDeclarativeScopeProvider;
+import tools.refinery.language.model.problem.*;
+
+import java.util.*;
+
+@Singleton
+public class DerivedVariableComputer {
+        @Inject
+        private LinkingHelper linkingHelper;
+
+        @Inject
+        private IQualifiedNameConverter qualifiedNameConverter;
+
+        @Inject
+        @Named(AbstractDeclarativeScopeProvider.NAMED_DELEGATE)
+        private IScopeProvider scopeProvider;
+
+        public void installDerivedVariables(Problem problem, Set<String> nodeNames) {
+                for (Statement statement : problem.getStatements()) {
+                        if (statement instanceof ParametricDefinition definition) {
+                                installDerivedParametricDefinitionState(definition, nodeNames);
+                        }
+                }
+        }
+
+        protected void installDerivedParametricDefinitionState(ParametricDefinition definition, Set<String> nodeNames) {
+                Set<String> knownVariables = new HashSet<>(nodeNames);
+                for (Parameter parameter : definition.getParameters()) {
+                        String name = parameter.getName();
+                        if (name != null) {
+                                knownVariables.add(name);
+                        }
+                }
+                if (definition instanceof PredicateDefinition predicateDefinition) {
+                        installDerivedPredicateDefinitionState(predicateDefinition, knownVariables);
+                } else if (definition instanceof FunctionDefinition functionDefinition) {
+                        installDerivedFunctionDefinitionState(functionDefinition, knownVariables);
+                } else if (definition instanceof RuleDefinition ruleDefinition) {
+                        installDerivedRuleDefinitionState(ruleDefinition, knownVariables);
+                } else {
+                        throw new IllegalArgumentException("Unknown ParametricDefinition: " + definition);
+                }
+        }
+
+        protected void installDerivedPredicateDefinitionState(PredicateDefinition definition, Set<String> knownVariables) {
+                for (Conjunction body : definition.getBodies()) {
+                        createVariablesForScope(new ImplicitVariableScope(body, knownVariables));
+                }
+        }
+
+        protected void installDerivedFunctionDefinitionState(FunctionDefinition definition, Set<String> knownVariables) {
+                for (Case body : definition.getCases()) {
+                        if (body instanceof Conjunction conjunction) {
+                                createVariablesForScope(new ImplicitVariableScope(conjunction, knownVariables));
+                        } else if (body instanceof Match match) {
+                                var condition = match.getCondition();
+                                if (condition != null) {
+                                        createVariablesForScope(new ImplicitVariableScope(match, match.getCondition(), knownVariables));
+                                }
+                        } else {
+                                throw new IllegalArgumentException("Unknown Case: " + body);
+                        }
+                }
+        }
+
+        protected void installDerivedRuleDefinitionState(RuleDefinition definition, Set<String> knownVariables) {
+                for (Conjunction precondition : definition.getPreconditions()) {
+                        createVariablesForScope(new ImplicitVariableScope(precondition, knownVariables));
+                }
+        }
+
+        protected void createVariablesForScope(ImplicitVariableScope scope) {
+                var queue = new ArrayDeque<ImplicitVariableScope>();
+                queue.addLast(scope);
+                while (!queue.isEmpty()) {
+                        var nextScope = queue.removeFirst();
+                        nextScope.createVariables(scopeProvider, linkingHelper, qualifiedNameConverter, queue);
+                }
+        }
+
+        public void discardDerivedVariables(Problem problem) {
+                for (Statement statement : problem.getStatements()) {
+                        if (statement instanceof ParametricDefinition parametricDefinition) {
+                                discardParametricDefinitionState(parametricDefinition);
+                        }
+                }
+        }
+
+        protected void discardParametricDefinitionState(ParametricDefinition definition) {
+                List<ExistentialQuantifier> existentialQuantifiers = new ArrayList<>();
+                List<VariableOrNodeExpr> variableOrNodeExprs = new ArrayList<>();
+                var treeIterator = definition.eAllContents();
+                // We must collect the nodes where we are discarding derived state and only discard them after the iteration,
+                // because modifying the containment hierarchy during iteration causes the TreeIterator to fail with
+                // IndexOutOfBoundsException.
+                while (treeIterator.hasNext()) {
+                        var child = treeIterator.next();
+                        var containingFeature = child.eContainingFeature();
+                        if (containingFeature == ProblemPackage.Literals.EXISTENTIAL_QUANTIFIER__IMPLICIT_VARIABLES ||
+                                        containingFeature == ProblemPackage.Literals.VARIABLE_OR_NODE_EXPR__SINGLETON_VARIABLE) {
+                                treeIterator.prune();
+                        } else if (child instanceof ExistentialQuantifier existentialQuantifier &&
+                                        !existentialQuantifier.getImplicitVariables().isEmpty()) {
+                                existentialQuantifiers.add(existentialQuantifier);
+                        } else if (child instanceof VariableOrNodeExpr variableOrNodeExpr &&
+                                        variableOrNodeExpr.getSingletonVariable() != null) {
+                                variableOrNodeExprs.add(variableOrNodeExpr);
+                        }
+                }
+                for (var existentialQuantifier : existentialQuantifiers) {
+                        existentialQuantifier.getImplicitVariables().clear();
+                }
+                for (var variableOrNodeExpr : variableOrNodeExprs) {
+                        variableOrNodeExpr.setSingletonVariable(null);
+                }
+        }
+}