1 files changed, 80 insertions, 0 deletions
diff --git a/lib/syntax_suggest/visitor.rb b/lib/syntax_suggest/visitor.rb
new file mode 100644
index 0000000000..6e25f7239c
--- /dev/null
+++ b/lib/syntax_suggest/visitor.rb
@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+module SyntaxSuggest
+  # Walks the Prism AST to extract structural info that cannot be reliably determined from tokens
+  # alone.
+  #
+  # Such as the location of lines that must be logically joined so the search algorithm will
+  # treat them as one. Example:
+  #
+  #   source = <<~RUBY
+  #     User                        # 1
+  #       .where(name: "Earlopain") # 2
+  #       .first                    # 3
+  #   RUBY
+  #   ast, _tokens = Prism.parse_lex(source).value
+  #   visitor = Visitor.new
+  #   visitor.visit(ast)
+  #   visitor.consecutive_lines # => Set[2, 1]
+  #
+  # This output means that line 1 and line 2 need to be joined with their next line.
+  #
+  # And determining the location of "endless" method definitions. For example:
+  #
+  #   source = <<~RUBY
+  #     def cube(x)
+  #       x * x * x
+  #     end
+  #     def square(x) = x * x # 1
+  #   RUBY
+  #
+  #   ast, _tokens = Prism.parse_lex(source).value
+  #   visitor = Visitor.new
+  #   visitor.visit(ast)
+  #   visitor.endless_def_keyword_offsets # => Set[28]
+  class Visitor < Prism::Visitor
+    attr_reader :endless_def_keyword_offsets, :consecutive_lines
+
+    def initialize
+      @endless_def_keyword_offsets = Set.new
+      @consecutive_lines = Set.new
+    end
+
+    # Called by Prism::Visitor for every method-call node in the AST
+    # (e.g. `foo.bar`, `foo.bar.baz`).
+    def visit_call_node(node)
+      receiver_loc = node.receiver&.location
+      call_operator_loc = node.call_operator_loc
+      message_loc = node.message_loc
+      if receiver_loc && call_operator_loc && message_loc
+        # dot-leading (dot on the next line)
+        #   foo        # line 1 - consecutive
+        #     .bar     # line 2
+        if receiver_loc.end_line != call_operator_loc.start_line && call_operator_loc.start_line == message_loc.start_line
+          (receiver_loc.end_line..call_operator_loc.start_line - 1).each do |line|
+            @consecutive_lines << line
+          end
+        end
+
+        # dot-trailing (dot on the same line as the receiver)
+        #   foo.       # line 1 - consecutive
+        #     bar      # line 2
+        if receiver_loc.end_line == call_operator_loc.start_line && call_operator_loc.start_line != message_loc.start_line
+          (call_operator_loc.start_line..message_loc.start_line - 1).each do |line|
+            @consecutive_lines << line
+          end
+        end
+      end
+      super
+    end
+
+    # Called by Prism::Visitor for every `def` node in the AST.
+    # Records the keyword start location for endless method definitions
+    # like `def foo = 123`. These are valid without a matching `end`,
+    # so Token must exclude them when deciding if a line is a keyword.
+    def visit_def_node(node)
+      @endless_def_keyword_offsets << node.def_keyword_loc.start_offset if node.equal_loc
+      super
+    end
+  end
+end