Compiler: fixes and improvements to closured variables

spec/compiler/semantic/closure_spec.cr

@@ -549,4 +549,200 @@ describe "Semantic: closure" do
     call = node.expressions[-2].as(Call)
     call.target_def.self_closured?.should be_false
   end
+
+  it "correctly detects previous var as closured (#5609)" do
+    assert_error %(
+      def block(&block)
+        block.call
+      end
+      def times
+        yield
+        yield
+      end
+      x = 1
+      times do
+        if x.is_a?(Int32)
+          x &+ 2
+        end
+        block do
+          x = "hello"
+        end
+      end
+      ),
+      "undefined method '&+' for String"
+  end
+
+  it "doesn't assign all types to metavar if closured but only assigned to once" do
+    semantic(%(
+      def capture(&block)
+        block
+      end
+      x = 1 == 2 ? 1 : nil
+      if x
+        capture do
+          x &+ 1
+        end
+      end
+      ))
+  end
+
+  it "does assign all types to metavar if closured but only assigned to once in a loop" do
+    assert_error %(
+      def capture(&block)
+        block
+      end
+      while 1 == 1
+        x = 1 == 2 ? 1 : nil
+        if x
+          capture do
+            x &+ 1
+          end
+        end
+      end
+      ),
+      "undefined method '&+'"
+  end
+
+  it "does assign all types to metavar if closured but only assigned to once in a loop through block" do
+    assert_error %(
+      def capture(&block)
+        block
+      end
+
+      def loop
+        while 1 == 1
+          yield
+        end
+      end
+
+      x = 1
+      loop do
+        x = 1 == 2 ? 1 : nil
+        if x
+          capture do
+            x &+ 1
+          end
+        end
+      end
+      ),
+      "undefined method '&+'"
+  end
+
+  it "does assign all types to metavar if closured but only assigned to once in a loop through captured block" do
+    assert_error %(
+      def capture(&block)
+        block
+      end
+
+      def loop(&block)
+        while 1 == 1
+          block.call
+        end
+      end
+
+      x = 1
+      loop do
+        x = 1 == 2 ? 1 : nil
+        if x
+          capture do
+            x &+ 1
+          end
+        end
+      end
+      ),
+      "undefined method '&+'"
+  end
+
+  it "doesn't assign all types to metavar if closured but declared inside block and never re-assigned" do
+    assert_no_errors %(
+      def capture(&block)
+        block
+      end
+
+      def loop(&block)
+        yield
+      end
+
+      loop do
+        x = 1 == 2 ? 1 : nil
+        if x
+          capture do
+            x &+ 1
+          end
+        end
+      end
+      )
+  end
+
+  it "doesn't assign all types to metavar if closured but declared inside block and re-assigned inside the same context before the closure" do
+    assert_no_errors %(
+      def capture(&block)
+        block
+      end
+
+      def loop(&block)
+        yield
+      end
+
+      loop do
+        x = 1 == 2 ? 1 : nil
+        x = 1 == 2 ? 1 : nil
+        if x
+          capture do
+            x &+ 1
+          end
+        end
+      end
+      )
+  end
+
+  it "is considered as closure if assigned once but comes from a method arg" do
+    assert_error %(
+      def capture(&block)
+        block
+      end
+      def foo(x)
+        capture do
+          x &+ 1
+        end
+        x = 1 == 2 ? 1 : nil
+      end
+      foo(1)
+      ),
+      "undefined method '&+'"
+  end
+
+  it "considers var as closure-readonly if it was assigned multiple times before it was closured" do
+    assert_no_errors(%(
+      def capture(&block)
+        block
+      end
+
+      x = "hello"
+      x = 1
+
+      capture do
+        x &+ 1
+      end
+      ))
+  end
+
+  it "correctly captures type of closured block arg" do
+    assert_type(%(
+      def capture(&block)
+        block.call
+      end
+      def foo
+        yield nil
+      end
+      z = nil
+      foo do |x|
+        capture do
+          x = 1
+        end
+        z = x
+      end
+      z
+      )) { nilable int32 }
+  end
 end

src/compiler/crystal/semantic/ast.cr

@@ -431,14 +431,40 @@ module Crystal
 
     # A variable is closured if it's used in a ProcLiteral context
     # where it wasn't created.
-    property? closured = false
+    getter? closured = false
 
     # Is this metavar assigned a value?
     property? assigned_to = false
 
+    # Is this metavar closured in a mutable way?
+    # This means it's closured and it got a value assigned to it more than once.
+    # If that's the case, when it's closured then all local variable related to
+    # it will also be bound to it.
+    property? mutably_closured = false
+
+    # Local variables associated with this meta variable.
+    # Can be Var or MetaVar.
+    property(local_vars) { [] of ASTNode }
+
     def initialize(@name : String, @type : Type? = nil)
     end
 
+    # Marks this variable as closured.
+    def mark_as_closured
+      @closured = true
+
+      return unless mutably_closured?
+
+      local_vars = @local_vars
+      return unless local_vars
+
+      # If a meta var is not readonly and it became a closure we must
+      # bind all previously related local vars to it so that
+      # they get all types assigned to it.
+      local_vars.each &.bind_to self
+      local_vars = nil
+    end
+
     # True if this variable belongs to the given context
     # but must be allocated in a closure.
     def closure_in?(context)
@@ -450,6 +476,11 @@ module Crystal
       @context.same?(context)
     end
 
+    # Is this metavar associated with any local vars?
+    def local_vars?
+      @local_vars
+    end
+
     def ==(other : self)
       name == other.name
     end
@@ -466,6 +497,7 @@ module Crystal
       end
       io << " (nil-if-read)" if nil_if_read?
       io << " (closured)" if closured?
+      io << " (mutably-closured)" if mutably_closured?
       io << " (assigned-to)" if assigned_to?
       io << " (object id: #{object_id})"
     end