Merge pull request #543 from Shopify/at-fix-deadcode-removal

Fix removal of def nodes with a rescue

lib/spoom/deadcode/remover.rb

@@ -564,7 +564,7 @@ def find(source, location, kind)
               raise ParserError, "Error while parsing #{location.file}: #{message}"
             end
 
-            visitor = new(location)
+            visitor = new(location, kind)
             visitor.visit(result.value)
 
             node = visitor.node
@@ -617,10 +617,11 @@ def node_match_kind?(node, kind)
         sig { returns(T::Array[Prism::Node]) }
         attr_reader :nodes_nesting
 
-        sig { params(location: Location).void }
-        def initialize(location)
+        sig { params(location: Location, kind: T.nilable(Definition::Kind)).void }
+        def initialize(location, kind)
           super()
           @location = location
+          @kind = kind
           @node = T.let(nil, T.nilable(Prism::Node))
           @nodes_nesting = T.let([], T::Array[Prism::Node])
         end
@@ -635,6 +636,9 @@ def visit(node)
             # We found the node we're looking for at `@location`
             @node = node
 
+            # The node we found matches the kind we're looking for, we can stop here
+            return if @kind && self.class.node_match_kind?(node, @kind)
+
             # There may be a more precise child inside the node that also matches `@location`, let's visit them
             @nodes_nesting << node
             super(node)

test/spoom/deadcode/remover_test.rb

@@ -1027,6 +1027,26 @@ def bar; end
         RB
       end
 
+      def test_deadcode_remover_removes_method_with_rescue
+        res = remove(<<~RB, "foo")
+          def bar; end
+
+          def foo
+            puts "foo"
+          rescue => error
+            puts error
+          end
+
+          def baz; end
+        RB
+
+        assert_equal(<<~RB, res)
+          def bar; end
+
+          def baz; end
+        RB
+      end
+
       def test_deadcode_remover_does_not_remove_singleton_class_if_more_than_one_node
         res = remove(<<~RB, "foo")
           class Foo