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