From db7f4a657749d224da936cefa1b339f3b57209d9 Mon Sep 17 00:00:00 2001
From: Soutaro Matsumoto <matsumoto@soutaro.com>
Date: Thu, 25 Apr 2024 13:27:28 +0900
Subject: [PATCH] Support module declarations with self types

---
 lib/rbs/inline/ast/declarations.rb  | 12 ++++++++++++
 lib/rbs/inline/parser.rb            | 11 +++++++++++
 lib/rbs/inline/writer.rb            |  4 +++-
 sig/rbs/inline/ast/declarations.rbs |  4 ++++
 test/rbs/inline/writer_test.rb      | 18 ++++++++++++++++++
 5 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/lib/rbs/inline/ast/declarations.rb b/lib/rbs/inline/ast/declarations.rb
index 44790ce..eba6779 100644
--- a/lib/rbs/inline/ast/declarations.rb
+++ b/lib/rbs/inline/ast/declarations.rb
@@ -76,16 +76,28 @@ class ModuleDecl < Base
           attr_reader :node
           attr_reader :members
           attr_reader :comments
+          attr_reader :inner_comments
 
           def initialize(node, comments)
             @node = node
             @comments = comments
             @members = []
+            @inner_comments = []
           end
 
           def module_name
             type_name(node.constant_path)
           end
+
+          def module_selfs
+            inner_comments.flat_map do |comment|
+              comment.annotations.filter_map do |ann|
+                if ann.is_a?(AST::Annotations::ModuleSelf)
+                  ann
+                end
+              end
+            end
+          end
         end
       end
     end
diff --git a/lib/rbs/inline/parser.rb b/lib/rbs/inline/parser.rb
index 73aeab8..f2a62c4 100644
--- a/lib/rbs/inline/parser.rb
+++ b/lib/rbs/inline/parser.rb
@@ -92,6 +92,17 @@ def visit_module_node(node)
         push_class_module_decl(module_decl) do
           visit node.body
         end
+
+        annotation_range = node.location.start_line .. node.location.end_line
+        annotations = comments.each_value.select do |annotation|
+          range = annotation.line_range
+          annotation_range.begin < range.begin && range.end < annotation_range.end
+        end
+
+        annotations.each do |annot|
+          comments.delete(annot.line_range.end)
+          module_decl.inner_comments.push(annot)
+        end
       end
 
       def visit_def_node(node)
diff --git a/lib/rbs/inline/writer.rb b/lib/rbs/inline/writer.rb
index 193c724..c509850 100644
--- a/lib/rbs/inline/writer.rb
+++ b/lib/rbs/inline/writer.rb
@@ -103,11 +103,13 @@ def translate_module_decl(decl)
           end
         end
 
+        self_types = decl.module_selfs.map { _1.constraint }.compact
+
         RBS::AST::Declarations::Module.new(
           name: decl.module_name,
           type_params: [],
           members: members,
-          self_types: [],
+          self_types: self_types,
           annotations: [],
           location: nil,
           comment: comment
diff --git a/sig/rbs/inline/ast/declarations.rbs b/sig/rbs/inline/ast/declarations.rbs
index 9eacfc2..db96a82 100644
--- a/sig/rbs/inline/ast/declarations.rbs
+++ b/sig/rbs/inline/ast/declarations.rbs
@@ -38,9 +38,13 @@ module RBS
 
           attr_reader comments: AnnotationParser::ParsingResult?
 
+          attr_reader inner_comments: Array[AnnotationParser::ParsingResult]
+
           def initialize: (Prism::ModuleNode, AnnotationParser::ParsingResult?) -> void
 
           %a{pure} def module_name: () -> TypeName?
+
+          %a{pure} def module_selfs: () -> Array[Annotations::ModuleSelf]
         end
       end
     end
diff --git a/test/rbs/inline/writer_test.rb b/test/rbs/inline/writer_test.rb
index f4aad53..d8bb404 100644
--- a/test/rbs/inline/writer_test.rb
+++ b/test/rbs/inline/writer_test.rb
@@ -6,6 +6,7 @@ class RBS::Inline::WriterTest < Minitest::Test
   include RBS::Inline
 
   def translate(src)
+    src = "# rbs_inline: enabled\n\n" + src
     uses, decls = Parser.parse(Prism.parse(src, filepath: "a.rb"))
     Writer.write(uses, decls)
   end
@@ -348,4 +349,21 @@ def length: ...
       end
     RBS
   end
+
+  def test_module_self
+    output = translate(<<~RUBY)
+      module Foo
+        # @rbs module-self BasicObject
+
+        def foo
+        end
+      end
+    RUBY
+
+    assert_equal <<~RBS, output
+      module Foo : BasicObject
+        def foo: () -> untyped
+      end
+    RBS
+  end
 end