Comments on index entries

lib/syntax_tree/index.rb

@@ -20,46 +20,128 @@ def initialize(line, column)
 
     # This entry represents a class definition using the class keyword.
     class ClassDefinition
-      attr_reader :nesting, :name, :location
+      attr_reader :nesting, :name, :location, :comments
 
-      def initialize(nesting, name, location)
+      def initialize(nesting, name, location, comments)
         @nesting = nesting
         @name = name
         @location = location
+        @comments = comments
       end
     end
 
     # This entry represents a module definition using the module keyword.
     class ModuleDefinition
-      attr_reader :nesting, :name, :location
+      attr_reader :nesting, :name, :location, :comments
 
-      def initialize(nesting, name, location)
+      def initialize(nesting, name, location, comments)
         @nesting = nesting
         @name = name
         @location = location
+        @comments = comments
       end
     end
 
     # This entry represents a method definition using the def keyword.
     class MethodDefinition
-      attr_reader :nesting, :name, :location
+      attr_reader :nesting, :name, :location, :comments
 
-      def initialize(nesting, name, location)
+      def initialize(nesting, name, location, comments)
         @nesting = nesting
         @name = name
         @location = location
+        @comments = comments
       end
     end
 
     # This entry represents a singleton method definition using the def keyword
     # with a specified target.
     class SingletonMethodDefinition
-      attr_reader :nesting, :name, :location
+      attr_reader :nesting, :name, :location, :comments
 
-      def initialize(nesting, name, location)
+      def initialize(nesting, name, location, comments)
         @nesting = nesting
         @name = name
         @location = location
+        @comments = comments
+      end
+    end
+
+    # When you're using the instruction sequence backend, this class is used to
+    # lazily parse comments out of the source code.
+    class FileComments
+      # We use the ripper library to pull out source comments.
+      class Parser < Ripper
+        attr_reader :comments
+
+        def initialize(*)
+          super
+          @comments = {}
+        end
+
+        def on_comment(value)
+          comments[lineno] = value.chomp
+        end
+      end
+
+      # This represents the Ruby source in the form of a file. When it needs to
+      # be read we'll read the file.
+      class FileSource
+        attr_reader :filepath
+
+        def initialize(filepath)
+          @filepath = filepath
+        end
+
+        def source
+          File.read(filepath)
+        end
+      end
+
+      # This represents the Ruby source in the form of a string. When it needs
+      # to be read the string is returned.
+      class StringSource
+        attr_reader :source
+
+        def initialize(source)
+          @source = source
+        end
+      end
+
+      attr_reader :source
+
+      def initialize(source)
+        @source = source
+      end
+
+      def comments
+        @comments ||= Parser.new(source.source).tap(&:parse).comments
+      end
+    end
+
+    # This class handles parsing comments from Ruby source code in the case that
+    # we use the instruction sequence backend. Because the instruction sequence
+    # backend doesn't provide comments (since they are dropped) we provide this
+    # interface to lazily parse them out.
+    class EntryComments
+      include Enumerable
+      attr_reader :file_comments, :location
+
+      def initialize(file_comments, location)
+        @file_comments = file_comments
+        @location = location
+      end
+
+      def each(&block)
+        line = location.line - 1
+        result = []
+
+        while line >= 0 && (comment = file_comments.comments[line])
+          result.unshift(comment)
+          line -= 1
+        end
+
+        result.each(&block)
       end
     end
 
@@ -74,16 +156,22 @@ class ISeqBackend
       VM_DEFINECLASS_FLAG_HAS_SUPERCLASS = 0x10
 
       def index(source)
-        index_iseq(RubyVM::InstructionSequence.compile(source).to_a)
+        index_iseq(
+          RubyVM::InstructionSequence.compile(source).to_a,
+          FileComments.new(FileComments::StringSource.new(source))
+        )
       end
 
       def index_file(filepath)
-        index_iseq(RubyVM::InstructionSequence.compile_file(filepath).to_a)
+        index_iseq(
+          RubyVM::InstructionSequence.compile_file(filepath).to_a,
+          FileComments.new(FileComments::FileSource.new(filepath))
+        )
       end
 
       private
 
-      def index_iseq(iseq)
+      def index_iseq(iseq, file_comments)
         results = []
         queue = [[iseq, []]]
 
@@ -106,11 +194,23 @@ def index_iseq(iseq)
               elsif flags & VM_DEFINECLASS_TYPE_MODULE > 0
                 code_location = class_iseq[4][:code_location]
                 location = Location.new(code_location[0], code_location[1])
-                results << ModuleDefinition.new(current_nesting, name, location)
+
+                results << ModuleDefinition.new(
+                  current_nesting,
+                  name,
+                  location,
+                  EntryComments.new(file_comments, location)
+                )
               else
                 code_location = class_iseq[4][:code_location]
                 location = Location.new(code_location[0], code_location[1])
-                results << ClassDefinition.new(current_nesting, name, location)
+
+                results << ClassDefinition.new(
+                  current_nesting,
+                  name,
+                  location,
+                  EntryComments.new(file_comments, location)
+                )
               end
 
               queue << [class_iseq, current_nesting + [name]]
@@ -122,14 +222,21 @@ def index_iseq(iseq)
               results << SingletonMethodDefinition.new(
                 current_nesting,
                 name,
-                location
+                location,
+                EntryComments.new(file_comments, location)
               )
             when :definesmethod
               _, name, method_iseq = insn
 
               code_location = method_iseq[4][:code_location]
               location = Location.new(code_location[0], code_location[1])
-              results << MethodDefinition.new(current_nesting, name, location)
+
+              results << MethodDefinition.new(
+                current_nesting,
+                name,
+                location,
+                EntryComments.new(file_comments, location)
+              )
             end
           end
         end
@@ -143,21 +250,27 @@ def index_iseq(iseq)
     # supported on all runtimes.
     class ParserBackend
       class IndexVisitor < Visitor
-        attr_reader :results, :nesting
+        attr_reader :results, :nesting, :statements
 
         def initialize
           @results = []
           @nesting = []
+          @statements = nil
         end
 
         def visit_class(node)
           name = visit(node.constant).to_sym
           location =
             Location.new(node.location.start_line, node.location.start_column)
 
-          results << ClassDefinition.new(nesting.dup, name, location)
-          nesting << name
+          results << ClassDefinition.new(
+            nesting.dup,
+            name,
+            location,
+            comments_for(node)
+          )
 
+          nesting << name
           super
           nesting.pop
         end
@@ -172,9 +285,19 @@ def visit_def(node)
             Location.new(node.location.start_line, node.location.start_column)
 
           results << if node.target.nil?
-            MethodDefinition.new(nesting.dup, name, location)
+            MethodDefinition.new(
+              nesting.dup,
+              name,
+              location,
+              comments_for(node)
+            )
           else
-            SingletonMethodDefinition.new(nesting.dup, name, location)
+            SingletonMethodDefinition.new(
+              nesting.dup,
+              name,
+              location,
+              comments_for(node)
+            )
           end
         end
 
@@ -183,9 +306,14 @@ def visit_module(node)
           location =
             Location.new(node.location.start_line, node.location.start_column)
 
-          results << ModuleDefinition.new(nesting.dup, name, location)
-          nesting << name
+          results << ModuleDefinition.new(
+            nesting.dup,
+            name,
+            location,
+            comments_for(node)
+          )
 
+          nesting << name
           super
           nesting.pop
         end
@@ -194,6 +322,30 @@ def visit_program(node)
           super
           results
         end
+
+        def visit_statements(node)
+          @statements = node
+          super
+        end
+
+        private
+
+        def comments_for(node)
+          comments = []
+
+          body = statements.body
+          line = node.location.start_line - 1
+          index = body.index(node) - 1
+
+          while index >= 0 && body[index].is_a?(Comment) &&
+                  (line - body[index].location.start_line < 2)
+            comments.unshift(body[index].value)
+            line = body[index].location.start_line
+            index -= 1
+          end
+
+          comments
+        end
       end
 
       def index(source)

test/index_test.rb

@@ -18,6 +18,13 @@ def test_module_nested
       end
     end
 
+    def test_module_comments
+      index_each("# comment1\n# comment2\nmodule Foo; end") do |entry|
+        assert_equal :Foo, entry.name
+        assert_equal ["# comment1", "# comment2"], entry.comments.to_a
+      end
+    end
+
     def test_class
       index_each("class Foo; end") do |entry|
         assert_equal :Foo, entry.name
@@ -32,6 +39,13 @@ def test_class_nested
       end
     end
 
+    def test_class_comments
+      index_each("# comment1\n# comment2\nclass Foo; end") do |entry|
+        assert_equal :Foo, entry.name
+        assert_equal ["# comment1", "# comment2"], entry.comments.to_a
+      end
+    end
+
     def test_method
       index_each("def foo; end") do |entry|
         assert_equal :foo, entry.name
@@ -46,6 +60,13 @@ def test_method_nested
       end
     end
 
+    def test_method_comments
+      index_each("# comment1\n# comment2\ndef foo; end") do |entry|
+        assert_equal :foo, entry.name
+        assert_equal ["# comment1", "# comment2"], entry.comments.to_a
+      end
+    end
+
     private
 
     def index_each(source)