refactor(generator): scope overrides

1 files changed, 61 insertions, 0 deletions

diff --git a/subprojects/generator/src/main/java/tools/refinery/generator/ProblemLoader.java b/subprojects/generator/src/main/java/tools/refinery/generator/ProblemLoader.java
index abfbe7e6..20ea8132 100644
--- a/subprojects/generator/src/main/java/tools/refinery/generator/ProblemLoader.java
+++ b/subprojects/generator/src/main/java/tools/refinery/generator/ProblemLoader.java
@@ -17,11 +17,18 @@ import org.eclipse.xtext.util.LazyStringInputStream;
 import org.eclipse.xtext.validation.CheckMode;
 import org.eclipse.xtext.validation.IResourceValidator;
 import tools.refinery.language.model.problem.Problem;
+import tools.refinery.language.model.problem.Relation;
+import tools.refinery.language.model.problem.ScopeDeclaration;
 import tools.refinery.store.util.CancellationToken;
 
+import java.io.ByteArrayOutputStream;
 import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
 
 public class ProblemLoader {
@@ -104,4 +111,58 @@ public class ProblemLoader {
                 }
                 return problem;
         }
+
+        public Problem loadScopeConstraints(Problem problem, List<String> extraScopes,
+                                                                                List<String> overrideScopes) throws IOException {
+                var allScopes = new ArrayList<>(extraScopes);
+                allScopes.addAll(overrideScopes);
+                if (allScopes.isEmpty()) {
+                        return problem;
+                }
+                int originalStatementCount = problem.getStatements().size();
+                var builder = new StringBuilder();
+                var problemResource = problem.eResource();
+                try (var outputStream = new ByteArrayOutputStream()) {
+                        problemResource.save(outputStream, Map.of());
+                        builder.append(outputStream.toString(StandardCharsets.UTF_8));
+                }
+                builder.append('\n');
+                for (var scope : allScopes) {
+                        builder.append("scope ").append(scope).append(".\n");
+                }
+                var modifiedProblem = loadString(builder.toString(), problemResource.getURI());
+                var modifiedStatements = modifiedProblem.getStatements();
+                int modifiedStatementCount = modifiedStatements.size();
+                if (modifiedStatementCount != originalStatementCount + allScopes.size()) {
+                        throw new IllegalArgumentException("Failed to parse scope constraints");
+                }
+                // Override scopes remove any scope constraint from the original problem with the same target type.
+                var overriddenScopes = new HashSet<Relation>();
+                for (int i = modifiedStatementCount - overrideScopes.size(); i < modifiedStatementCount; i++) {
+                        var statement = modifiedStatements.get(i);
+                        if (!(statement instanceof ScopeDeclaration scopeDeclaration)) {
+                                throw new IllegalStateException("Invalid scope constraint: " + statement);
+                        }
+                        for (var typeScope : scopeDeclaration.getTypeScopes()) {
+                                overriddenScopes.add(typeScope.getTargetType());
+                        }
+                }
+                int statementIndex = 0;
+                var iterator = modifiedStatements.iterator();
+                // Scope overrides only affect type scopes from the original problem and leave type scopes added on the
+                // command line intact.
+                while (statementIndex < originalStatementCount && iterator.hasNext()) {
+                        var statement = iterator.next();
+                        if (statement instanceof ScopeDeclaration scopeDeclaration) {
+                                var typeScopes = scopeDeclaration.getTypeScopes();
+                                typeScopes.removeIf(typeScope -> overriddenScopes.contains(typeScope.getTargetType()));
+                                // Scope declarations with no type scopes are invalid, so we have to remove them.
+                                if (typeScopes.isEmpty()) {
+                                        iterator.remove();
+                                }
+                        }
+                        statementIndex++;
+                }
+                return modifiedProblem;
+        }
 }