Make Lint/SharedVarInFiber rule account for loop { ... }

spec/ameba/ast/branch_spec.cr

@@ -298,6 +298,34 @@ module Ameba::AST
       end
     end
 
+    context "Crystal::Call" do
+      context "loop" do
+        it "constructs a branch in block" do
+          branch = branch_of_assign_in_def <<-CRYSTAL
+            def method(a)
+              loop do
+                b = (a = 1)
+              end
+            end
+            CRYSTAL
+          branch.to_s.should eq "b = (a = 1)"
+        end
+      end
+
+      context "other" do
+        it "skips constructing a branch in block" do
+          branch = branch_of_assign_in_def <<-CRYSTAL
+            def method(a)
+              1.upto(10) do
+                b = (a = 1)
+              end
+            end
+            CRYSTAL
+          branch.should be_nil
+        end
+      end
+    end
+
     describe "#initialize" do
       it "creates new branch" do
         nodes = as_nodes <<-CRYSTAL
@@ -358,6 +386,30 @@ module Ameba::AST
         branch = Branch.new nodes.assign_nodes.first, branchable
         branch.in_loop?.should be_false
       end
+
+      context "Crystal::Call" do
+        it "returns true if branch is in a loop" do
+          nodes = as_nodes <<-CRYSTAL
+            loop do
+              a = 1
+            end
+            CRYSTAL
+          branchable = Branchable.new nodes.call_nodes.first
+          branch = Branch.new nodes.assign_nodes.first, branchable
+          branch.in_loop?.should be_true
+        end
+
+        it "returns false if branch is not in a loop" do
+          nodes = as_nodes <<-CRYSTAL
+            1.upto(10) do
+              a = 1
+            end
+            CRYSTAL
+          branchable = Branchable.new nodes.call_nodes.first
+          branch = Branch.new nodes.assign_nodes.first, branchable
+          branch.in_loop?.should be_false
+        end
+      end
     end
   end
 end

spec/ameba/rule/lint/shared_var_in_fiber_spec.cr

@@ -39,7 +39,7 @@ module Ameba::Rule::Lint
         CRYSTAL
     end
 
-    it "reports if there is a shared var in spawn" do
+    it "reports if there is a shared var in spawn (while)" do
       source = expect_issue subject, <<-CRYSTAL
         i = 0
         while i < 10
@@ -56,6 +56,24 @@ module Ameba::Rule::Lint
       expect_no_corrections source
     end
 
+    it "reports if there is a shared var in spawn (loop)" do
+      source = expect_issue subject, <<-CRYSTAL
+        i = 0
+        loop do
+          break if i >= 10
+          spawn do
+            puts(i)
+               # ^ error: Shared variable `i` is used in fiber
+          end
+          i += 1
+        end
+
+        Fiber.yield
+        CRYSTAL
+
+      expect_no_corrections source
+    end
+
     it "reports reassigned reference to shared var in spawn" do
       source = expect_issue subject, <<-CRYSTAL
         channel = Channel(String).new

spec/spec_helper.cr

@@ -259,6 +259,7 @@ module Ameba
       Crystal::MacroLiteral,
       Crystal::Expressions,
       Crystal::ControlExpression,
+      Crystal::Call,
     }
 
     def initialize(node)

src/ameba/ast/branch.cr

@@ -1,3 +1,5 @@
+require "./util"
+
 module Ameba::AST
   # Represents the branch in Crystal code.
   # Branch is a part of a branchable statement.
@@ -67,6 +69,8 @@ module Ameba::AST
 
     # :nodoc:
     private class BranchVisitor < Crystal::Visitor
+      include Util
+
       @current_branch : Crystal::ASTNode?
 
       property branchable : Branchable?
@@ -79,7 +83,7 @@ module Ameba::AST
         on_branchable_start(node, branches)
       end
 
-      private def on_branchable_start(node, branches : Array | Tuple)
+      private def on_branchable_start(node, branches : Enumerable)
         @branchable = Branchable.new(node, @branchable)
 
         branches.each do |branch_node|
@@ -172,6 +176,18 @@ module Ameba::AST
       def end_visit(node : Crystal::MacroFor)
         on_branchable_end node
       end
+
+      def visit(node : Crystal::Call)
+        if loop?(node) && (block = node.block)
+          on_branchable_start node, block.body
+        end
+      end
+
+      def end_visit(node : Crystal::Call)
+        if loop?(node) && node.block
+          on_branchable_end node
+        end
+      end
     end
   end
 end