spoom 1.2.4 → 1.3.0

data/lib/spoom/deadcode/remover.rb

@@ -13,7 +13,7 @@ module Spoom
         @context = context
       end
 
-      sig { params(kind: Definition::Kind, location: Location).void }
+      sig { params(kind: T.nilable(Definition::Kind), location: Location).returns(String) }
       def remove_location(kind, location)
         file = location.file
 
@@ -23,7 +23,7 @@ module Spoom
 
         node_remover = NodeRemover.new(@context.read(file), kind, location)
         node_remover.apply_edit
-        @context.write!(file, node_remover.new_source)
+        node_remover.new_source
       end
 
       class NodeRemover
@@ -32,7 +32,7 @@ module Spoom
         sig { returns(String) }
         attr_reader :new_source
 
-        sig { params(source: String, kind: Definition::Kind, location: Location).void }
+        sig { params(source: String, kind: T.nilable(Definition::Kind), location: Location).void }
         def initialize(source, kind, location)
           @old_source = source
           @new_source = T.let(source.dup, String)
@@ -52,11 +52,15 @@ module Spoom
 
           node = @node_context.node
           case node
-          when SyntaxTree::ClassDeclaration, SyntaxTree::ModuleDeclaration, SyntaxTree::DefNode
+          when Prism::ClassNode, Prism::ModuleNode, Prism::DefNode
             delete_node_and_comments_and_sigs(@node_context)
-          when SyntaxTree::Const, SyntaxTree::ConstPathField
+          when Prism::ConstantWriteNode, Prism::ConstantOperatorWriteNode,
+                Prism::ConstantAndWriteNode, Prism::ConstantOrWriteNode,
+                Prism::ConstantPathWriteNode, Prism::ConstantPathOperatorWriteNode,
+                Prism::ConstantPathAndWriteNode, Prism::ConstantPathOrWriteNode,
+                Prism::ConstantTargetNode
             delete_constant_assignment(@node_context)
-          when SyntaxTree::SymbolLiteral # for attr accessors
+          when Prism::SymbolNode # for attr accessors
             delete_attr_accessor(@node_context)
           else
             raise Error, "Unsupported node type: #{node.class}"
@@ -67,19 +71,26 @@ module Spoom
 
         sig { params(context: NodeContext).void }
         def delete_constant_assignment(context)
-          # Pop the Varfield node from the nesting nodes
-          if context.node.is_a?(SyntaxTree::Const)
-            context = context.parent_context
+          case context.node
+          when Prism::ConstantWriteNode, Prism::ConstantOperatorWriteNode,
+               Prism::ConstantAndWriteNode, Prism::ConstantOrWriteNode,
+               Prism::ConstantPathWriteNode, Prism::ConstantPathOperatorWriteNode,
+               Prism::ConstantPathAndWriteNode, Prism::ConstantPathOrWriteNode
+            # Nesting node is an assign, it means only one constant is assigned on the line
+            # so we can remove the whole assign
+            delete_node_and_comments_and_sigs(context)
+            return
           end
 
+          # We're assigning multiple constants, we need to remove only the useless node
           parent_context = context.parent_context
           parent_node = parent_context.node
-          if parent_node.is_a?(SyntaxTree::Assign)
+          if parent_node.is_a?(Prism::ConstantWriteNode)
             # Nesting node is an assign, it means only one constant is assigned on the line
             # so we can remove the whole assign
             delete_node_and_comments_and_sigs(parent_context)
             return
-          elsif parent_node.is_a?(SyntaxTree::MLHS) && parent_node.parts.size == 1
+          elsif parent_node.is_a?(Prism::MultiWriteNode) && parent_node.lefts.size == 1
             # Nesting node is a single left hand side, it means only one constant is assigned
             # so we can remove the whole line
             delete_node_and_comments_and_sigs(parent_context.parent_context)
@@ -102,36 +113,36 @@ module Spoom
             # BAZ = 42
             # ~~~
             delete_lines(node.location.start_line, node.location.end_line)
-          elsif prev_node && next_node
+          elsif prev_node && next_node.is_a?(Prism::ConstantTargetNode)
             # We have a node before and after one the same line, just remove the part of the line
             #
             # ~~~
             # FOO, BAR, BAZ = 42 # we need to remove BAR
             # ~~~
-            replace_chars(prev_node.location.end_char, next_node.location.start_char, ", ")
+            replace_chars(prev_node.location.end_offset, next_node.location.start_offset, ", ")
           elsif prev_node
             # We have a node before, on the same line, but no node after, just remove the part of the line
             #
             # ~~~
             # FOO, BAR = 42 # we need to remove BAR
             # ~~~
-            nesting_context = parent_context.parent_context
-            nesting_assign = T.cast(nesting_context.node, T.any(SyntaxTree::MAssign, SyntaxTree::MLHSParen))
-            case nesting_assign
-            when SyntaxTree::MAssign
-              replace_chars(prev_node.location.end_char, nesting_assign.value.location.start_char, " = ")
-            when SyntaxTree::MLHSParen
-              nesting_context = nesting_context.parent_context
-              nesting_assign = T.cast(nesting_context.node, SyntaxTree::MAssign)
-              replace_chars(prev_node.location.end_char, nesting_assign.value.location.start_char, ") = ")
+            nesting_assign = T.cast(parent_context.node, Prism::MultiWriteNode)
+
+            rparen_loc = nesting_assign.rparen_loc
+            if rparen_loc
+              # We have an assign with parenthesis, we need to remove the part of the line until the closing parenthesis
+              delete_chars(prev_node.location.end_offset, rparen_loc.start_offset)
+            else
+              # We don't have a parenthesis, we need to remove the part of the line until the operator
+              replace_chars(prev_node.location.end_offset, nesting_assign.operator_loc.start_offset, " ")
             end
-          elsif next_node
+          elsif next_node.is_a?(Prism::ConstantTargetNode)
             # We don't have a node before but a node after on the same line, just remove the part of the line
             #
             # ~~~
             # FOO, BAR = 42 # we need to remove FOO
             # ~~~
-            delete_chars(node.location.start_char, next_node.location.start_char)
+            delete_chars(node.location.start_offset, next_node.location.start_offset)
           else
             # Should have been removed as a single MLHS node
             raise "Unexpected case while removing constant assignment"
@@ -142,10 +153,10 @@ module Spoom
         def delete_attr_accessor(context)
           args_context = context.parent_context
           send_context = args_context.parent_context
-          send_context = send_context.parent_context if send_context.node.is_a?(SyntaxTree::ArgParen)
+          send_context = send_context.parent_context if send_context.node.is_a?(Prism::ArgumentsNode)
 
-          send_node = T.cast(send_context.node, T.any(SyntaxTree::Command, SyntaxTree::CallNode))
-          need_accessor = context.node_string(send_node.message) == "attr_accessor"
+          send_node = T.cast(send_context.node, Prism::CallNode)
+          need_accessor = send_node.name == :attr_accessor
 
           if args_context.node.child_nodes.size == 1
             # Only one accessor is defined, we can remove the whole node
@@ -175,21 +186,21 @@ module Spoom
             # ~~~
             # attr_reader :foo, :bar, :baz # we need to remove bar
             # ~~~
-            replace_chars(prev_node.location.end_char, next_node.location.start_char, ", ")
+            replace_chars(prev_node.location.end_offset, next_node.location.start_offset, ", ")
           elsif prev_node
             # We have a node before, on the same line, but no node after, just remove the part of the line
             #
             # ~~~
             # attr_reader :foo, :bar, :baz # we need to remove baz
             # ~~~
-            delete_chars(prev_node.location.end_char, context.node.location.end_char)
+            delete_chars(prev_node.location.end_offset, context.node.location.end_offset)
           elsif next_node
             # We don't have a node before but a node after on the same line, just remove the part of the line
             #
             # ~~~
             # attr_reader :foo, :bar, :baz # we need to remove foo
             # ~~~
-            delete_chars(context.node.location.start_char, next_node.location.start_char)
+            delete_chars(context.node.location.start_offset, next_node.location.start_offset)
           else
             raise "Unexpected case while removing attr_accessor"
           end
@@ -199,13 +210,13 @@ module Spoom
 
         sig do
           params(
-            node: SyntaxTree::Node,
+            node: Prism::Node,
             send_context: NodeContext,
             was_removed: T::Boolean,
           ).void
         end
         def insert_accessor(node, send_context, was_removed:)
-          name = @node_context.node_string(node)
+          name = node.slice
           code = case @kind
           when Definition::Kind::AttrReader
             "attr_writer #{name}"
@@ -221,10 +232,10 @@ module Spoom
           node_after = send_context.next_node
 
           if was_removed
-            first_node = send_context.attached_comments_and_sigs.first || send_context.node
+            first_node = send_context.attached_sigs.first || send_context.node
             at_line = first_node.location.start_line - 1
 
-            prev_context = NodeContext.new(@old_source, first_node, send_context.nesting)
+            prev_context = NodeContext.new(@old_source, @node_context.comments, first_node, send_context.nesting)
             node_before = prev_context.previous_node
 
             new_line_before = node_before && send_context.node.location.start_line - node_before.location.end_line < 2
@@ -251,21 +262,51 @@ module Spoom
           start_line = context.node.location.start_line
           end_line = context.node.location.end_line
 
-          # Adjust the lines to remove to include the comments
-          nodes = context.attached_comments_and_sigs
-          if nodes.any?
-            start_line = T.must(nodes.first).location.start_line
+          # TODO: remove once Prism location are fixed
+          node = context.node
+          case node
+          when Prism::ConstantWriteNode, Prism::ConstantOperatorWriteNode,
+                Prism::ConstantAndWriteNode, Prism::ConstantOrWriteNode,
+                Prism::ConstantPathWriteNode, Prism::ConstantPathOperatorWriteNode,
+                Prism::ConstantPathAndWriteNode, Prism::ConstantPathOrWriteNode
+            value = node.value
+            if value.is_a?(Prism::StringNode)
+              end_line = value.closing_loc&.start_line || value.location.end_line
+            end
           end
 
+          # Adjust the lines to remove to include sigs attached to the node
+          first_node = context.attached_sigs.first || context.node
+          start_line = first_node.location.start_line if first_node
+
+          # Adjust the lines to remove to include comments attached to the node
+          first_comment = context.attached_comments(first_node).first
+          start_line = first_comment.location.start_line if first_comment
+
           # Adjust the lines to remove to include previous blank lines
-          prev_context = NodeContext.new(@old_source, nodes.first || context.node, context.nesting)
-          before = prev_context.previous_node
+          prev_context = NodeContext.new(@old_source, @node_context.comments, first_node, context.nesting)
+          before = T.let(prev_context.previous_node, T.nilable(T.any(Prism::Node, Prism::Comment)))
+
+          # There may be an unrelated comment between the current node and the one before
+          # if there is, we only want to delete lines up to the last comment found
+          if before
+            to_node = first_comment || node
+            comment = @node_context.comments_between_lines(before.location.end_line, to_node.location.start_line).last
+            before = comment if comment
+          end
+
           if before && before.location.end_line < start_line - 1
             # There is a node before and a blank line
             start_line = before.location.end_line + 1
-          elsif before.nil? && context.parent_node.location.start_line < start_line - 1
-            # There is no node before, but a blank line
-            start_line = context.parent_node.location.start_line + 1
+          elsif before.nil?
+            # There is no node before, check if there is a blank line
+            parent_context = context.parent_context
+            # With Prism the StatementsNode location starts at the first line of the first node
+            parent_context = parent_context.parent_context if parent_context.node.is_a?(Prism::StatementsNode)
+            if parent_context.node.location.start_line < start_line - 1
+              # There is a blank line before the node
+              start_line = parent_context.node.location.start_line + 1
+            end
           end
 
           # Adjust the lines to remove to include following blank lines
@@ -296,41 +337,24 @@ module Spoom
           @new_source[start_char...end_char] = replacement
         end
 
-        sig { params(line_number: Integer, start_column: Integer, end_column: Integer).void }
-        def delete_line_part(line_number, start_column, end_column)
-          lines = []
-          @new_source.lines.each_with_index do |line, index|
-            current_line = index + 1
-
-            lines << if line_number == current_line
-              T.must(line[0...start_column]) + T.must(line[end_column..-1])
-            else
-              line
-            end
-          end
-          @new_source = lines.join
-        end
-
-        sig { params(node: SyntaxTree::MethodAddBlock, name: String, kind: Definition::Kind).returns(String) }
+        sig { params(node: Prism::CallNode, name: String, kind: T.nilable(Definition::Kind)).returns(String) }
         def transform_sig(node, name:, kind:)
           type = T.let(nil, T.nilable(String))
 
-          statements = node.block.bodystmt
-          statements = statements.statements if statements.is_a?(SyntaxTree::BodyStmt)
+          block = T.cast(node.block, Prism::BlockNode)
+          statements = T.cast(block.body, Prism::StatementsNode)
 
           statements.body.each do |call|
-            next unless call.is_a?(SyntaxTree::CallNode)
-            next unless @node_context.node_string(call.message) == "returns"
+            next unless call.is_a?(Prism::CallNode)
+            next unless call.name == :returns
 
             args = call.arguments
-            args = args.arguments if args.is_a?(SyntaxTree::ArgParen)
-
-            next unless args.is_a?(SyntaxTree::Args)
+            next unless args
 
-            first = args.parts.first
+            first = args.arguments.first
             next unless first
 
-            type = @node_context.node_string(first)
+            type = first.slice
           end
 
           name = name.delete_prefix(":")
@@ -348,20 +372,31 @@ module Spoom
       class NodeContext
         extend T::Sig
 
-        sig { returns(SyntaxTree::Node) }
+        sig { returns(T::Hash[Integer, Prism::Comment]) }
+        attr_reader :comments
+
+        sig { returns(Prism::Node) }
         attr_reader :node
 
-        sig { returns(T::Array[SyntaxTree::Node]) }
+        sig { returns(T::Array[Prism::Node]) }
         attr_accessor :nesting
 
-        sig { params(source: String, node: SyntaxTree::Node, nesting: T::Array[SyntaxTree::Node]).void }
-        def initialize(source, node, nesting)
+        sig do
+          params(
+            source: String,
+            comments: T::Hash[Integer, Prism::Comment],
+            node: Prism::Node,
+            nesting: T::Array[Prism::Node],
+          ).void
+        end
+        def initialize(source, comments, node, nesting)
           @source = source
+          @comments = comments
           @node = node
           @nesting = nesting
         end
 
-        sig { returns(SyntaxTree::Node) }
+        sig { returns(Prism::Node) }
         def parent_node
           parent = @nesting.last
           raise "No parent for node #{node}" unless parent
@@ -375,100 +410,119 @@ module Spoom
           parent = nesting.pop
           raise "No parent context for node #{@node}" unless parent
 
-          NodeContext.new(@source, parent, nesting)
+          NodeContext.new(@source, @comments, parent, nesting)
         end
 
-        sig { returns(T::Array[SyntaxTree::Node]) }
+        sig { returns(T::Array[Prism::Node]) }
         def previous_nodes
           parent = parent_node
+          child_nodes = parent.child_nodes.compact
 
-          index = parent.child_nodes.index(@node)
+          index = child_nodes.index(@node)
           raise "Node #{@node} not found in parent #{parent}" unless index
 
-          parent.child_nodes[0...index].reject { |child| child.is_a?(SyntaxTree::VoidStmt) }
+          T.must(child_nodes[0...index])
         end
 
-        sig { returns(T.nilable(SyntaxTree::Node)) }
+        sig { returns(T.nilable(Prism::Node)) }
         def previous_node
           previous_nodes.last
         end
 
-        sig { returns(T::Array[SyntaxTree::Node]) }
+        sig { returns(T::Array[Prism::Node]) }
         def next_nodes
           parent = parent_node
+          child_nodes = parent.child_nodes.compact
 
-          index = parent.child_nodes.index(node)
+          index = child_nodes.index(node)
           raise "Node #{@node} not found in nesting node #{parent}" unless index
 
-          parent.child_nodes[(index + 1)..-1].reject { |node| node.is_a?(SyntaxTree::VoidStmt) }
+          T.must(child_nodes.compact[(index + 1)..-1])
         end
 
-        sig { returns(T.nilable(SyntaxTree::Node)) }
+        sig { returns(T.nilable(Prism::Node)) }
         def next_node
           next_nodes.first
         end
 
         sig { returns(T.nilable(NodeContext)) }
         def sclass_context
-          sclass = T.let(nil, T.nilable(SyntaxTree::SClass))
+          sclass = T.let(nil, T.nilable(Prism::SingletonClassNode))
 
           nesting = @nesting.dup
           until nesting.empty? || sclass
             node = nesting.pop
-            next unless node.is_a?(SyntaxTree::SClass)
+            next unless node.is_a?(Prism::SingletonClassNode)
 
             sclass = node
           end
 
-          return unless sclass.is_a?(SyntaxTree::SClass)
+          return unless sclass.is_a?(Prism::SingletonClassNode)
 
-          nodes = sclass.bodystmt.statements.body.reject do |node|
-            node.is_a?(SyntaxTree::VoidStmt) || node.is_a?(SyntaxTree::Comment) ||
-              sorbet_signature?(node) || sorbet_extend_sig?(node)
+          body = sclass.body
+          return NodeContext.new(@source, @comments, sclass, nesting) unless body.is_a?(Prism::StatementsNode)
+
+          nodes = body.child_nodes.reject do |node|
+            sorbet_signature?(node) || sorbet_extend_sig?(node)
           end
 
           if nodes.size <= 1
-            return NodeContext.new(@source, sclass, nesting)
+            return NodeContext.new(@source, @comments, sclass, nesting)
           end
 
           nil
         end
 
-        sig { params(node: T.nilable(SyntaxTree::Node)).returns(T::Boolean) }
+        sig { params(node: T.nilable(Prism::Node)).returns(T::Boolean) }
         def sorbet_signature?(node)
-          return false unless node.is_a?(SyntaxTree::MethodAddBlock)
+          node.is_a?(Prism::CallNode) && node.name == :sig
+        end
 
-          call = node.call
-          return false unless call.is_a?(SyntaxTree::CallNode)
+        sig { params(node: T.nilable(Prism::Node)).returns(T::Boolean) }
+        def sorbet_extend_sig?(node)
+          return false unless node.is_a?(Prism::CallNode)
+          return false unless node.name == :extend
 
-          ident = call.message
-          return false unless ident.is_a?(SyntaxTree::Ident)
+          args = node.arguments
+          return false unless args
+          return false unless args.arguments.size == 1
 
-          ident.value == "sig"
+          args.arguments.first&.slice == "T::Sig"
         end
 
-        sig { params(node: T.nilable(SyntaxTree::Node)).returns(T::Boolean) }
-        def sorbet_extend_sig?(node)
-          return false unless node.is_a?(SyntaxTree::Command)
-          return false unless node_string(node.message) == "extend"
-          return false unless node.arguments.parts.size == 1
+        sig { params(start_line: Integer, end_line: Integer).returns(T::Array[Prism::Comment]) }
+        def comments_between_lines(start_line, end_line)
+          comments = T.let([], T::Array[Prism::Comment])
 
-          node_string(T.must(node.arguments.parts.first)) == "T::Sig"
+          (start_line + 1).upto(end_line - 1) do |line|
+            comment = @comments[line]
+            comments << comment if comment
+          end
+
+          comments
         end
 
-        sig { params(comment: SyntaxTree::Node, node: SyntaxTree::Node).returns(T::Boolean) }
-        def comment_for_node?(comment, node)
-          return false unless comment.is_a?(SyntaxTree::Comment)
+        sig { params(node: Prism::Node).returns(T::Array[Prism::Comment]) }
+        def attached_comments(node)
+          comments = T.let([], T::Array[Prism::Comment])
+
+          start_line = node.location.start_line - 1
+          start_line.downto(1) do |line|
+            comment = @comments[line]
+            break unless comment
 
-          comment.location.end_line == node.location.start_line - 1
+            comments << comment
+          end
+
+          comments.reverse
         end
 
-        sig { returns(T::Array[SyntaxTree::Node]) }
-        def attached_comments_and_sigs
-          nodes = T.let([], T::Array[SyntaxTree::Node])
+        sig { returns(T::Array[Prism::Node]) }
+        def attached_sigs
+          nodes = T.let([], T::Array[Prism::Node])
 
           previous_nodes.reverse_each do |prev_node|
-            break unless comment_for_node?(prev_node, nodes.last || node) || sorbet_signature?(prev_node)
+            break unless sorbet_signature?(prev_node)
 
             nodes << prev_node
           end
@@ -476,13 +530,13 @@ module Spoom
           nodes.reverse
         end
 
-        sig { returns(T.nilable(SyntaxTree::MethodAddBlock)) }
+        sig { returns(T.nilable(Prism::CallNode)) }
         def attached_sig
           previous_nodes.reverse_each do |node|
-            if node.is_a?(SyntaxTree::Comment)
+            if node.is_a?(Prism::Comment)
               next
             elsif sorbet_signature?(node)
-              return T.cast(node, SyntaxTree::MethodAddBlock)
+              return T.cast(node, Prism::CallNode)
             else
               break
             end
@@ -490,79 +544,92 @@ module Spoom
 
           nil
         end
-
-        sig { params(node: T.any(Symbol, SyntaxTree::Node)).returns(String) }
-        def node_string(node)
-          case node
-          when Symbol
-            node.to_s
-          else
-            T.must(@source[node.location.start_char...node.location.end_char])
-          end
-        end
       end
 
-      class NodeFinder < SyntaxTree::Visitor
+      class NodeFinder < Visitor
         extend T::Sig
 
         class << self
           extend T::Sig
 
-          sig { params(source: String, location: Location, kind: Definition::Kind).returns(NodeContext) }
+          sig { params(source: String, location: Location, kind: T.nilable(Definition::Kind)).returns(NodeContext) }
           def find(source, location, kind)
-            tree = SyntaxTree.parse(source)
+            result = Prism.parse(source)
+
+            unless result.success?
+              message = result.errors.map do |e|
+                "#{e.message} (at #{e.location.start_line}:#{e.location.start_column})."
+              end.join(" ")
+
+              raise ParserError, "Error while parsing #{location.file}: #{message}"
+            end
 
             visitor = new(location)
-            visitor.visit(tree)
+            visitor.visit(result.value)
 
             node = visitor.node
             unless node
               raise Error, "Can't find node at #{location}"
             end
 
-            unless node_match_kind?(node, kind)
+            if kind && !node_match_kind?(node, kind)
               raise Error, "Can't find node at #{location}, expected #{kind} but got #{node.class}"
             end
 
-            NodeContext.new(source, node, visitor.nodes_nesting)
+            comments_by_line = T.let(
+              result.comments.to_h do |comment|
+                [comment.location.start_line, comment]
+              end,
+              T::Hash[Integer, Prism::Comment],
+            )
+
+            NodeContext.new(source, comments_by_line, node, visitor.nodes_nesting)
           end
 
-          sig { params(node: SyntaxTree::Node, kind: Definition::Kind).returns(T::Boolean) }
+          sig { params(node: Prism::Node, kind: Definition::Kind).returns(T::Boolean) }
           def node_match_kind?(node, kind)
             case kind
             when Definition::Kind::AttrReader, Definition::Kind::AttrWriter
-              node.is_a?(SyntaxTree::SymbolLiteral)
+              node.is_a?(Prism::SymbolNode)
             when Definition::Kind::Class
-              node.is_a?(SyntaxTree::ClassDeclaration)
+              node.is_a?(Prism::ClassNode)
             when Definition::Kind::Constant
-              node.is_a?(SyntaxTree::Const) || node.is_a?(SyntaxTree::ConstPathField)
+              node.is_a?(Prism::ConstantWriteNode) ||
+                node.is_a?(Prism::ConstantAndWriteNode) ||
+                node.is_a?(Prism::ConstantOrWriteNode) ||
+                node.is_a?(Prism::ConstantOperatorWriteNode) ||
+                node.is_a?(Prism::ConstantPathWriteNode) ||
+                node.is_a?(Prism::ConstantPathAndWriteNode) ||
+                node.is_a?(Prism::ConstantPathOrWriteNode) ||
+                node.is_a?(Prism::ConstantPathOperatorWriteNode) ||
+                node.is_a?(Prism::ConstantTargetNode)
             when Definition::Kind::Method
-              node.is_a?(SyntaxTree::DefNode)
+              node.is_a?(Prism::DefNode)
             when Definition::Kind::Module
-              node.is_a?(SyntaxTree::ModuleDeclaration)
+              node.is_a?(Prism::ModuleNode)
             end
           end
         end
 
-        sig { returns(T.nilable(SyntaxTree::Node)) }
+        sig { returns(T.nilable(Prism::Node)) }
         attr_reader :node
 
-        sig { returns(T::Array[SyntaxTree::Node]) }
-        attr_accessor :nodes_nesting
+        sig { returns(T::Array[Prism::Node]) }
+        attr_reader :nodes_nesting
 
         sig { params(location: Location).void }
         def initialize(location)
           super()
           @location = location
-          @node = T.let(nil, T.nilable(SyntaxTree::Node))
-          @nodes_nesting = T.let([], T::Array[SyntaxTree::Node])
+          @node = T.let(nil, T.nilable(Prism::Node))
+          @nodes_nesting = T.let([], T::Array[Prism::Node])
         end
 
-        sig { override.params(node: T.nilable(SyntaxTree::Node)).void }
+        sig { override.params(node: T.nilable(Prism::Node)).void }
         def visit(node)
           return unless node
 
-          location = location_from_node(node)
+          location = Location.from_prism(@location.file, node.location)
 
           if location == @location
             # We found the node we're looking for at `@location`
@@ -578,38 +645,6 @@ module Spoom
             super(node)
           end
         end
-
-        private
-
-        # TODO: remove once SyntaxTree location are fixed
-        sig { params(node: SyntaxTree::Node).returns(Location) }
-        def location_from_node(node)
-          case node
-          when SyntaxTree::Program, SyntaxTree::BodyStmt
-            # Patch SyntaxTree node locations to use the one of their children
-            location_from_children(node, node.statements.body)
-          when SyntaxTree::Statements
-            # Patch SyntaxTree node locations to use the one of their children
-            location_from_children(node, node.body)
-          else
-            Location.from_syntax_tree(@location.file, node.location)
-          end
-        end
-
-        # TODO: remove once SyntaxTree location are fixed
-        sig { params(node: SyntaxTree::Node, nodes: T::Array[SyntaxTree::Node]).returns(Location) }
-        def location_from_children(node, nodes)
-          first = T.must(nodes.first)
-          last = T.must(nodes.last)
-
-          Location.new(
-            @location.file,
-            first.location.start_line,
-            first.location.start_column,
-            last.location.end_line,
-            last.location.end_column,
-          )
-        end
       end
     end
   end