From 0fd5890738616bfac1383dad0003f17ed0b63ac2 Mon Sep 17 00:00:00 2001
From: "V. Elenhaupt" <3624712+veelenga@users.noreply.github.com>
Date: Thu, 22 Nov 2018 10:38:32 +0200
Subject: [PATCH] Extend UnreachableCode rule: handle control flow (#83)

---
 spec/ameba/ast/flow_expression_spec.cr        |  32 +-
 spec/ameba/ast/util_spec.cr                   | 237 +++++++++
 .../visitors/flow_expression_visitor_spec.cr  |   4 +-
 spec/ameba/ast/visitors/node_visitor_spec.cr  |  14 +-
 spec/ameba/rule/lint/unreachable_code_spec.cr | 462 +++++++++++++++++-
 spec/spec_helper.cr                           |   2 +-
 src/ameba/ast/branchable.cr                   |   2 +-
 src/ameba/ast/flow_expression.cr              |  56 +--
 src/ameba/ast/util.cr                         |  97 ++++
 .../ast/visitors/flow_expression_visitor.cr   |  71 +--
 src/ameba/rule/lint/unreachable_code.cr       |   4 +-
 11 files changed, 880 insertions(+), 101 deletions(-)

diff --git a/spec/ameba/ast/flow_expression_spec.cr b/spec/ameba/ast/flow_expression_spec.cr
index fe982463..65c2d22e 100644
--- a/spec/ameba/ast/flow_expression_spec.cr
+++ b/spec/ameba/ast/flow_expression_spec.cr
@@ -5,42 +5,37 @@ module Ameba::AST
     describe "#initialize" do
       it "creates a new flow expression" do
         node = as_node("return 22")
-        parent_node = as_node("def foo; return 22; end")
-        flow_expression = FlowExpression.new node, parent_node
+        flow_expression = FlowExpression.new node, false
         flow_expression.node.should_not be_nil
-        flow_expression.parent_node.should_not be_nil
+        flow_expression.in_loop?.should eq false
       end
 
       describe "#delegation" do
         it "delegates to_s to @node" do
           node = as_node("return 22")
-          parent_node = as_node("def foo; return 22; end")
-          flow_expression = FlowExpression.new node, parent_node
+          flow_expression = FlowExpression.new node, false
           flow_expression.to_s.should eq node.to_s
         end
 
         it "delegates location to @node" do
           node = as_node %(break if true)
-          parent_node = as_node("def foo; return 22 if true; end")
-          flow_expression = FlowExpression.new node, parent_node
+          flow_expression = FlowExpression.new node, false
           flow_expression.location.should eq node.location
         end
       end
 
-      describe "#find_unreachable_node" do
-        it "returns first unreachable node" do
+      describe "#unreachable_nodes" do
+        it "returns unreachable nodes" do
           nodes = as_nodes %(
             def foobar
               return
               a = 1
-              a + 1
+              a = 2
             end
           )
-          node = nodes.control_expression_nodes.first
-          assign_node = nodes.assign_nodes.first
-          def_node = nodes.def_nodes.first
-          flow_expression = FlowExpression.new node, def_node
-          flow_expression.find_unreachable_node.should eq assign_node
+          node = nodes.expressions_nodes.first
+          flow_expression = FlowExpression.new node, false
+          flow_expression.unreachable_nodes.should eq nodes.assign_nodes
         end
 
         it "returns nil if there is no unreachable node" do
@@ -50,10 +45,9 @@ module Ameba::AST
               return a
             end
           )
-          node = nodes.control_expression_nodes.first
-          def_node = nodes.def_nodes.first
-          flow_expression = FlowExpression.new node, def_node
-          flow_expression.find_unreachable_node.should eq nil
+          node = nodes.expressions_nodes.first
+          flow_expression = FlowExpression.new node, false
+          flow_expression.unreachable_nodes.empty?.should eq true
         end
       end
     end
diff --git a/spec/ameba/ast/util_spec.cr b/spec/ameba/ast/util_spec.cr
index 9deaf6c5..53e1f274 100644
--- a/spec/ameba/ast/util_spec.cr
+++ b/spec/ameba/ast/util_spec.cr
@@ -71,5 +71,242 @@ module Ameba::AST
         ])
       end
     end
+
+    describe "#flow_command?" do
+      it "returns true if this is return" do
+        node = as_node("return 22")
+        subject.flow_command?(node, false).should eq true
+      end
+
+      it "returns true if this is a break in a loop" do
+        node = as_node("break")
+        subject.flow_command?(node, true).should eq true
+      end
+
+      it "returns false if this is a break out of loop" do
+        node = as_node("break")
+        subject.flow_command?(node, false).should eq false
+      end
+
+      it "returns true if this is a next in a loop" do
+        node = as_node("next")
+        subject.flow_command?(node, true).should eq true
+      end
+
+      it "returns false if this is a next out of loop" do
+        node = as_node("next")
+        subject.flow_command?(node, false).should eq false
+      end
+
+      it "returns true if this is raise" do
+        node = as_node("raise e")
+        subject.flow_command?(node, false).should eq true
+      end
+
+      it "returns true if this is exit" do
+        node = as_node("exit")
+        subject.flow_command?(node, false).should eq true
+      end
+
+      it "returns true if this is abort" do
+        node = as_node("abort")
+        subject.flow_command?(node, false).should eq true
+      end
+
+      it "returns false otherwise" do
+        node = as_node("foobar")
+        subject.flow_command?(node, false).should eq false
+      end
+    end
+
+    describe "#flow_expression?" do
+      it "returns true if this is a flow command" do
+        node = as_node("return")
+        subject.flow_expression?(node, true).should eq true
+      end
+
+      it "returns true if this is if-else consumed by flow expressions" do
+        node = as_node %(
+          if foo
+            return :foo
+          else
+            return :bar
+          end
+        )
+        subject.flow_expression?(node, false).should eq true
+      end
+
+      it "returns true if this is unless-else consumed by flow expressions" do
+        node = as_node %(
+          unless foo
+            return :foo
+          else
+            return :bar
+          end
+        )
+        subject.flow_expression?(node).should eq true
+      end
+
+      it "returns true if this is case consumed by flow expressions" do
+        node = as_node %(
+          case
+          when 1
+            return 1
+          when 2
+            return 2
+          else
+            return 3
+          end
+        )
+        subject.flow_expression?(node).should eq true
+      end
+
+      it "returns true if this is exception handler consumed by flow expressions" do
+        node = as_node %(
+          begin
+            raise "exp"
+          rescue e
+            return e
+          end
+        )
+        subject.flow_expression?(node).should eq true
+      end
+
+      it "returns true if this while consumed by flow expressions" do
+        node = as_node %(
+          while true
+            return
+          end
+        )
+        subject.flow_expression?(node).should eq true
+      end
+
+      it "returns false if this while with break" do
+        node = as_node %(
+          while true
+            break
+          end
+        )
+        subject.flow_expression?(node).should eq false
+      end
+
+      it "returns true if this until consumed by flow expressions" do
+        node = as_node %(
+          until false
+            return
+          end
+        )
+        subject.flow_expression?(node).should eq true
+      end
+
+      it "returns false if this until with break" do
+        node = as_node %(
+          until false
+            break
+          end
+        )
+        subject.flow_expression?(node).should eq false
+      end
+
+      it "returns true if this expressions consumed by flow expressions" do
+        node = as_node %(
+          exp1
+          exp2
+          return
+        )
+        subject.flow_expression?(node).should eq true
+      end
+
+      it "returns false otherwise" do
+        node = as_node %(
+          exp1
+          exp2
+        )
+        subject.flow_expression?(node).should eq false
+      end
+    end
+
+    describe "#raise?" do
+      it "returns true if this is a raise method call" do
+        node = as_node "raise e"
+        subject.raise?(node).should eq true
+      end
+
+      it "returns false if it has a receiver" do
+        node = as_node "obj.raise e"
+        subject.raise?(node).should eq false
+      end
+
+      it "returns false if size of the arguments doesn't match" do
+        node = as_node "raise"
+        subject.raise?(node).should eq false
+      end
+    end
+
+    describe "#exit?" do
+      it "returns true if this is a exit method call" do
+        node = as_node "exit"
+        subject.exit?(node).should eq true
+      end
+
+      it "returns true if this is a exit method call with one argument" do
+        node = as_node "exit 1"
+        subject.exit?(node).should eq true
+      end
+
+      it "returns false if it has a receiver" do
+        node = as_node "obj.exit"
+        subject.exit?(node).should eq false
+      end
+
+      it "returns false if size of the arguments doesn't match" do
+        node = as_node "exit 1, 1"
+        subject.exit?(node).should eq false
+      end
+    end
+
+    describe "#abort?" do
+      it "returns true if this is an abort method call" do
+        node = as_node "abort"
+        subject.abort?(node).should eq true
+      end
+
+      it "returns true if this is an abort method call with one argument" do
+        node = as_node "abort \"message\""
+        subject.abort?(node).should eq true
+      end
+
+      it "returns true if this is an abort method call with two arguments" do
+        node = as_node "abort \"message\", 1"
+        subject.abort?(node).should eq true
+      end
+
+      it "returns false if it has a receiver" do
+        node = as_node "obj.abort"
+        subject.abort?(node).should eq false
+      end
+
+      it "returns false if size of the arguments doesn't match" do
+        node = as_node "abort 1, 1, 1"
+        subject.abort?(node).should eq false
+      end
+    end
+
+    describe "#loop?" do
+      it "returns true if this is a loop method call" do
+        node = as_node "loop"
+        subject.loop?(node).should eq true
+      end
+
+      it "returns false if it has a receiver" do
+        node = as_node "obj.loop"
+        subject.loop?(node).should eq false
+      end
+
+      it "returns false if size of the arguments doesn't match" do
+        node = as_node "loop 1"
+        subject.loop?(node).should eq false
+      end
+    end
   end
 end
diff --git a/spec/ameba/ast/visitors/flow_expression_visitor_spec.cr b/spec/ameba/ast/visitors/flow_expression_visitor_spec.cr
index b8ce338a..2e41ce53 100644
--- a/spec/ameba/ast/visitors/flow_expression_visitor_spec.cr
+++ b/spec/ameba/ast/visitors/flow_expression_visitor_spec.cr
@@ -25,7 +25,7 @@ module Ameba::AST
           end
         end
       )
-      rule.expressions.size.should eq 2
+      rule.expressions.size.should eq 3
     end
 
     it "properly creates nested flow expressions" do
@@ -40,7 +40,7 @@ module Ameba::AST
           )
         end
       )
-      rule.expressions.size.should eq 3
+      rule.expressions.size.should eq 4
     end
 
     it "creates an expression for break" do
diff --git a/spec/ameba/ast/visitors/node_visitor_spec.cr b/spec/ameba/ast/visitors/node_visitor_spec.cr
index 4946adfe..38535a9c 100644
--- a/spec/ameba/ast/visitors/node_visitor_spec.cr
+++ b/spec/ameba/ast/visitors/node_visitor_spec.cr
@@ -5,14 +5,12 @@ module Ameba::AST
   source = Source.new ""
 
   describe NodeVisitor do
-    {% for name in NODES %}
-      describe "{{name}}" do
-        it "allow to visit {{name}} node" do
-          visitor = NodeVisitor.new rule, source
-          nodes = Crystal::Parser.new("").parse
-          nodes.accept visitor
-        end
+    describe "visit" do
+      it "allow to visit ASTNode" do
+        visitor = NodeVisitor.new rule, source
+        nodes = Crystal::Parser.new("").parse
+        nodes.accept visitor
       end
-    {% end %}
+    end
   end
 end
diff --git a/spec/ameba/rule/lint/unreachable_code_spec.cr b/spec/ameba/rule/lint/unreachable_code_spec.cr
index e3ef9a89..20b8909d 100644
--- a/spec/ameba/rule/lint/unreachable_code_spec.cr
+++ b/spec/ameba/rule/lint/unreachable_code_spec.cr
@@ -41,6 +41,23 @@ module Ameba::Rule::Lint
         subject.catch(s).should be_valid
       end
 
+      it "doesn't report if there is no else in if" do
+        s = Source.new %(
+          if a > 0
+            return :positive
+          end
+          :reachable
+        )
+        subject.catch(s).should be_valid
+      end
+
+      it "doesn't report return in on-line if" do
+        s = Source.new %(
+          return :positive if a > 0
+        )
+        subject.catch(s).should be_valid
+      end
+
       it "doesn't report if return is used in a block" do
         s = Source.new %(
           def foo
@@ -57,7 +74,7 @@ module Ameba::Rule::Lint
         subject.catch(s).should be_valid
       end
 
-      pending "reports if there is unreachable code after if-then-else" do
+      it "reports if there is unreachable code after if-then-else" do
         s = Source.new %(
           def foo
             if a > 0
@@ -71,7 +88,448 @@ module Ameba::Rule::Lint
         )
         subject.catch(s).should_not be_valid
         issue = s.issues.first
-        issue.location.to_s.should eq ":8:4"
+        issue.location.to_s.should eq ":8:3"
+      end
+
+      it "reports if there is unreachable code after if-then-else-if" do
+        s = Source.new %(
+          def foo
+            if a > 0
+              return :positive
+            elsif a != 0
+              return :negative
+            else
+              return :zero
+            end
+
+            :unreachable
+          end
+        )
+        subject.catch(s).should_not be_valid
+        issue = s.issues.first
+        issue.location.to_s.should eq ":10:3"
+      end
+
+      it "doesn't report if there is no unreachable code after if-then-else" do
+        s = Source.new %(
+          def foo
+            if a > 0
+              return :positive
+            else
+              return :negative
+            end
+          end
+        )
+        subject.catch(s).should be_valid
+      end
+
+      it "doesn't report if there is no unreachable in inner branch" do
+        s = Source.new %(
+          def foo
+            if a > 0
+              return :positive if a != 1
+            else
+              return :negative
+            end
+
+            :not_unreachable
+          end
+        )
+        subject.catch(s).should be_valid
+      end
+
+      it "doesn't report if there is no unreachable in exception handler" do
+        s = Source.new %(
+          def foo
+            puts :bar
+          rescue Exception
+            raise "Error!"
+          end
+        )
+        subject.catch(s).should be_valid
+      end
+
+      it "doesn't report if there is multiple conditions with return" do
+        s = Source.new %(
+          if :foo
+            if :bar
+              return :foobar
+            else
+              return :foobaz
+            end
+          elsif :fox
+            return :foofox
+          end
+
+          return :reachable
+        )
+        subject.catch(s).should be_valid
+      end
+
+      it "reports if there is unreachable code after unless" do
+        s = Source.new %(
+          unless :foo
+            return :bar
+          else
+            return :foo
+          end
+
+          :unreachable
+        )
+        subject.catch(s).should_not be_valid
+        issue = s.issues.first
+        issue.location.to_s.should eq ":7:1"
+      end
+
+      it "doesn't report if there is no unreachable code after unless" do
+        s = Source.new %(
+          unless :foo
+            return :bar
+          end
+
+          :reachable
+        )
+        subject.catch(s).should be_valid
+      end
+    end
+
+    context "binary op" do
+      it "reports unreachable code in a binary operator" do
+        s = Source.new %(
+          (return 22) && puts "a"
+        )
+        subject.catch(s).should_not be_valid
+
+        issue = s.issues.first
+        issue.location.to_s.should eq ":1:16"
+      end
+
+      it "reports unreachable code in inner binary operator" do
+        s = Source.new %(
+          do_something || (return 22) && puts "a"
+        )
+        subject.catch(s).should_not be_valid
+
+        issue = s.issues.first
+        issue.location.to_s.should eq ":1:32"
+      end
+
+      it "reports unreachable code after the binary op" do
+        s = Source.new %(
+          (return 22) && break
+          :unreachable
+        )
+        subject.catch(s).should_not be_valid
+
+        issue = s.issues.first
+        issue.location.to_s.should eq ":2:1"
+      end
+
+      it "doesn't report if return is not the right" do
+        s = Source.new %(
+          puts "a" && return
+        )
+        subject.catch(s).should be_valid
+      end
+
+      it "doesn't report unreachable code in multiple binary expressions" do
+        s = Source.new %(
+          foo || bar || baz
+        )
+        subject.catch(s).should be_valid
+      end
+    end
+
+    context "case" do
+      it "reports if there is unreachable code after case" do
+        s = Source.new %(
+          def foo
+            case cond
+            when 1
+              something
+              return
+            when 2
+              something2
+              return
+            else
+              something3
+              return
+            end
+            :unreachable
+          end
+        )
+        subject.catch(s).should_not be_valid
+        issue = s.issues.first
+        issue.location.to_s.should eq ":13:3"
+      end
+
+      it "doesn't report if case does not have else" do
+        s = Source.new %(
+          def foo
+            case cond
+            when 1
+              something
+              return
+            when 2
+              something2
+              return
+            end
+            :reachable
+          end
+        )
+        subject.catch(s).should be_valid
+      end
+
+      it "doesn't report if one when does not return" do
+        s = Source.new %(
+          def foo
+            case cond
+            when 1
+              something
+              return
+            when 2
+              something2
+            else
+              something3
+              return
+            end
+            :reachable
+          end
+        )
+        subject.catch(s).should be_valid
+      end
+    end
+
+    context "exception handler" do
+      it "reports unreachable code if it returns in body and rescues" do
+        s = Source.new %(
+          def foo
+            begin
+              return false
+            rescue Error
+              return false
+            rescue Exception
+              return false
+            end
+            :unreachable
+          end
+        )
+        subject.catch(s).should_not be_valid
+
+        issue = s.issues.first
+        issue.location.to_s.should eq ":9:3"
+      end
+
+      it "reports unreachable code if it returns in rescues and else" do
+        s = Source.new %(
+          def foo
+            begin
+              do_something
+            rescue Error
+              return :error
+            else
+              return true
+            end
+            :unreachable
+          end
+        )
+        subject.catch(s).should_not be_valid
+
+        issue = s.issues.first
+        issue.location.to_s.should eq ":9:3"
+      end
+
+      it "doesn't report if there is no else and ensure doesn't return" do
+        s = Source.new %(
+          def foo
+            begin
+              return false
+            rescue Error
+              puts "error"
+            rescue Exception
+              return false
+            end
+            :reachable
+          end
+        )
+        subject.catch(s).should be_valid
+      end
+
+      it "doesn't report if there is no else and body doesn't return" do
+        s = Source.new %(
+          def foo
+            begin
+              do_something
+            rescue Error
+              return true
+            rescue Exception
+              return false
+            end
+            :reachable
+          end
+        )
+        subject.catch(s).should be_valid
+      end
+
+      it "doesn't report if there is else and ensure doesn't return" do
+        s = Source.new %(
+          def foo
+            begin
+              do_something
+            rescue Error
+              puts "yo"
+            else
+              return true
+            end
+            :reachable
+          end
+        )
+        subject.catch(s).should be_valid
+      end
+
+      it "doesn't report if there is else and it doesn't return" do
+        s = Source.new %(
+          def foo
+            begin
+              do_something
+            rescue Error
+              return false
+            else
+              puts "yo"
+            end
+            :reachable
+          end
+        )
+        subject.catch(s).should be_valid
+      end
+
+      it "reports if there is unreachable code in rescue" do
+        s = Source.new %(
+          def method
+          rescue
+            return 22
+            :unreachable
+          end
+        )
+
+        subject.catch(s).should_not be_valid
+
+        issue = s.issues.first
+        issue.location.to_s.should eq ":4:3"
+      end
+    end
+
+    context "while/until" do
+      it "reports if there is unreachable code after while" do
+        s = Source.new %(
+          def method
+            while something
+              if :foo
+                return :foo
+              else
+                return :foobar
+              end
+            end
+            :unreachable
+          end
+        )
+
+        subject.catch(s).should_not be_valid
+
+        issue = s.issues.first
+        issue.location.to_s.should eq ":9:3"
+      end
+
+      it "reports if there is unreachable code after until" do
+        s = Source.new %(
+          def method
+            until something
+              if :foo
+                return :foo
+              else
+                return :foobar
+              end
+            end
+            :unreachable
+          end
+        )
+
+        subject.catch(s).should_not be_valid
+
+        issue = s.issues.first
+        issue.location.to_s.should eq ":9:3"
+      end
+
+      it "doesn't report if there is reachable code after while with break" do
+        s = Source.new %(
+          while something
+            break
+          end
+          :reachable
+        )
+
+        subject.catch(s).should be_valid
+      end
+    end
+
+    context "rescue" do
+      it "reports unreachable code in rescue" do
+        s = Source.new %(
+          begin
+
+          rescue e
+            raise e
+            :unreachable
+          end
+        )
+
+        subject.catch(s).should_not be_valid
+
+        issue = s.issues.first
+        issue.location.to_s.should eq ":5:3"
+      end
+
+      it "doesn't report if there is no unreachable code in rescue" do
+        s = Source.new %(
+          begin
+
+          rescue e
+            raise e
+          end
+        )
+
+        subject.catch(s).should be_valid
+      end
+    end
+
+    context "when" do
+      it "reports unreachable code in when" do
+        s = Source.new %(
+          case
+          when valid?
+            return 22
+            :unreachable
+          else
+
+          end
+        )
+
+        subject.catch(s).should_not be_valid
+
+        issue = s.issues.first
+        issue.location.to_s.should eq ":4:3"
+      end
+
+      it "doesn't report if there is no unreachable code in when" do
+        s = Source.new %(
+          case
+          when valid?
+            return 22
+          else
+          end
+        )
+
+        subject.catch(s).should be_valid
       end
     end
 
diff --git a/spec/spec_helper.cr b/spec/spec_helper.cr
index e99a87c0..e5a98b1a 100644
--- a/spec/spec_helper.cr
+++ b/spec/spec_helper.cr
@@ -129,7 +129,7 @@ module Ameba
       Crystal::If,
       Crystal::While,
       Crystal::MacroLiteral,
-      Crystal::ControlExpression,
+      Crystal::Expressions,
     ]
 
     def initialize(node)
diff --git a/src/ameba/ast/branchable.cr b/src/ameba/ast/branchable.cr
index 035cb46e..74cd91f3 100644
--- a/src/ameba/ast/branchable.cr
+++ b/src/ameba/ast/branchable.cr
@@ -1,7 +1,7 @@
 module Ameba::AST
   # A generic entity to represent a branchable Crystal node.
   # For example, `Crystal::If`, `Crystal::Unless`, `Crystal::While`
-  # are branchable.
+  # are branchables.
   #
   # ```
   # white a > 100 # Branchable A
diff --git a/src/ameba/ast/flow_expression.cr b/src/ameba/ast/flow_expression.cr
index 8c7bae16..344e6012 100644
--- a/src/ameba/ast/flow_expression.cr
+++ b/src/ameba/ast/flow_expression.cr
@@ -1,3 +1,5 @@
+require "./util"
+
 module Ameba::AST
   # Represents a flow expression in Crystal code.
   # For example,
@@ -14,12 +16,14 @@ module Ameba::AST
   # a parent node, which allows easily search through the related statement
   # (i.e. find unreachable code)
   class FlowExpression
+    include Util
+
+    # Is true only if some of the nodes parents is a loop.
+    getter? in_loop : Bool
+
     # The actual node of the flow expression.
     getter node : Crystal::ASTNode
 
-    # Parent ast node.
-    getter parent_node : Crystal::ASTNode
-
     delegate to_s, to: @node
     delegate location, to: @node
 
@@ -28,10 +32,10 @@ module Ameba::AST
     # ```
     # FlowExpression.new(node, parent_node)
     # ```
-    def initialize(@node, @parent_node)
+    def initialize(@node, @in_loop)
     end
 
-    # Returns first node which can't be reached because of a flow expression.
+    # Returns nodes which can't be reached because of a flow expression inside.
     # For example:
     #
     # ```
@@ -41,39 +45,21 @@ module Ameba::AST
     #
     #   a + 2 # => unreachable assign node
     # end
-    # ```
-    #
-    def find_unreachable_node
-      UnreachableNodeVisitor.new(node, parent_node)
-        .tap(&.accept parent_node)
-        .unreachable_nodes
-        .first?
-    end
+    def unreachable_nodes
+      unreachable_nodes = [] of Crystal::ASTNode
 
-    # :nodoc:
-    private class UnreachableNodeVisitor < Crystal::Visitor
-      getter unreachable_nodes = Array(Crystal::ASTNode).new
-      @after_control_flow_node = false
-      @branch : AST::Branch?
-
-      def initialize(@node : Crystal::ASTNode, parent_node)
-        @branch = Branch.of(@node, parent_node)
+      case current_node = node
+      when Crystal::Expressions
+        control_flow_found = false
+        current_node.expressions.each do |exp|
+          unreachable_nodes << exp if control_flow_found
+          control_flow_found ||= flow_expression?(exp, in_loop?)
+        end
+      when Crystal::BinaryOp
+        unreachable_nodes << current_node.right if flow_expression?(current_node.left, in_loop?)
       end
 
-      def visit(node : Crystal::ASTNode)
-        if node.class == @node.class &&
-           node.location == @node.location
-          @after_control_flow_node = true
-          return false
-        end
-
-        if @after_control_flow_node && !node.nop? && @branch.nil?
-          @unreachable_nodes << node
-          return false
-        end
-
-        true
-      end
+      unreachable_nodes
     end
   end
 end
diff --git a/src/ameba/ast/util.cr b/src/ameba/ast/util.cr
index d50d3551..b00b3a09 100644
--- a/src/ameba/ast/util.cr
+++ b/src/ameba/ast/util.cr
@@ -40,4 +40,101 @@ module Ameba::AST::Util
 
     node_lines
   end
+
+  # Returns true if node is a flow command, false - otherwise.
+  # Node represents a flow command if it is a control expression,
+  # or special call node that interrupts execution (i.e. raise, exit, abort).
+  def flow_command?(node, in_loop?)
+    case node
+    when Crystal::Return
+      true
+    when Crystal::Break, Crystal::Next
+      in_loop?
+    when Crystal::Call
+      raise?(node) || exit?(node) || abort?(node)
+    else
+      false
+    end
+  end
+
+  # Returns true if node is a flow expression, false if not.
+  # Node represents a flow expression if it is full-filed by a flow command.
+  #
+  # For example, this node is a flow expressions, because each branch contains
+  # a flow command `return`:
+  #
+  # ```
+  # if a > 0
+  #   return :positive
+  # elsif a < 0
+  #   return :negative
+  # else
+  #   return :zero
+  # end
+  # ```
+  #
+  # This node is a not a flow expression:
+  #
+  # ```
+  # if a > 0
+  #   return :positive
+  # end
+  # ```
+  #
+  # That's because not all branches return(i.e. `else` is missing).
+  #
+  def flow_expression?(node, in_loop? = false)
+    return true if flow_command? node, in_loop?
+
+    case node
+    when Crystal::If, Crystal::Unless
+      flow_expressions? [node.then, node.else], in_loop?
+    when Crystal::BinaryOp
+      flow_expression? node.left, in_loop?
+    when Crystal::Case
+      flow_expressions? [node.whens, node.else].flatten, in_loop?
+    when Crystal::ExceptionHandler
+      flow_expressions? [node.else || node.body, node.rescues].flatten, in_loop?
+    when Crystal::While, Crystal::Until
+      flow_expression? node.body, in_loop?
+    when Crystal::Rescue, Crystal::When
+      flow_expression? node.body, in_loop?
+    when Crystal::Expressions
+      node.expressions.any? { |exp| flow_expression? exp, in_loop? }
+    else
+      false
+    end
+  end
+
+  private def flow_expressions?(nodes, in_loop?)
+    nodes.all? { |exp| flow_expression? exp, in_loop? }
+  end
+
+  # Returns true if node represents `raise` method call.
+  def raise?(node)
+    node.is_a?(Crystal::Call) &&
+      node.name == "raise" && node.args.size == 1 && node.obj.nil?
+  end
+
+  # Returns true if node represents `exit` method call.
+  def exit?(node)
+    node.is_a?(Crystal::Call) &&
+      node.name == "exit" && node.args.size <= 1 && node.obj.nil?
+  end
+
+  # Returns true if node represents `abort` method call.
+  def abort?(node)
+    node.is_a?(Crystal::Call) &&
+      node.name == "abort" && node.args.size <= 2 && node.obj.nil?
+  end
+
+  # Returns true if node represents a loop.
+  def loop?(node)
+    case node
+    when Crystal::While, Crystal::Until
+      true
+    when Crystal::Call
+      node.name == "loop" && node.args.size == 0 && node.obj.nil?
+    end
+  end
 end
diff --git a/src/ameba/ast/visitors/flow_expression_visitor.cr b/src/ameba/ast/visitors/flow_expression_visitor.cr
index d632a70b..f75e0d64 100644
--- a/src/ameba/ast/visitors/flow_expression_visitor.cr
+++ b/src/ameba/ast/visitors/flow_expression_visitor.cr
@@ -1,60 +1,67 @@
+require "../util"
 require "./base_visitor"
 
 module Ameba::AST
   # AST Visitor that traverses all the flow expressions.
   class FlowExpressionVisitor < BaseVisitor
-    @node_stack = Array(Crystal::ASTNode).new
-    @flow_expression : FlowExpression?
+    include Util
 
+    @loop_stack = [] of Crystal::ASTNode
+
+    # Creates a new flow expression visitor.
     def initialize(@rule, @source)
       @source.ast.accept self
     end
 
+    # :nodoc:
     def visit(node)
-      @node_stack.push node
-    end
-
-    def end_visit(node)
-      if @flow_expression.nil?
-        @node_stack.pop unless @node_stack.empty?
-      else
-        @flow_expression = nil
+      if flow_expression?(node, in_loop?)
+        @rule.test @source, node, FlowExpression.new(node, in_loop?)
       end
-    end
-
-    def visit(node : Crystal::ControlExpression)
-      on_flow_expression_start(node)
 
       true
     end
 
+    # :nodoc:
+    def visit(node : Crystal::While)
+      on_loop_started(node)
+    end
+
+    # :nodoc:
+    def visit(node : Crystal::Until)
+      on_loop_started(node)
+    end
+
+    # :nodoc:
     def visit(node : Crystal::Call)
-      if raise?(node) || exit?(node) || abort?(node)
-        on_flow_expression_start(node)
-      else
-        @node_stack.push node
-      end
-
-      true
+      on_loop_started(node) if loop?(node)
     end
 
-    private def on_flow_expression_start(node)
-      if parent_node = @node_stack.last?
-        @flow_expression = FlowExpression.new(node, parent_node)
-        @rule.test @source, node, @flow_expression
-      end
+    # :nodoc:
+    def end_visit(node : Crystal::While)
+      on_loop_ended(node)
     end
 
-    private def raise?(node)
-      node.name == "raise" && node.args.size == 1 && node.obj.nil?
+    # :nodoc:
+    def end_visit(node : Crystal::Until)
+      on_loop_ended(node)
     end
 
-    private def exit?(node)
-      node.name == "exit" && node.args.size <= 1 && node.obj.nil?
+    # :nodoc:
+    def end_visit(node : Crystal::Call)
+      on_loop_ended(node) if loop?(node)
     end
 
-    private def abort?(node)
-      node.name == "abort" && node.args.size <= 2 && node.obj.nil?
+    private def on_loop_started(node)
+      @loop_stack.push(node)
+    end
+
+    private def on_loop_ended(node)
+      @loop_stack.pop?
+    end
+
+    private def in_loop?
+      @loop_stack.any?
     end
   end
 end
diff --git a/src/ameba/rule/lint/unreachable_code.cr b/src/ameba/rule/lint/unreachable_code.cr
index e069054d..ded09bf5 100644
--- a/src/ameba/rule/lint/unreachable_code.cr
+++ b/src/ameba/rule/lint/unreachable_code.cr
@@ -43,6 +43,8 @@ module Ameba::Rule::Lint
   # ```
   #
   struct UnreachableCode < Base
+    include AST::Util
+
     properties do
       description "Reports unreachable code"
     end
@@ -54,7 +56,7 @@ module Ameba::Rule::Lint
     end
 
     def test(source, node, flow_expression : AST::FlowExpression)
-      if unreachable_node = flow_expression.find_unreachable_node
+      if unreachable_node = flow_expression.unreachable_nodes.first?
         issue_for unreachable_node, MSG
       end
     end