From 57898fd797e7b938cce1e7d48b43912cc4a0eb0b Mon Sep 17 00:00:00 2001
From: Sijawusz Pur Rahnama <sija@sija.pl>
Date: Thu, 28 Dec 2023 14:15:12 +0100
Subject: [PATCH] =?UTF-8?q?Make=20`BranchVisitor`=20treat=20`loop=20{=20?=
 =?UTF-8?q?=E2=80=A6=20}`=20calls=20as=20branchable?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 spec/ameba/ast/branch_spec.cr | 52 +++++++++++++++++++++++++++++++++++
 spec/spec_helper.cr           |  1 +
 src/ameba/ast/branch.cr       | 18 +++++++++++-
 3 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/spec/ameba/ast/branch_spec.cr b/spec/ameba/ast/branch_spec.cr
index 4da80472..3197eec4 100644
--- a/spec/ameba/ast/branch_spec.cr
+++ b/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
diff --git a/spec/spec_helper.cr b/spec/spec_helper.cr
index 094e0192..04dfa74c 100644
--- a/spec/spec_helper.cr
+++ b/spec/spec_helper.cr
@@ -259,6 +259,7 @@ module Ameba
       Crystal::MacroLiteral,
       Crystal::Expressions,
       Crystal::ControlExpression,
+      Crystal::Call,
     }
 
     def initialize(node)
diff --git a/src/ameba/ast/branch.cr b/src/ameba/ast/branch.cr
index 73ee8dab..30929880 100644
--- a/src/ameba/ast/branch.cr
+++ b/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