Refactor rbs annotate to use Rewriter instead of Writer

Replace AST mutation + Writer serialization with Rewriter-based
comment editing, preserving non-documentation comments and original
formatting in annotated RBS files.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: soutaro

lib/rbs/annotate/rdoc_annotator.rb

@@ -14,24 +14,21 @@ def initialize(source:)
       end
 
       def annotate_file(path, preserve:)
-        content = path.read()
+        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

sig/annotate/rdoc_annotater.rbs

@@ -10,7 +10,7 @@ module RBS
 
       def annotate_file: (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 +50,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 +70,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