Parse module-self annotation

ruby · soutaro · Jan 16, 2025 · Jan 16, 2025 · Jan 16, 2025 · Jan 16, 2025
commit 67a38faada84c00b768b96cd6d5b958d6f5cff6b
diff --git a/lib/rbs/ast/ruby/declarations.rb b/lib/rbs/ast/ruby/declarations.rb
@@ -150,20 +150,60 @@ def initialize(node)
         end
 
         class ModuleDecl < Base
+          class SelfConstraint
+            attr_reader :name
+            attr_reader :args
+            attr_reader :annotation
+
+            def initialize(name, args, annotation)
+              @name = name
+              @args = args
+              @annotation = annotation
+            end
+
+            def map_type_name
+              SelfConstraint.new(
+                yield(name),
+                args.map { |type| type.map_type_name { yield(_1) } },
+                annotation
+              ) #: self
+            end
+
+            def self.build(annotations)
+              self_annotations = [] #: Array[Annotation::ModuleSelfAnnotation]
+              other_annotations = [] #: Array[Annotation::leading_annotation]
+
+              annotations.each do |annotation|
+                if annotation.is_a?(Annotation::ModuleSelfAnnotation)
+                  self_annotations << annotation
+                else
+                  other_annotations << annotation
+                end
+              end
+
+              [
+                self_annotations.map { new(_1.type_name, _1.type_args, _1) },
+                other_annotations
+              ]
+            end
+          end
+
           attr_reader :node
           attr_reader :location
           attr_reader :module_name
           attr_reader :module_name_location
           attr_reader :generics
           attr_reader :members
+          attr_reader :self_constraints
 
-          def initialize(buffer, node, location:, module_name:, module_name_location:, generics:)
+          def initialize(buffer, node, location:, module_name:, module_name_location:, generics:, self_constraints:)
             super(buffer)
             @node = node
             @location = location
             @module_name = module_name
             @module_name_location = module_name_location
             @generics = generics
+            @self_constraints = self_constraints
             @members = []
           end
 
@@ -174,7 +214,7 @@ def type_params
           end
 
           def self_types
-            []
+            self_constraints
           end
 
           def each_member(&block)

diff --git a/lib/rbs/environment.rb b/lib/rbs/environment.rb
@@ -578,6 +578,7 @@ def resolve_ruby_declaration(resolver, map, decl, context:, prefix:)
           module_name: module_name,
           module_name_location: decl.module_name_location,
           generics: decl.generics.map_type_name {|name| absolute_type_name(resolver, map, name, context: inner_context) },
+          self_constraints: decl.self_constraints.map { _1.map_type_name {|name| absolute_type_name(resolver, map, name, context: inner_context) } },
         ).tap do |resolved|
           prefix = module_name.to_namespace
           decl.members.each do |member|

diff --git a/lib/rbs/inline_parser.rb b/lib/rbs/inline_parser.rb
@@ -278,19 +278,29 @@ def visit_module_node(node)
         end
 
         generics, leading_comments = AST::Ruby::Declarations::GenericsTypeParams.build(leading_comments)
+        push_type_params(generics.type_params) do
+          self_constraints, leading_comments = AST::Ruby::Declarations::ModuleDecl::SelfConstraint.build(leading_comments)
 
-        decl = AST::Ruby::Declarations::ModuleDecl.new(
-          buffer,
-          node,
-          location: buffer.rbs_location(node.location),
-          module_name: module_name,
-          module_name_location: module_name_location,
-          generics: generics
-        )
+          if leading_block
+            leading_comments.each do |comment|
+              report_unused_annotation(leading_block, comment)
+            end
+          end
 
-        insert_class_module_decl(decl)
-        push_decl_context(decl) do
-          visit node.body
+          decl = AST::Ruby::Declarations::ModuleDecl.new(
+            buffer,
+            node,
+            location: buffer.rbs_location(node.location),
+            module_name: module_name,
+            module_name_location: module_name_location,
+            generics: generics,
+            self_constraints: self_constraints
+          )
+
+          insert_class_module_decl(decl)
+          push_decl_context(decl) do
+            visit node.body
+          end
         end
       end
 
@@ -312,13 +322,7 @@ def visit_def_node(node)
           when AST::Ruby::CommentBlock::AnnotationSyntaxError
             report_annotation_syntax_error(block, annot)
           else
-            if location = block.translate_comment_location(annot.location).first
-              comment, start_pos, end_pos = location
-              start_pos = start_pos + comment.location.start_character_offset
-              end_pos = end_pos + comment.location.start_character_offset
-              loc = Location.new(buffer, start_pos, end_pos)
-              diagnostics << Diagnostics::UnusedAnnotation.new(loc, "Unused annotation")
-            end
+            report_unused_annotation(block, annot)
           end
         end
 
@@ -577,6 +581,16 @@ def report_annotation_syntax_error(block, error)
           diagnostics << Diagnostics::AnnotationSyntaxError.new(loc, "Annotation syntax error: #{error.error.message}")
         end
       end
+
+      def report_unused_annotation(block, annot)
+        if location = block.translate_comment_location(annot.location).first
+          comment, start_pos, end_pos = location
+          start_pos = start_pos + comment.location.start_character_offset
+          end_pos = end_pos + comment.location.start_character_offset
+          loc = Location.new(buffer, start_pos, end_pos)
+          diagnostics << Diagnostics::UnusedAnnotation.new(loc, "Unused annotation")
+        end
+      end
     end
   end
 end
diff --git a/sig/ast/ruby/declarations.rbs b/sig/ast/ruby/declarations.rbs
@@ -116,6 +116,22 @@ module RBS
         end
 
         class ModuleDecl < Base
+          class SelfConstraint
+            attr_reader name: TypeName
+
+            attr_reader args: Array[Types::t]
+
+            attr_reader annotation: Annotation::ModuleSelfAnnotation
+
+            def initialize: (TypeName, Array[Types::t], Annotation::ModuleSelfAnnotation) -> void
+
+            # Note that type names in `annotation` doesn't update
+            #
+            def map_type_name: () { (TypeName) -> TypeName } -> self
+
+            def self.build: (Array[Annotation::leading_annotation]) -> [Array[instance], Array[Annotation::leading_annotation]]
+          end
+
           type member = Members::t | t | SingletonClassDecl
 
           attr_reader node: Prism::ModuleNode
@@ -128,6 +144,8 @@ module RBS
 
           attr_reader generics: GenericsTypeParams
 
+          attr_reader self_constraints: Array[SelfConstraint]
+
           attr_reader members: Array[member]
 
           alias name module_name
@@ -139,11 +157,12 @@ module RBS
             module_name: TypeName,
             module_name_location: Location,
             generics: GenericsTypeParams,
+            self_constraints: Array[SelfConstraint]
           ) -> void
 
           def type_params: () -> Array[AST::TypeParam]
 
-          def self_types: () -> Array[untyped]
+          def self_types: () -> Array[SelfConstraint]
 
           def each_member: () { (Members::t) -> void } -> void
                          | () -> Enumerator[Members::t]

diff --git a/sig/environment/module_entry.rbs b/sig/environment/module_entry.rbs
@@ -43,7 +43,7 @@ module RBS
       #
       def validate_type_params: () -> void
 
-      def self_types: () -> Array[AST::Declarations::Module::Self]
+      def self_types: () -> Array[AST::Declarations::Module::Self | AST::Ruby::Declarations::ModuleDecl::SelfConstraint]
     end
   end
 end
diff --git a/sig/errors.rbs b/sig/errors.rbs
@@ -102,14 +102,17 @@ module RBS
   class NoSelfTypeFoundError < DefinitionError
     include DetailedMessageable
 
+    type self_constraint = AST::Declarations::Module::Self
+                         | AST::Ruby::Declarations::ModuleDecl::SelfConstraint
+
     attr_reader type_name: TypeName
     attr_reader location: Location[untyped, untyped]?
 
     def initialize: (type_name: TypeName, location: Location[untyped, untyped]?) -> void
 
     # The type name in `self` is automatically normalized
     #
-    def self.check!: (AST::Declarations::Module::Self, env: Environment) -> void
+    def self.check!: (self_constraint, env: Environment) -> void
   end
 
   class NoMixinFoundError < DefinitionError

diff --git a/sig/inline_parser.rbs b/sig/inline_parser.rbs
@@ -143,6 +143,8 @@ module RBS
       def constant_node?: (Prism::Node) -> (Prism::ConstantReadNode | Prism::ConstantPathNode | nil)
 
       def report_annotation_syntax_error: (AST::Ruby::CommentBlock, AST::Ruby::CommentBlock::AnnotationSyntaxError) -> void
+
+      def report_unused_annotation: (AST::Ruby::CommentBlock, AST::Ruby::Annotation::t) -> void
     end
   end
 end
diff --git a/test/rbs/ancestor_builder_test.rb b/test/rbs/ancestor_builder_test.rb
@@ -140,12 +140,12 @@ module Foo
 
         builder.one_instance_ancestors(type_name("::Foo")).tap do |a|
           assert_equal type_name("::Foo"), a.type_name
-          assert_equal [:A, :B], a.params
-        end
-
-        builder.one_singleton_ancestors(type_name("::Foo")).tap do |a|
-          assert_equal type_name("::Foo"), a.type_name
-          assert_nil a.params
+          assert_equal(
+            [
+              Ancestor::Instance.new(name: type_name("::BasicObject"), args: [], source: nil)
+            ],
+            a.self_types
+          )
         end
       end
     end

diff --git a/test/rbs/inline_parser_test.rb b/test/rbs/inline_parser_test.rb
@@ -475,4 +475,24 @@ module Foo
       assert_equal "out B < Integer = untyped", mod.generics.type_params[1].to_s
     end
   end
+
+  def test_module_decl__self_constraints
+    buffer, result = parse_ruby(<<~RUBY)
+      # @rbs module-self BasicObject -- type parameter of A
+      module Foo
+      end
+    RUBY
+
+    ret = RBS::InlineParser.parse(buffer, result)
+
+    ret.declarations[0].tap do |mod|
+      assert_instance_of RBS::AST::Ruby::Declarations::ModuleDecl, mod
+
+      mod.self_constraints[0].tap do |constraint|
+        assert_instance_of RBS::AST::Ruby::Declarations::ModuleDecl::SelfConstraint, constraint
+        assert_equal "BasicObject", constraint.type_name.to_s
+        assert_equal [], constraint.type_args
+      end
+    end
+  end
 end