Add RBS::Rewriter and use it in rbs annotate

lib/rbs.rb

@@ -51,6 +51,7 @@
 require "rbs/resolver/type_name_resolver"
 require "rbs/ast/comment"
 require "rbs/writer"
+require "rbs/rewriter"
 require "rbs/prototype/helpers"
 require "rbs/prototype/rbi"
 require "rbs/prototype/rb"

lib/rbs/annotate/rdoc_annotator.rb

@@ -13,25 +13,22 @@ def initialize(source:)
         @include_filename = true
       end
 
-      def annotate_file(path, preserve:)
-        content = path.read()
+      def annotate_file(path, preserve: true)
+        buffer, _, decls = Parser.parse_signature(path.read())
 
-        _, _, decls = Parser.parse_signature(content)
+        rewriter = Rewriter.new(buffer)
+        annotate_decls(decls, rewriter)
 
-        annotate_decls(decls)
-
-        path.open("w") do |io|
-          Writer.new(out: io).preserve!(preserve: preserve).write(decls)
-        end
+        path.write(rewriter.string)
       end
 
-      def annotate_decls(decls, outer: [])
+      def annotate_decls(decls, rewriter, outer: [])
         decls.each do |decl|
           case decl
           when AST::Declarations::Class, AST::Declarations::Module
-            annotate_class(decl, outer: outer)
+            annotate_class(decl, rewriter, outer: outer)
           when AST::Declarations::Constant
-            annotate_constant(decl, outer: outer)
+            annotate_constant(decl, rewriter, outer: outer)
           end
         end
       end
@@ -244,44 +241,44 @@ def doc_for_attribute(typename, attr_name, require: nil, singleton:, tester:)
         end
       end
 
-      def annotate_class(decl, outer:)
+      def annotate_class(decl, rewriter, outer:)
         annots = annotations(decl)
 
         full_name = resolve_name(decl.name, outer: outer)
         unless annots.skip?
           text = resolve_doc_source(annots.copy_annotation, tester: annots) { doc_for_class(full_name, tester: annots) }
         end
 
-        replace_comment(decl, text)
+        replace_comment(decl, text, rewriter)
 
         unless annots.skip_all?
           outer_ = outer + [decl.name.to_namespace]
 
           decl.each_member do |member|
             case member
             when AST::Members::MethodDefinition
-              annotate_method(full_name, member)
+              annotate_method(full_name, member, rewriter)
             when AST::Members::Alias
-              annotate_alias(full_name, member)
+              annotate_alias(full_name, member, rewriter)
             when AST::Members::AttrReader, AST::Members::AttrAccessor, AST::Members::AttrWriter
-              annotate_attribute(full_name, member)
+              annotate_attribute(full_name, member, rewriter)
             end
           end
 
-          annotate_decls(decl.each_decl.to_a, outer: outer_)
+          annotate_decls(decl.each_decl.to_a, rewriter, outer: outer_)
         end
       end
 
-      def annotate_constant(const, outer:)
+      def annotate_constant(const, rewriter, outer:)
         annots = Annotations.new([])
 
         full_name = resolve_name(const.name, outer: outer)
         text = doc_for_constant(full_name, tester: annots)
 
-        replace_comment(const, text)
+        replace_comment(const, text, rewriter)
       end
 
-      def annotate_alias(typename, als)
+      def annotate_alias(typename, als, rewriter)
         annots = annotations(als)
 
         unless annots.skip?
@@ -295,7 +292,7 @@ def annotate_alias(typename, als)
           end
         end
 
-        replace_comment(als, text)
+        replace_comment(als, text, rewriter)
       end
 
       def join_docs(docs, separator: "----")
@@ -311,7 +308,7 @@ def join_docs(docs, separator: "----")
         end
       end
 
-      def annotate_method(typename, method)
+      def annotate_method(typename, method, rewriter)
         annots = annotations(method)
 
         unless annots.skip?
@@ -337,10 +334,10 @@ def annotate_method(typename, method)
           }
         end
 
-        replace_comment(method, text)
+        replace_comment(method, text, rewriter)
       end
 
-      def annotate_attribute(typename, attr)
+      def annotate_attribute(typename, attr, rewriter)
         annots = annotations(attr)
 
         unless annots.skip?
@@ -368,18 +365,17 @@ def annotate_attribute(typename, attr)
           end
         end
 
-        replace_comment(attr, text)
+        replace_comment(attr, text, rewriter)
       end
 
-      def replace_comment(commented, string)
+      def replace_comment(commented, string, rewriter)
         if string
           if string.empty?
-            commented.instance_variable_set(:@comment, nil)
+            rewriter.delete_comment(commented.comment) if commented.comment
+          elsif commented.comment
+            rewriter.replace_comment(commented.comment, content: string)
           else
-            commented.instance_variable_set(
-              :@comment,
-              AST::Comment.new(location: nil, string: string)
-            )
+            rewriter.add_comment(commented.location || raise, *commented.annotations.filter_map(&:location), content: string)
           end
         end
       end

lib/rbs/cli.rb

@@ -962,8 +962,6 @@ def run_annotate(args, options)
       source = RBS::Annotate::RDocSource.new()
       annotator = RBS::Annotate::RDocAnnotator.new(source: source)
 
-      preserve = true
-
       OptionParser.new do |opts|
         opts.banner = <<-EOB
 Usage: rbs annotate [options...] [files...]
@@ -984,7 +982,7 @@ def run_annotate(args, options)
         opts.on("-d", "--dir DIRNAME", "Load RDoc from DIRNAME") {|d| source.extra_dirs << Pathname(d) }
         opts.on("--[no-]arglists", "Generate arglists section (defaults to true)") {|b| annotator.include_arg_lists = b }
         opts.on("--[no-]filename", "Include source file name in the documentation (defaults to true)") {|b| annotator.include_filename = b }
-        opts.on("--[no-]preserve", "Try preserve the format of the original file (defaults to true)") {|b| preserve = b }
+        opts.on("--[no-]preserve", "[Deprecated] It always preserves the format") { stdout.puts "The `--preserve` option is deprecated. The tool always preserves the format of RBS files." }
       end.parse!(args)
 
       source.load()
@@ -994,11 +992,11 @@ def run_annotate(args, options)
         if path.directory?
           Pathname.glob((path + "**/*.rbs").to_s).each do |path|
             stdout.puts "Processing #{path}..."
-            annotator.annotate_file(path, preserve: preserve)
+            annotator.annotate_file(path)
           end
         else
           stdout.puts "Processing #{path}..."
-          annotator.annotate_file(path, preserve: preserve)
+          annotator.annotate_file(path)
         end
       end
 

lib/rbs/rewriter.rb

@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+
+module RBS
+  class Rewriter
+    attr_reader :buffer
+
+    def initialize(buffer)
+      raise "Rewriter only supports toplevel buffers" if buffer.parent
+
+      @buffer = buffer
+      @rewrites = []
+    end
+
+    def rewrite(location, string)
+      @rewrites.each do |existing_location, _|
+        if location.start_pos < existing_location.end_pos && existing_location.start_pos < location.end_pos
+          raise "Overlapping rewrites: #{existing_location} and #{location}"
+        end
+      end
+
+      @rewrites << [location, string]
+      self
+    end
+
+    def add_comment(*locations, content:)
+      earliest = locations.min_by(&:start_pos) or raise "At least one location is required"
+      insert_pos = earliest.start_pos
+      indent = " " * earliest.start_column
+
+      formatted = format_comment(content, indent)
+
+      loc = Location.new(buffer, insert_pos, insert_pos)
+      rewrite(loc, "#{formatted}\n#{indent}")
+    end
+
+    def replace_comment(comment, content:)
+      location = comment.location or raise "Comment must have a location"
+      indent = " " * location.start_column
+
+      rewrite(location, format_comment(content, indent))
+    end
+
+    def delete_comment(comment)
+      location = comment.location or raise "Comment must have a location"
+      line_start = location.start_pos - location.start_column
+      line_end = location.end_pos + 1
+      loc = Location.new(buffer, line_start, line_end)
+      rewrite(loc, "")
+    end
+
+    def string
+      result = buffer.content.dup
+
+      @rewrites.sort_by { |location, _| location.start_pos }.reverse_each do |location, replacement|
+        result[location.start_pos...location.end_pos] = replacement
+      end
+
+      result
+    end
+
+    private
+
+    def format_comment(content, indent)
+      content.lines.map do |line|
+        line = line.chomp
+        line.empty? ? "#" : "# #{line}"
+      end.join("\n#{indent}")
+    end
+  end
+end

sig/annotate/rdoc_annotater.rbs

@@ -8,9 +8,10 @@ module RBS
 
       def initialize: (source: RDocSource) -> void
 
-      def annotate_file: (Pathname, preserve: bool) -> void
+      def annotate_file: (Pathname) -> void
+                       | %a{deprecated: it preserves the format by default} (Pathname, preserve: bool) -> void
 
-      def annotate_decls: (Array[AST::Declarations::t], ?outer: Array[Namespace]) -> void
+      def annotate_decls: (Array[AST::Declarations::t], Rewriter, ?outer: Array[Namespace]) -> void
 
       interface _PathTester
         def test_path: (String) -> bool
@@ -50,15 +51,15 @@ module RBS
 
       def resolve_doc_source: (Annotations::Copy?, tester: _PathTester) { () -> String? } -> String?
 
-      def annotate_class: (AST::Declarations::Class | AST::Declarations::Module, outer: Array[Namespace]) -> void
+      def annotate_class: (AST::Declarations::Class | AST::Declarations::Module, Rewriter, outer: Array[Namespace]) -> void
 
-      def annotate_constant: (AST::Declarations::Constant, outer: Array[Namespace]) -> void
+      def annotate_constant: (AST::Declarations::Constant, Rewriter, outer: Array[Namespace]) -> void
 
-      def annotate_method: (TypeName, AST::Members::MethodDefinition) -> void
+      def annotate_method: (TypeName, AST::Members::MethodDefinition, Rewriter) -> void
 
-      def annotate_alias: (TypeName, AST::Members::Alias) -> void
+      def annotate_alias: (TypeName, AST::Members::Alias, Rewriter) -> void
 
-      def annotate_attribute: (TypeName, AST::Members::AttrReader | AST::Members::AttrWriter | AST::Members::AttrAccessor) -> void
+      def annotate_attribute: (TypeName, AST::Members::AttrReader | AST::Members::AttrWriter | AST::Members::AttrAccessor, Rewriter) -> void
 
       def annotations: (_Annotated) -> Annotations
 
@@ -70,10 +71,12 @@ module RBS
       # - If empty string is given as `comment`, it deletes the original comment.
       # - If `nil` is given as `comment`, it keeps the original comment.
       #
-      def replace_comment: (Object & _Commented, String? comment) -> void
+      def replace_comment: (_Commented, String?, Rewriter) -> void
 
       interface _Commented
-        def comment: () -> AST::Comment?
+        %a{pure} def comment: () -> AST::Comment?
+        %a{pure} def location: () -> Location[untyped, untyped]?
+        def annotations: () -> Array[AST::Annotation]
       end
 
       def resolve_name: (TypeName, outer: Array[Namespace]) -> TypeName

sig/rewriter.rbs

@@ -0,0 +1,45 @@
+module RBS
+  # Rewriter performs targeted character-range replacements on Buffer content.
+  #
+  # Unlike Writer which regenerates entire source from AST, Rewriter preserves
+  # everything outside the rewritten ranges, including non-documentation comments.
+  #
+  # Rewrite requests are buffered and applied all at once when `#string` is called.
+  #
+  class Rewriter
+    attr_reader buffer: Buffer
+
+    # Initialize with a toplevel buffer.
+    # Raises if the buffer has a parent (non-toplevel).
+    #
+    def initialize: (Buffer buffer) -> void
+
+    # Register a rewrite request for the given location.
+    #
+    def rewrite: (Location[untyped, untyped] location, String string) -> self
+
+    # Add a new comment before the earliest of the given locations.
+    #
+    def add_comment: (*Location[untyped, untyped] locations, content: String) -> self
+
+    # Replace an existing comment's content.
+    #
+    def replace_comment: (AST::Comment comment, content: String) -> self
+
+    # Delete an existing comment.
+    #
+    def delete_comment: (AST::Comment comment) -> self
+
+    # Apply all buffered rewrites and return the resulting string.
+    #
+    # Raises if any rewrite ranges overlap.
+    #
+    def string: () -> String
+
+    private
+
+    def format_comment: (String content, String indent) -> String
+
+    @rewrites: Array[[Location[untyped, untyped], String]]
+  end
+end