Store references sources in free variables

5 files changed, 39 insertions, 34 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/FreeVariable.java b/src/main/java/org/shedlang/compiler/ast/typed/FreeVariable.java
index ff56adf5..f1791a9a 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/FreeVariable.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/FreeVariable.java
@@ -1,4 +1,8 @@
 package org.shedlang.compiler.ast.typed;
 
-public record FreeVariable(int index, String name) {
+import org.shedlang.compiler.sources.Source;
+
+import java.util.Set;
+
+public record FreeVariable(int index, String name, Set<Source> referenceSources) {
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/FreeVariables.java b/src/main/java/org/shedlang/compiler/ast/typed/FreeVariables.java
index 37f044cc..b905f5c4 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/FreeVariables.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/FreeVariables.java
@@ -1,8 +1,11 @@
 package org.shedlang.compiler.ast.typed;
 
+import org.pcollections.Empty;
+import org.pcollections.OrderedPMap;
 import org.pcollections.OrderedPSet;
+import org.pcollections.PSet;
+import org.shedlang.compiler.sources.Source;
 
-import java.util.Arrays;
 import java.util.Iterator;
 import java.util.stream.Stream;
 
@@ -11,21 +14,6 @@ public class FreeVariables implements Iterable<FreeVariable> {
         return new FreeVariables(OrderedPSet.empty());
     }
 
-    public static FreeVariables of(String... freeVariableNames) {
-        return of(Arrays.asList(freeVariableNames));
-    }
-
-    public static FreeVariables of(Iterable<String> freeVariableNames) {
-        var freeVariables = OrderedPSet.<FreeVariable>empty();
-        var freeVariableIndex = 0;
-        for (var freeVariableName : freeVariableNames) {
-            var freeVariable = new FreeVariable(freeVariableIndex, freeVariableName);
-            freeVariables = freeVariables.plus(freeVariable);
-            freeVariableIndex++;
-        }
-        return new FreeVariables(freeVariables);
-    }
-
     private final OrderedPSet<FreeVariable> freeVariables;
 
     private FreeVariables(OrderedPSet<FreeVariable> freeVariables) {
@@ -51,22 +39,35 @@ public class FreeVariables implements Iterable<FreeVariable> {
     }
 
     public static Builder builder() {
-        return new Builder(OrderedPSet.empty());
+        return new Builder(OrderedPMap.empty());
     }
 
     public static class Builder {
-        private final OrderedPSet<String> freeVariables;
+        private final OrderedPMap<String, PSet<Source>> freeVariables;
 
-        private Builder(OrderedPSet<String> freeVariables) {
+        private Builder(OrderedPMap<String, PSet<Source>> freeVariables) {
             this.freeVariables = freeVariables;
         }
 
-        public Builder plus(String name) {
-            return new Builder(this.freeVariables.plus(name));
+        public Builder plus(String name, Source referenceSource) {
+            return new Builder(this.freeVariables.plus(
+                name,
+                this.freeVariables.getOrDefault(name, Empty.set()).plus(referenceSource)
+            ));
         }
 
         public FreeVariables build() {
-            return FreeVariables.of(this.freeVariables);
+            var freeVariables = OrderedPSet.<FreeVariable>empty();
+            var freeVariableIndex = 0;
+            for (var entry : this.freeVariables.entrySet()) {
+                var name = entry.getKey();
+                var referenceSources = entry.getValue();
+
+                var freeVariable = new FreeVariable(freeVariableIndex, name, referenceSources);
+                freeVariables = freeVariables.plus(freeVariable);
+                freeVariableIndex++;
+            }
+            return new FreeVariables(freeVariables);
         }
     }
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
index 3a156671..2b14e345 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
@@ -148,7 +148,7 @@ public record TypedFunctionDefinitionNode(
                 this.namedParams,
                 this.effect,
                 this.returnType,
-                this.freeVariables.plus(freeVariable),
+                this.freeVariables.plus(freeVariable, Sources.NULL),
                 this.body
             );
         }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index f4c0af7f..f3018fdb 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -57,7 +57,7 @@ public class TypeChecker {
     }
 
     private static TypeCheckerResult<Type> runtimeLookup(String name, Source source, TypeCheckerContext context) {
-        return context.runtimeLookup(name).map(type -> {
+        return context.runtimeLookup(name, source).map(type -> {
             if (type.isEmpty()) {
                 throw new UnknownVariableError(name, source);
             }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index f64292d4..9fef9c61 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -45,7 +45,7 @@ public class TypeCheckerContext {
 
     private sealed interface StackElement {
         boolean isBoundInElement(String name);
-        StackElement withFreeVariable(String name);
+        StackElement withFreeVariable(String name, Source referenceSource);
     }
 
     private static final class FreeVariableCapture implements StackElement {
@@ -63,8 +63,8 @@ public class TypeCheckerContext {
         }
 
         @Override
-        public StackElement withFreeVariable(String name) {
-            return new FreeVariableCapture(this.freeVariables.plus(name));
+        public StackElement withFreeVariable(String name, Source referenceSource) {
+            return new FreeVariableCapture(this.freeVariables.plus(name, referenceSource));
         }
     }
 
@@ -98,13 +98,13 @@ public class TypeCheckerContext {
         }
 
         @Override
-        public LexicalContext withFreeVariable(String name) {
+        public LexicalContext withFreeVariable(String name, Source referenceSource) {
             return new LexicalContext(
                 this.environment,
                 this.effect,
                 this.returnType,
                 this.boundVariables,
-                this.freeVariables.plus(name),
+                this.freeVariables.plus(name, referenceSource),
                 this.state
             );
         }
@@ -215,14 +215,14 @@ public class TypeCheckerContext {
         }
     }
 
-    public TypeCheckerResult<Optional<Variable>> runtimeLookup(String name) {
+    public TypeCheckerResult<Optional<Variable>> runtimeLookup(String name, Source referenceSource) {
         var variable = Optional.ofNullable(environment().get(name));
 
-        var context = withVariable(name);
+        var context = withVariable(name, referenceSource);
         return context.toResult(variable);
     }
 
-    private TypeCheckerContext withVariable(String name) {
+    private TypeCheckerContext withVariable(String name, Source referenceSource) {
         var newLexicalContextStack = this.stack;
 
         for (var index = 0; index < this.stack.size(); index++) {
@@ -232,7 +232,7 @@ public class TypeCheckerContext {
             } else{
                 newLexicalContextStack = newLexicalContextStack.with(
                     index,
-                    stackElement.withFreeVariable(name)
+                    stackElement.withFreeVariable(name, referenceSource)
                 );
             }
         }