Extract DependencyChecker

2 files changed, 112 insertions, 64 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/DependencyChecker.java b/src/main/java/org/shedlang/compiler/typechecker/DependencyChecker.java
new file mode 100644
index 00000000..b1d7c1ac
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/typechecker/DependencyChecker.java
@@ -0,0 +1,90 @@
+package org.shedlang.compiler.typechecker;
+
+import org.pcollections.Empty;
+import org.pcollections.PMap;
+import org.pcollections.PVector;
+import org.shedlang.compiler.ast.typed.FreeVariable;
+import org.shedlang.compiler.ast.untyped.UntypedVariableDeclarationNode;
+import org.shedlang.compiler.typechecker.errors.ReferenceBeforeDefinitionError;
+import org.shedlang.compiler.util.Iterables;
+
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.Set;
+
+class DependencyChecker {
+    static final DependencyChecker EMPTY = new DependencyChecker(Empty.map(), Empty.vector());
+
+    private record ImmediateDependencies(
+        Collection<UntypedVariableDeclarationNode> variableDeclarations,
+        Set<FreeVariable> dependencies
+    ) {}
+
+    private final PMap<String, Set<FreeVariable>> deferredDependencies;
+    private final PVector<ImmediateDependencies> immediateDependencies;
+
+    private DependencyChecker(
+        PMap<String, Set<FreeVariable>> deferredDependencies,
+        PVector<ImmediateDependencies> immediateDependencies
+    ) {
+        this.deferredDependencies = deferredDependencies;
+        this.immediateDependencies = immediateDependencies;
+    }
+
+    DependencyChecker addDeferredDependencies(String name, Set<FreeVariable> dependencies) {
+        return new DependencyChecker(
+            // TODO: check that name is not already in this.deferredDependencies
+            this.deferredDependencies.plus(name, dependencies),
+            this.immediateDependencies
+        );
+    }
+
+    DependencyChecker addImmediateDependencies(
+        Collection<UntypedVariableDeclarationNode> variableDeclarations,
+        Set<FreeVariable> dependencies
+    ) {
+        return new DependencyChecker(
+            this.deferredDependencies,
+            this.immediateDependencies.plus(new ImmediateDependencies(variableDeclarations, dependencies))
+        );
+    }
+
+    void check() {
+        var state = new Object() {
+            HashSet<String> defined = new HashSet<>();
+
+            void define(String name) {
+                defined.add(name);
+            }
+
+            void require(FreeVariable dependency) {
+                if (defined.contains(dependency.name())) {
+                    return;
+                }
+
+                var subdependencies = deferredDependencies.get(dependency.name());
+
+                if (subdependencies == null) {
+                    throw new ReferenceBeforeDefinitionError(
+                        dependency.name(),
+                        Iterables.first(dependency.referenceSources())
+                    );
+                }
+
+                for (var subdependency : subdependencies) {
+                    require(subdependency);
+                }
+            }
+        };
+
+        for (var immediateDependencies : this.immediateDependencies) {
+            for (var dependency : immediateDependencies.dependencies) {
+                state.require(dependency);
+            }
+
+            for (var declaredVariable : immediateDependencies.variableDeclarations) {
+                state.define(declaredVariable.name());
+            }
+        }
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 1cc7a7d6..3fbf9bb6 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -7,7 +7,6 @@ import org.shedlang.compiler.typechecker.errors.*;
 import org.shedlang.compiler.typechecker.typesets.TypeSetTypeLevelValue;
 import org.shedlang.compiler.typechecker.typesets.TypeSets;
 import org.shedlang.compiler.types.*;
-import org.shedlang.compiler.util.Iterables;
 import org.shedlang.compiler.util.P;
 import org.shedlang.compiler.util.PCollectors;
 
@@ -595,7 +594,12 @@ public class TypeChecker {
             .map(declaration -> declaration.name())
             .collect(Collectors.toSet());
 
-        for (var typeCheckSteps : typeCheckStepSets) {
+        var dependencyChecker = DependencyChecker.EMPTY;
+
+        for (var i = 0; i < untyped.body().size(); i++) {
+            var untypedStatement = untyped.body().get(i);
+            var typeCheckSteps = typeCheckStepSets.get(i);
+
             context = context.enterFreeVariableCapture();
 
             var typedStatementResult = typeCheckSteps.define(context);
@@ -611,10 +615,24 @@ public class TypeChecker {
             var dependencies = freeVariables.stream()
                 .filter(freeVariable -> variableNames.contains(freeVariable.name()))
                 .collect(Collectors.toSet());
-            dependenciesByBodyStatementIndex.add(dependencies);
+
+            switch (untypedStatement) {
+                case UntypedFunctionDefinitionNode untypedFunctionDefinition -> {
+                    dependencyChecker = dependencyChecker.addDeferredDependencies(
+                        untypedFunctionDefinition.name(),
+                        dependencies
+                    );
+                }
+                case UntypedValNode untypedVal -> {
+                    dependencyChecker = dependencyChecker.addImmediateDependencies(
+                        declaredVariables(untypedVal, P.vector()),
+                        dependencies
+                    );
+                }
+            }
         }
 
-        checkDependenciesBetweenNamespaceStatements(untyped, dependenciesByBodyStatementIndex);
+        dependencyChecker.check();
 
         var fieldTypes = typedBody.stream()
             .flatMap(bodyStatement ->
@@ -639,66 +657,6 @@ public class TypeChecker {
             .toResult(typed);
     }
 
-    private static void checkDependenciesBetweenNamespaceStatements(
-        UntypedNamespaceNode untyped,
-        List<Set<FreeVariable>> dependenciesByBodyStatementIndex
-    ) {
-        var deferredDependenciesByFunctionName = new HashMap<String, Set<FreeVariable>>();
-        for (var i = 0; i < untyped.body().size(); i++) {
-            var untypedStatement = untyped.body().get(i);
-            if (untypedStatement instanceof UntypedFunctionDefinitionNode functionDefinition) {
-                var dependencies = dependenciesByBodyStatementIndex.get(i);
-                deferredDependenciesByFunctionName.put(functionDefinition.name(), dependencies);
-            }
-        }
-
-        var state = new Object() {
-            HashSet<String> defined = new HashSet<>();
-
-            void define(String name) {
-                defined.add(name);
-            }
-
-            void require(FreeVariable dependency) {
-                if (defined.contains(dependency.name())) {
-                    return;
-                }
-
-                var subdependencies = deferredDependenciesByFunctionName.get(dependency.name());
-
-                if (subdependencies == null) {
-                    throw new ReferenceBeforeDefinitionError(
-                        dependency.name(),
-                        Iterables.first(dependency.referenceSources())
-                    );
-                }
-
-                for (var subdependency : subdependencies) {
-                    require(subdependency);
-                }
-            }
-        };
-
-        for (var i = 0; i < untyped.body().size(); i++) {
-            var untypedStatement = untyped.body().get(i);
-            var dependencies = dependenciesByBodyStatementIndex.get(i);
-
-            switch (untypedStatement) {
-                case UntypedFunctionDefinitionNode functionDefinition -> {
-                }
-                case UntypedValNode val -> {
-                    for (var dependency : dependencies) {
-                        state.require(dependency);
-                    }
-
-                    for (var declaredVariable : declaredVariables(val, P.vector())) {
-                        state.define(declaredVariable.name());
-                    }
-                }
-            }
-        }
-    }
-
     static TypeCheckSteps<? extends TypedNamespaceStatementNode> typeCheckNamespaceStatement(
         UntypedNamespaceStatementNode untyped
     ) {