Correct unbound variable handly in multi-line

lib/keisan/ast/assignment.rb

@@ -49,11 +49,13 @@ def evaluate_assignments(context = nil)
       end
 
       def unbound_variables(context = nil)
-        variables = super(context)
+        context ||= Context.new
+
         if is_variable_definition?
-          variables.delete(children.first.name)
+          variable_assignment_unbound_variables(context)
         else
-          variables
+          # TODO: Should update to handle function / list assignment.
+          super(context)
         end
       end
 
@@ -70,6 +72,10 @@ def is_variable_definition?
         children.first.is_a?(Variable)
       end
 
+      def variable_name
+        children.first.name
+      end
+
       def is_function_definition?
         children.first.is_a?(Function)
       end
@@ -78,6 +84,10 @@ def is_list_assignment?
         children.first.is_a?(List)
       end
 
+      def rhs_unbound_variables(context = nil)
+        children.last.unbound_variables(context)
+      end
+
       private
 
       def evaluate_variable_assignment(context, lhs, rhs)
@@ -96,6 +106,16 @@ def evaluate_list_assignment(context, lhs, rhs)
       def evaluate_cell_assignment(context, lhs, rhs)
         CellAssignment.new(self, context, lhs, rhs).evaluate
       end
+
+      def variable_assignment_unbound_variables(context)
+        rhs = rhs_unbound_variables(context)
+        # If the right-side is fully defined, then this is a valid assignment.
+        if rhs.empty?
+          Set.new
+        else
+          rhs | Set.new([variable_name])
+        end
+      end
     end
   end
 end

lib/keisan/ast/multi_line.rb

@@ -20,9 +20,41 @@ def simplify(context = nil)
         evaluate(context)
       end
 
+      def unbound_variables(context = nil)
+        context ||= Context.new
+        defined_variables = Set.new
+
+        children.inject(Set.new) do |unbound_variables, child|
+          if child.is_a?(Assignment) && child.is_variable_definition?
+            child_unbound_variables = variable_assignment_unbound_variables(context, child, defined_variables)
+            if child_unbound_variables.empty?
+              defined_variables.add(child.variable_name)
+              unbound_variables
+            else
+              unbound_variables | child_unbound_variables
+            end
+          else
+            unbound_variables | (child.unbound_variables(context) - defined_variables)
+          end
+        end
+      end
+
       def to_s
         children.map(&:to_s).join(";")
       end
+
+      private
+
+      def variable_assignment_unbound_variables(context, assignment, defined_variables)
+        rhs_child_unbound_variables = assignment.rhs_unbound_variables(context) - defined_variables
+
+        # If there are no unbound variables, this is a properly bound assignment.
+        if rhs_child_unbound_variables.empty?
+          Set.new
+        else
+          Set.new([assignment.variable_name]) | rhs_child_unbound_variables
+        end
+      end
     end
   end
 end

spec/keisan/ast/assignment_spec.rb

@@ -395,4 +395,72 @@
       end
     end
   end
+
+  describe "unbound_variables" do
+    let(:context) { Keisan::Context.new }
+
+    context "single-line" do
+      context "empty context" do
+        it "assigns to itself" do
+          ast = Keisan::AST.parse("x = x")
+          expect(ast.unbound_variables(context)).to eq Set.new(["x"])
+        end
+
+        it "has assigned variables bound" do
+          ast = Keisan::AST.parse("x = 5")
+          expect(ast.unbound_variables(context)).to eq Set.new
+        end
+
+        it "has variable set to unknown unbound" do
+          ast = Keisan::AST.parse("x = y")
+          expect(ast.unbound_variables(context)).to eq Set.new(["x", "y"])
+        end
+      end
+
+      context "with one defintion in context" do
+        let(:context) {
+          Keisan::Context.new.tap do |context|
+            context.register_variable!("x", 1)
+          end
+        }
+
+        it "can self-assign" do
+          ast = Keisan::AST.parse("x = x")
+          expect(ast.unbound_variables(context)).to eq Set.new
+        end
+
+        it "re-assigns" do
+          ast = Keisan::AST.parse("x = 5")
+          expect(ast.unbound_variables(context)).to eq Set.new
+        end
+
+        it "has assigned variables bound" do
+          ast = Keisan::AST.parse("y = x")
+          expect(ast.unbound_variables(context)).to eq Set.new
+        end
+
+        it "has assigned variables bound" do
+          ast = Keisan::AST.parse("y = x + z")
+          expect(ast.unbound_variables(context)).to eq Set.new(["y", "z"])
+        end
+      end
+    end
+
+    context "multi-line" do
+      it "binds assigned variables" do
+        ast = Keisan::AST.parse("x = 1; y = 2; x + y")
+        expect(ast.unbound_variables(context)).to eq Set.new
+      end
+
+      it "recognizes when one variable is assigned to" do
+        ast = Keisan::AST.parse("x = 3; y = 4; x + y")
+        expect(ast.unbound_variables(context)).to eq Set.new
+      end
+
+      it "recognizes when one variable is assigned to and one is not" do
+        ast = Keisan::AST.parse("x = 5; y = x + z; y")
+        expect(ast.unbound_variables(context)).to eq Set.new(["y", "z"])
+      end
+    end
+  end
 end

spec/keisan/ast/multi_line_spec.rb

@@ -0,0 +1,57 @@
+require "spec_helper"
+
+RSpec.describe Keisan::AST::MultiLine do
+  describe "unbound_variables" do
+    let(:context) { Keisan::Context.new }
+
+    context "with empty context" do
+      it "binds assigned variables" do
+        ast = Keisan::AST.parse("x = 1; y = 2; x + y")
+        expect(ast.unbound_variables(context)).to eq Set.new
+      end
+
+      it "recognizes when one variable is assigned to" do
+        ast = Keisan::AST.parse("x = 3; y = 4; x + y")
+        expect(ast.unbound_variables(context)).to eq Set.new
+      end
+
+      it "recognizes when one variable is assigned to and one is not" do
+        ast = Keisan::AST.parse("x = 5; y = x + z; y")
+        expect(ast.unbound_variables(context)).to eq Set.new(["y", "z"])
+      end
+
+      it "propagates unbound variable" do
+        ast = Keisan::AST.parse("x = 1; y = x")
+        expect(ast.unbound_variables(context)).to eq Set.new
+      end
+    end
+
+    context "with one variable defined" do
+      let(:context) {
+        Keisan::Context.new.tap do |context|
+          context.register_variable!("x", 1)
+        end
+      }
+
+      it "handles case when y is defined" do
+        ast = Keisan::AST.parse("y = 5; x + y")
+        expect(ast.unbound_variables(context)).to eq Set.new
+      end
+
+      it "handles case when y is dependent on x" do
+        ast = Keisan::AST.parse("y = x; x + y")
+        expect(ast.unbound_variables(context)).to eq Set.new
+      end
+
+      it "propagates unbound variable" do
+        ast = Keisan::AST.parse("y = x + z; x + y")
+        expect(ast.unbound_variables(context)).to eq Set.new(["y", "z"])
+      end
+
+      it "propagates unbound variable" do
+        ast = Keisan::AST.parse("x = 1; y = x")
+        expect(ast.unbound_variables(context)).to eq Set.new
+      end
+    end
+  end
+end