Make BranchVisitor treat loop { … } calls as branchable

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/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