Add CommentAssociation

Author: soutaro

lib/rbs.rb

@@ -33,6 +33,7 @@
 require "rbs/ast/ruby/members"
 require "rbs/source"
 require "rbs/inline_parser"
+require "rbs/inline_parser/comment_association"
 require "rbs/environment"
 require "rbs/environment/use_map"
 require "rbs/environment/class_entry"

lib/rbs/inline_parser/comment_association.rb

@@ -0,0 +1,115 @@
+module RBS
+  class InlineParser
+    class CommentAssociation
+      attr_reader :blocks, :associated_blocks, :start_line_map, :end_line_map
+
+      def initialize(blocks)
+        @blocks = blocks.sort_by {|block| block.start_line }
+        @associated_blocks = Set[].compare_by_identity
+
+        @start_line_map = {}
+        @end_line_map = {}
+
+        blocks.each do |block|
+          if block.leading?
+            end_line_map[block.end_line] = block
+          else
+            start_line_map[block.start_line] = block
+          end
+        end
+      end
+
+      def self.build(buffer, result)
+        blocks = AST::Ruby::CommentBlock.build(buffer, result.comments)
+        new(blocks)
+      end
+
+      class Reference
+        attr_reader :block
+
+        def initialize(block, association)
+          @block = block
+          @associated_blocks = association
+        end
+
+        def associate!
+          @associated_blocks << block
+          self
+        end
+
+        def associated?
+          @associated_blocks.include?(block)
+        end
+      end
+
+      def leading_block(node)
+        start_line = node.location.start_line
+
+        if block = end_line_map.fetch(start_line - 1, nil)
+          Reference.new(block, associated_blocks)
+        end
+      end
+
+      def leading_block!(node)
+        if ref = leading_block(node)
+          unless ref.associated?
+            ref.associate!.block
+          end
+        end
+      end
+
+      def trailing_block(node)
+        location =
+          if node.is_a?(Prism::Node)
+            node.location
+          else
+            node
+          end #: Prism::Location
+        end_line = location.end_line
+        if block = start_line_map.fetch(end_line, nil)
+          Reference.new(block, associated_blocks)
+        end
+      end
+
+      def trailing_block!(node)
+        if ref = trailing_block(node)
+          unless ref.associated?
+            ref.associate!.block
+          end
+        end
+      end
+
+      def each_enclosed_block(node)
+        if block_given?
+          start_line = node.location.start_line
+          end_line = node.location.end_line
+
+          if start_line+1 < end_line
+            ((start_line + 1)...end_line).each do |line|
+              if block = end_line_map.fetch(line, nil)
+                unless associated_blocks.include?(block)
+                  associated_blocks << block
+                  yield block
+                end
+              end
+            end
+          end
+        else
+          enum_for :each_enclosed_block, node
+        end
+      end
+
+      def each_unassociated_block
+        if block_given?
+          blocks.each do |block|
+            unless associated_blocks.include?(block)
+              yield block
+            end
+          end
+        else
+          enum_for :each_unassociated_block
+        end
+      end
+    end
+  end
+end

sig/inline_parser/comment_association.rbs

@@ -0,0 +1,71 @@
+use RBS::AST::Ruby::CommentBlock
+
+module RBS
+  class InlineParser
+    # CommentAssociation manages the association between `Prism::Node` and `CommentBlock`
+    #
+    class CommentAssociation
+      attr_reader blocks: Array[CommentBlock]
+
+      attr_reader start_line_map: Hash[Integer, CommentBlock]
+
+      attr_reader end_line_map: Hash[Integer, CommentBlock]
+
+      # CommentBlocks that are already associated to a node, which cannot be associated to another node again
+      #
+      attr_reader associated_blocks: Set[CommentBlock]
+
+      def self.build: (Buffer, Prism::Result) -> instance
+
+      def initialize: (Array[CommentBlock]) -> void
+
+      class Reference
+        attr_reader block: CommentBlock
+
+        @associated_blocks: Set[CommentBlock]
+
+        def initialize: (CommentBlock, Set[CommentBlock]) -> void
+
+        def associate!: () -> self
+
+        def associated?: () -> bool
+      end
+
+      # Returns an unassociated CommentBlock that can be associated to given node
+      #
+      # Automatically updates association status.
+      #
+      def leading_block!: (Prism::Node) -> CommentBlock?
+
+      # Returns a Reference that is associated to given node
+      #
+      # Updates association explicitly through the reference.
+      #
+      def leading_block: (Prism::Node) -> Reference?
+
+      # Returns a CommentBlock that is associated to given node, or by its location
+      #
+      # Update association status.
+      #
+      def trailing_block!: (Prism::Node | Prism::Location) -> CommentBlock?
+
+      # Returns a Reference that is associated to given node, or by its location
+      #
+      # Updates association explicitly through the reference.
+      #
+      def trailing_block: (Prism::Node | Prism::Location) -> Reference?
+
+      # Yields leading CommentBlocks that is enclosed in the given node
+      #
+      # Note that `enclosed_blocks` works only after all of the *leading* blocks inside the node is associated.
+      #
+      # Update association status.
+      #
+      def each_enclosed_block: (Prism::Node) { (CommentBlock) -> void } -> void
+                             | (Prism::Node) -> Enumerator[CommentBlock]
+
+      def each_unassociated_block: () { (CommentBlock) -> void } -> void
+                                 | () -> Enumerator[CommentBlock]
+    end
+  end
+end