1 files changed, 84 insertions, 0 deletions
diff --git a/sample/prism/locate_nodes.rb b/sample/prism/locate_nodes.rb
new file mode 100644
index 0000000000..7a51db4367
--- /dev/null
+++ b/sample/prism/locate_nodes.rb
@@ -0,0 +1,84 @@
+# This script locates a set of nodes determined by a line and column (in bytes).
+
+require "prism"
+require "pp"
+
+# This method determines if the given location covers the given line and column.
+# It's important to note that columns (and offsets) in prism are always in
+# bytes. This is because prism supports all 90 source encodings that Ruby
+# supports. You can always retrieve the column (or offset) of a location in
+# other units with other provided APIs, like #start_character_column or
+# #start_code_units_column.
+def covers?(location, line:, column:)
+  start_line = location.start_line
+  end_line = location.end_line
+
+  if start_line == end_line
+    # If the location only spans one line, then we only check if the line
+    # matches and that the column is covered by the column range.
+    line == start_line && (location.start_column...location.end_column).cover?(column)
+  else
+    # Otherwise, we check that it is on the start line and the column is greater
+    # than or equal to the start column, or that it is on the end line and the
+    # column is less than the end column, or that it is between the start and
+    # end lines.
+    (line == start_line && column >= location.start_column) ||
+      (line == end_line && column < location.end_column) ||
+      (line > start_line && line < end_line)
+  end
+end
+
+# This method descends down into the AST whose root is `node` and returns the
+# array of all of the nodes that cover the given line and column.
+def locate(node, line:, column:)
+  queue = [node]
+  result = []
+
+  # We could use a recursive method here instead if we wanted, but it's
+  # important to note that that will not work for ASTs that are nested deeply
+  # enough to cause a stack overflow.
+  while (node = queue.shift)
+    result << node
+
+    # Nodes have `child_nodes` and `compact_child_nodes`. `child_nodes` have
+    # consistent indices but include `nil` for optional fields that are not
+    # present, whereas `compact_child_nodes` has inconsistent indices but does
+    # not include `nil` for optional fields that are not present.
+    node.compact_child_nodes.find do |child|
+      queue << child if covers?(child.location, line: line, column: column)
+    end
+  end
+
+  result
+end
+
+result = Prism.parse_stream(DATA)
+locate(result.value, line: 4, column: 14).each_with_index do |node, index|
+  print " " * index
+  print node.class.name.split("::", 2).last
+  print " "
+  puts PP.pp(node.location, +"")
+end
+
+# =>
+# ProgramNode (1,0)-(7,3)
+#  StatementsNode (1,0)-(7,3)
+#   ModuleNode (1,0)-(7,3)
+#    StatementsNode (2,2)-(6,5)
+#     ClassNode (2,2)-(6,5)
+#      StatementsNode (3,4)-(5,7)
+#       DefNode (3,4)-(5,7)
+#        StatementsNode (4,6)-(4,21)
+#         CallNode (4,6)-(4,21)
+#          CallNode (4,6)-(4,15)
+#           ArgumentsNode (4,12)-(4,15)
+#            IntegerNode (4,12)-(4,15)
+
+__END__
+module Foo
+  class Bar
+    def baz
+      111 + 222 + 333
+    end
+  end
+end