Merge pull request #940 from Shopify/Alex/line-matching

Line-matching rewriter

Author: vinistock

lib/spoom/rbs.rb

@@ -69,7 +69,20 @@ def initialize(string, location)
     end
 
     class Annotation < Comment; end
-    class Signature < Comment; end
+
+    class Signature < Comment
+      # Locations of the `#|` continuation comment lines that make up a multiline signature,
+      # in addition to the `#:` line tracked by `location`.
+      #: Array[Prism::Location]
+      attr_reader :continuation_locations
+
+      #: (String, Prism::Location, ?continuation_locations: Array[Prism::Location]) -> void
+      def initialize(string, location, continuation_locations: [])
+        super(string, location)
+        @continuation_locations = continuation_locations
+      end
+    end
+
     class TypeAlias < Comment; end
 
     module ExtractRBSComments
@@ -99,16 +112,18 @@ def node_rbs_comments(node)
           elsif string.start_with?("#: ")
             string = string.delete_prefix("#:").strip
             location = comment.location
+            continuation_locations = [] #: Array[Prism::Location]
 
             continuation_comments.reverse_each do |continuation_comment|
               string = "#{string}#{continuation_comment.slice.delete_prefix("#|")}"
               location = location.join(continuation_comment.location)
+              continuation_locations << continuation_comment.location
             end
             continuation_comments.clear
 
             next if string.start_with?("type ")
 
-            res.signatures.prepend(Signature.new(string, location))
+            res.signatures.prepend(Signature.new(string, location, continuation_locations:))
           elsif string.start_with?("#|")
             continuation_comments << comment
           end

lib/spoom/sorbet/translate.rb

@@ -5,6 +5,7 @@
 
 require "spoom/source/rewriter"
 require "spoom/sorbet/translate/translator"
+require "spoom/sorbet/translate/validator"
 require "spoom/sorbet/translate/rbs_comments_to_sorbet_sigs"
 require "spoom/sorbet/translate/sorbet_assertions_to_rbs_comments"
 require "spoom/sorbet/translate/sorbet_sigs_to_rbs_comments"

lib/spoom/sorbet/translate/rbs_comments_to_sorbet_sigs.rb

@@ -54,3 +54,4 @@ def rewrite_if_needed(
 
 require "spoom/sorbet/translate/rbs_comments_to_sorbet_sigs/options"
 require "spoom/sorbet/translate/rbs_comments_to_sorbet_sigs/human_readable_translator"
+require "spoom/sorbet/translate/rbs_comments_to_sorbet_sigs/line_matching_translator"

lib/spoom/sorbet/translate/rbs_comments_to_sorbet_sigs/base_translator.rb

@@ -100,6 +100,8 @@ def visit_attr(node)
               location: "#{@file}:#{node.location.start_line}",
             )
 
+            known_annotations = nil #: Array[Spoom::RBS::Annotation]?
+
             signatures.each do |signature|
               attr_type = ::RBS::Parser.parse_type(signature.string)
               sig = RBI::Sig.new
@@ -118,17 +120,21 @@ def visit_attr(node)
 
               sig.return_type = @type_translator.translate(attr_type)
 
-              apply_member_annotations(comments.method_annotations, sig)
+              known_annotations = apply_member_annotations(comments.method_annotations, sig)
 
               @rewriter << Source::Replace.new(
                 signature.location.start_offset,
                 signature.location.end_offset,
-                sig.string(max_line_length: @max_line_length),
+                pad_out_line_count(of: sig.string(max_line_length: @max_line_length), to_height_of: signature),
               )
             rescue ::RBS::ParsingError, ::RBI::Error
               # Ignore signatures with errors
               next
             end
+
+            if known_annotations
+              rewrite_member_annotations(comments.method_annotations, known: known_annotations)
+            end
           end
 
           #: (Prism::DefNode, Spoom::RBS::Comments) -> void
@@ -146,6 +152,8 @@ def rewrite_def(def_node, comments)
             builder.visit(def_node)
             rbi_node = builder.tree.nodes.first #: as RBI::Method
 
+            known_annotations = nil #: Array[Spoom::RBS::Annotation]?
+
             signatures.each do |signature|
               begin
                 method_type = ::RBS::Parser.parse_method_type(signature.string)
@@ -163,7 +171,7 @@ def rewrite_def(def_node, comments)
 
               sig = translator.result
 
-              apply_member_annotations(comments.method_annotations, sig)
+              known_annotations = apply_member_annotations(comments.method_annotations, sig)
 
               # Sorbet runtime doesn't support `sig` on `method_added` or
               # `singleton_method_added`, so we always use `without_runtime` for them.
@@ -174,9 +182,13 @@ def rewrite_def(def_node, comments)
               @rewriter << Source::Replace.new(
                 signature.location.start_offset,
                 signature.location.end_offset,
-                sig.string(max_line_length: @max_line_length),
+                pad_out_line_count(of: sig.string(max_line_length: @max_line_length), to_height_of: signature),
               )
             end
+
+            if known_annotations
+              rewrite_member_annotations(comments.method_annotations, known: known_annotations)
+            end
           end
 
           #: (Array[Spoom::RBS::Signature], method_name: String, location: String) -> Array[Spoom::RBS::Signature]
@@ -187,27 +199,26 @@ def apply_overloads_strategy(signatures, method_name:, location:)
             when :translate_all
               signatures
             when :translate_last
-              kept = signatures.last #: as Spoom::RBS::Signature
               others = signatures[0...-1] #: as !nil
+              others.each { |signature| rewrite_discarded_overload(signature) }
 
-              # Delete all the signatures we didn't keep
-              others.each do |signature|
-                from = adjust_to_line_start(signature.location.start_offset)
-                to = adjust_to_line_end(signature.location.end_offset)
-                @rewriter << Source::Delete.new(from, to)
-              end
+              kept = signatures.last #: as Spoom::RBS::Signature
               [kept]
             else # :raise
               raise Error, "Method `#{method_name}` at #{location} has multiple overloaded signatures"
             end
           end
 
+          # Called for every overloaded method sig that we discard because it wasn't the last one.
+          # @abstract
+          #: (Spoom::RBS::Signature) -> void
+          def rewrite_discarded_overload(signature) = raise
+
           #: (PrismTypes::anyScopeNode) -> void
           def apply_class_annotations(node)
             comments = node_rbs_comments(node)
             return if comments.empty?
 
-            indent = " " * (node.location.start_column + 2)
             insert_pos = case node
             when Prism::ClassNode
               (node.superclass || node.constant_path).location.end_offset
@@ -217,16 +228,14 @@ def apply_class_annotations(node)
               node.expression.location.end_offset
             end
 
-            class_annotations = comments.class_annotations
-            if class_annotations.any?
+            # Only translate (and `extend T::Helpers`) when there's at least one *known* class
+            # annotation. A node with only unknown annotations (e.g. `@private`) is left untouched.
+            if comments.class_annotations.any?
               unless already_extends?(node, /^(::)?T::Helpers$/)
-                @rewriter << Source::Insert.new(insert_pos, "\n#{indent}extend T::Helpers\n")
+                extend_with("T::Helpers", into: node, at: insert_pos)
               end
 
-              class_annotations.reverse_each do |annotation|
-                from = adjust_to_line_start(annotation.location.start_offset)
-                to = adjust_to_line_end(annotation.location.end_offset)
-
+              comments.annotations.reverse_each do |annotation|
                 content = case annotation.string
                 when "@abstract"
                   "abstract!"
@@ -241,15 +250,13 @@ def apply_class_annotations(node)
                   rbs_type = @type_translator.translate(srb_type)
                   "requires_ancestor { #{rbs_type} }"
                 else
+                  apply_class_annotation(annotation, parent_node: node, insert_pos:, sorbet_replacement: nil)
                   next
                 end
 
-                @rewriter << Source::Delete.new(from, to)
-
-                newline = node.body.nil? ? "" : "\n"
-                @rewriter << Source::Insert.new(insert_pos, "\n#{indent}#{content}#{newline}")
+                apply_class_annotation(annotation, parent_node: node, insert_pos:, sorbet_replacement: content)
               rescue ::RBS::ParsingError, ::RBI::Error
-                # Ignore annotations with errors
+                apply_class_annotation(annotation, parent_node: node, insert_pos:, sorbet_replacement: nil)
                 next
               end
             end
@@ -261,14 +268,11 @@ def apply_class_annotations(node)
                 next unless signature.string.start_with?("[")
 
                 type_params = ::RBS::Parser.parse_type_params(signature.string)
+                rewrite_type_params_signature(signature, type_params:)
                 next if type_params.empty?
 
-                from = adjust_to_line_start(signature.location.start_offset)
-                to = adjust_to_line_end(signature.location.end_offset)
-                @rewriter << Source::Delete.new(from, to)
-
                 unless already_extends?(node, /^(::)?T::Generic$/)
-                  @rewriter << Source::Insert.new(insert_pos, "\n#{indent}extend T::Generic\n")
+                  extend_with("T::Generic", into: node, at: insert_pos)
                 end
 
                 type_params.each do |type_param|
@@ -293,8 +297,7 @@ def apply_class_annotations(node)
                     end
                   end
 
-                  newline = node.body.nil? ? "" : "\n"
-                  @rewriter << Source::Insert.new(insert_pos, "\n#{indent}#{type_member}#{newline}")
+                  insert_type_member(type_member, parent_node: node, insert_pos:)
                 rescue ::RBS::ParsingError, ::RBI::Error
                   # Ignore signatures with errors
                   next
@@ -303,8 +306,31 @@ def apply_class_annotations(node)
             end
           end
 
-          #: (Array[Spoom::RBS::Annotation], RBI::Sig) -> void
+          # @param is_known: true if this is an RBS annotation that we recognize
+          #                  false for some other `@`-prefixed thing, like a documentation `@param` tag.
+          # @abstract
+          #: (
+          #|   Spoom::RBS::Annotation,
+          #|   parent_node: PrismTypes::anyScopeNode,
+          #|   insert_pos: Integer,
+          #|   sorbet_replacement: String?
+          #| ) -> void
+          def apply_class_annotation(annotation, parent_node:, insert_pos:, sorbet_replacement:) = raise
+
+          # Rewrites the `#: [...]` type params comment (e.g. delete it, or mark it as translated).
+          # @abstract
+          #: (Spoom::RBS::Signature, type_params: Array[::RBS::AST::TypeParam]) -> void
+          def rewrite_type_params_signature(signature, type_params:) = raise
+
+          # Inserts a single `type_member` declaration into the class/module body.
+          # @abstract
+          #: (String type_member, parent_node: PrismTypes::anyScopeNode, insert_pos: Integer) -> void
+          def insert_type_member(type_member, parent_node:, insert_pos:) = raise
+
+          #: (Array[Spoom::RBS::Annotation], RBI::Sig) -> Array[Spoom::RBS::Annotation]
           def apply_member_annotations(annotations, sig)
+            known = [] #: Array[Spoom::RBS::Annotation]
+
             annotations.each do |annotation|
               case annotation.string
               when "@abstract"
@@ -323,10 +349,37 @@ def apply_member_annotations(annotations, sig)
                 sig.is_overridable = true
               when "@without_runtime"
                 sig.without_runtime = true
+              else
+                next
               end
+
+              known << annotation
             end
+
+            known
           end
 
+          # Rewrites the member annotation comments in the source. Called once per method,
+          # regardless of how many overloaded signatures share the annotations, to avoid
+          # emitting duplicate markers.
+          #
+          #: (Array[Spoom::RBS::Annotation], known: Array[Spoom::RBS::Annotation]) -> void
+          def rewrite_member_annotations(annotations, known:)
+            annotations.each do |annotation|
+              rewrite_annotation(annotation, is_known: known.include?(annotation))
+            end
+          end
+
+          # @param is_known: true if this is an RBS annotation that we recognize
+          #                  false for some other `@`-prefixed thing, like a documentation `@param` tag.
+          # @overridable
+          #: (Spoom::RBS::Annotation, is_known: bool) -> void
+          def rewrite_annotation(annotation, is_known:) = nil # no-op
+
+          # @abstract
+          #: (String mixin_name, into: PrismTypes::anyScopeNode, at: Integer) -> void
+          def extend_with(mixin_name, into:, at:) = raise
+
           #: (PrismTypes::anyScopeNode, Regexp) -> bool
           def already_extends?(node, constant_regex)
             node.child_nodes.any? do |c|
@@ -408,12 +461,22 @@ def apply_type_aliases(comments)
 
               @rewriter << Source::Delete.new(from, to)
               content = "#{indent}#{alias_name} = T.type_alias { #{sorbet_type.to_rbi} }\n"
+              content = pad_out_line_count(of: content, to_height_of: type_alias)
               @rewriter << Source::Insert.new(insert_pos, content)
             rescue ::RBS::ParsingError, ::RBI::Error
               # Ignore type aliases with errors
               next
             end
           end
+
+          # @overridable
+          #: (of: String, to_height_of: Spoom::RBS::Comment) -> String
+          def pad_out_line_count(of:, to_height_of:)
+            replacement = of
+
+            # no-op implementation
+            replacement
+          end
         end
       end
     end

lib/spoom/sorbet/translate/rbs_comments_to_sorbet_sigs/human_readable_translator.rb

@@ -9,6 +9,62 @@ module Sorbet
     module Translate
       module RBSCommentsToSorbetSigs
         class HumanReadableTranslator < BaseTranslator
+          private
+
+          # Deletes the discarded overload from the source codes
+          # @override
+          #: (Spoom::RBS::Signature) -> void
+          def rewrite_discarded_overload(signature)
+            from = adjust_to_line_start(signature.location.start_offset)
+            to = adjust_to_line_end(signature.location.end_offset)
+            @rewriter << Source::Delete.new(from, to)
+          end
+
+          # @override
+          #: (
+          #|   Spoom::RBS::Annotation,
+          #|   parent_node: PrismTypes::anyScopeNode,
+          #|   insert_pos: Integer,
+          #|   sorbet_replacement: String?
+          #| ) -> void
+          def apply_class_annotation(annotation, parent_node:, insert_pos:, sorbet_replacement:)
+            return unless sorbet_replacement # unknown annotation.
+
+            from = adjust_to_line_start(annotation.location.start_offset)
+            to = adjust_to_line_end(annotation.location.end_offset)
+
+            @rewriter << Source::Delete.new(from, to)
+
+            indent = " " * (parent_node.location.start_column + 2)
+            newline = parent_node.body.nil? ? "" : "\n"
+            @rewriter << Source::Insert.new(insert_pos, "\n#{indent}#{sorbet_replacement}#{newline}")
+          end
+
+          # @override
+          #: (Spoom::RBS::Signature, type_params: Array[::RBS::AST::TypeParam]) -> void
+          def rewrite_type_params_signature(signature, type_params:)
+            from = adjust_to_line_start(signature.location.start_offset)
+            to = adjust_to_line_end(signature.location.end_offset)
+            @rewriter << Source::Delete.new(from, to)
+          end
+
+          # @override
+          #: (String type_member, parent_node: PrismTypes::anyScopeNode, insert_pos: Integer) -> void
+          def insert_type_member(type_member, parent_node:, insert_pos:)
+            indent = " " * (parent_node.location.start_column + 2)
+            newline = parent_node.body.nil? ? "" : "\n"
+            @rewriter << Source::Insert.new(insert_pos, "\n#{indent}#{type_member}#{newline}")
+          end
+
+          # @override
+          #: (String mixin_name, into: Prism::Node, at: Integer) -> void
+          def extend_with(mixin_name, into:, at:)
+            indent = " " * (into.location.start_column + 2)
+            # `extend` is always followed by an annotation or `type_member`, so it always needs a
+            # trailing newline to separate them. Since it's never the last inserted line, that
+            # trailing newline can't leave a blank line before `end` (unlike the lines that follow).
+            @rewriter << Source::Insert.new(at, "\n#{indent}extend #{mixin_name}\n")
+          end
         end
       end
     end