Support module-self inline annotation

config.yml

@@ -784,6 +784,30 @@ nodes:
       - name: type_name_location
         c_type: rbs_location_range
         optional: true
+  - name: RBS::AST::Ruby::Annotations::ModuleSelfAnnotation
+    rust_name: ModuleSelfAnnotationNode
+    fields:
+      - name: prefix_location
+        c_type: rbs_location_range
+      - name: keyword_location
+        c_type: rbs_location_range
+      - name: colon_location
+        c_type: rbs_location_range
+      - name: name
+        c_type: rbs_type_name
+      - name: args
+        c_type: rbs_node_list
+      - name: open_bracket_location
+        c_type: rbs_location_range
+        optional: true
+      - name: close_bracket_location
+        c_type: rbs_location_range
+        optional: true
+      - name: args_comma_locations
+        c_type: rbs_location_range_list
+      - name: comment_location
+        c_type: rbs_location_range
+        optional: true
   - name: RBS::AST::Ruby::Annotations::ParamTypeAnnotation
     rust_name: ParamTypeAnnotationNode
     fields:

docs/inline.md

@@ -123,10 +123,38 @@ end
 
 This creates the types `::API`, `::API::V1`, and `::API::V1::Resources`.
 
+### `module-self` constraint
+
+The `module-self` constraint declares which classes or modules the module can be mixed into.
+
+```ruby
+# @rbs module-self: _Each[String]
+module Enumerable2
+end
+```
+
+This is equivalent to `module Enumerable2 : _Each[String]` in RBS, meaning `Enumerable2` can only be included in classes that satisfy the `_Each[String]` interface.
+
+Multiple `module-self` constraints can be declared with separate annotations:
+
+```ruby
+# @rbs module-self: _Each[String]
+# @rbs module-self: Comparable
+module StringCollection
+end
+```
+
+You can add a description after `--`:
+
+```ruby
+# @rbs module-self: Minitest::Test -- depending on assertion methods
+module TestHelper
+end
+```
+
 ### Current Limitations
 
 - Generic module definitions are not supported
-- Module self-type constraints are not supported
 
 ## Method Definitions
 

ext/rbs_extension/ast_translation.c

@@ -947,6 +947,27 @@ VALUE rbs_struct_to_ruby_value(rbs_translation_context_t ctx, rbs_node_t *instan
             &h
         );
     }
+    case RBS_AST_RUBY_ANNOTATIONS_MODULE_SELF_ANNOTATION: {
+        rbs_ast_ruby_annotations_module_self_annotation_t *node = (rbs_ast_ruby_annotations_module_self_annotation_t *) instance;
+
+        VALUE h = rb_hash_new();
+        rb_hash_aset(h, ID2SYM(rb_intern("location")), rbs_location_range_to_ruby_location(ctx, node->base.location));
+        rb_hash_aset(h, ID2SYM(rb_intern("prefix_location")), rbs_location_range_to_ruby_location(ctx, node->prefix_location));
+        rb_hash_aset(h, ID2SYM(rb_intern("keyword_location")), rbs_location_range_to_ruby_location(ctx, node->keyword_location));
+        rb_hash_aset(h, ID2SYM(rb_intern("colon_location")), rbs_location_range_to_ruby_location(ctx, node->colon_location));
+        rb_hash_aset(h, ID2SYM(rb_intern("name")), rbs_struct_to_ruby_value(ctx, (rbs_node_t *) node->name)); // rbs_type_name
+        rb_hash_aset(h, ID2SYM(rb_intern("args")), rbs_node_list_to_ruby_array(ctx, node->args));
+        rb_hash_aset(h, ID2SYM(rb_intern("open_bracket_location")), rbs_location_range_to_ruby_location(ctx, node->open_bracket_location));   // optional
+        rb_hash_aset(h, ID2SYM(rb_intern("close_bracket_location")), rbs_location_range_to_ruby_location(ctx, node->close_bracket_location)); // optional
+        rb_hash_aset(h, ID2SYM(rb_intern("args_comma_locations")), rbs_location_range_list_to_ruby_array(ctx, node->args_comma_locations));
+        rb_hash_aset(h, ID2SYM(rb_intern("comment_location")), rbs_location_range_to_ruby_location(ctx, node->comment_location)); // optional
+
+        return CLASS_NEW_INSTANCE(
+            RBS_AST_Ruby_Annotations_ModuleSelfAnnotation,
+            1,
+            &h
+        );
+    }
     case RBS_AST_RUBY_ANNOTATIONS_NODE_TYPE_ASSERTION: {
         rbs_ast_ruby_annotations_node_type_assertion_t *node = (rbs_ast_ruby_annotations_node_type_assertion_t *) instance;
 

ext/rbs_extension/class_constants.c

@@ -54,6 +54,7 @@ VALUE RBS_AST_Ruby_Annotations_DoubleSplatParamTypeAnnotation;
 VALUE RBS_AST_Ruby_Annotations_InstanceVariableAnnotation;
 VALUE RBS_AST_Ruby_Annotations_MethodTypesAnnotation;
 VALUE RBS_AST_Ruby_Annotations_ModuleAliasAnnotation;
+VALUE RBS_AST_Ruby_Annotations_ModuleSelfAnnotation;
 VALUE RBS_AST_Ruby_Annotations_NodeTypeAssertion;
 VALUE RBS_AST_Ruby_Annotations_ParamTypeAnnotation;
 VALUE RBS_AST_Ruby_Annotations_ReturnTypeAnnotation;
@@ -147,6 +148,7 @@ void rbs__init_constants(void) {
     IMPORT_CONSTANT(RBS_AST_Ruby_Annotations_InstanceVariableAnnotation, RBS_AST_Ruby_Annotations, "InstanceVariableAnnotation");
     IMPORT_CONSTANT(RBS_AST_Ruby_Annotations_MethodTypesAnnotation, RBS_AST_Ruby_Annotations, "MethodTypesAnnotation");
     IMPORT_CONSTANT(RBS_AST_Ruby_Annotations_ModuleAliasAnnotation, RBS_AST_Ruby_Annotations, "ModuleAliasAnnotation");
+    IMPORT_CONSTANT(RBS_AST_Ruby_Annotations_ModuleSelfAnnotation, RBS_AST_Ruby_Annotations, "ModuleSelfAnnotation");
     IMPORT_CONSTANT(RBS_AST_Ruby_Annotations_NodeTypeAssertion, RBS_AST_Ruby_Annotations, "NodeTypeAssertion");
     IMPORT_CONSTANT(RBS_AST_Ruby_Annotations_ParamTypeAnnotation, RBS_AST_Ruby_Annotations, "ParamTypeAnnotation");
     IMPORT_CONSTANT(RBS_AST_Ruby_Annotations_ReturnTypeAnnotation, RBS_AST_Ruby_Annotations, "ReturnTypeAnnotation");

ext/rbs_extension/class_constants.h

@@ -62,6 +62,7 @@ extern VALUE RBS_AST_Ruby_Annotations_DoubleSplatParamTypeAnnotation;
 extern VALUE RBS_AST_Ruby_Annotations_InstanceVariableAnnotation;
 extern VALUE RBS_AST_Ruby_Annotations_MethodTypesAnnotation;
 extern VALUE RBS_AST_Ruby_Annotations_ModuleAliasAnnotation;
+extern VALUE RBS_AST_Ruby_Annotations_ModuleSelfAnnotation;
 extern VALUE RBS_AST_Ruby_Annotations_NodeTypeAssertion;
 extern VALUE RBS_AST_Ruby_Annotations_ParamTypeAnnotation;
 extern VALUE RBS_AST_Ruby_Annotations_ReturnTypeAnnotation;

include/rbs/ast.h

@@ -97,44 +97,45 @@ enum rbs_node_type {
     RBS_AST_RUBY_ANNOTATIONS_INSTANCE_VARIABLE_ANNOTATION = 36,
     RBS_AST_RUBY_ANNOTATIONS_METHOD_TYPES_ANNOTATION = 37,
     RBS_AST_RUBY_ANNOTATIONS_MODULE_ALIAS_ANNOTATION = 38,
-    RBS_AST_RUBY_ANNOTATIONS_NODE_TYPE_ASSERTION = 39,
-    RBS_AST_RUBY_ANNOTATIONS_PARAM_TYPE_ANNOTATION = 40,
-    RBS_AST_RUBY_ANNOTATIONS_RETURN_TYPE_ANNOTATION = 41,
-    RBS_AST_RUBY_ANNOTATIONS_SKIP_ANNOTATION = 42,
-    RBS_AST_RUBY_ANNOTATIONS_SPLAT_PARAM_TYPE_ANNOTATION = 43,
-    RBS_AST_RUBY_ANNOTATIONS_TYPE_APPLICATION_ANNOTATION = 44,
-    RBS_AST_STRING = 45,
-    RBS_AST_TYPE_PARAM = 46,
-    RBS_METHOD_TYPE = 47,
-    RBS_NAMESPACE = 48,
-    RBS_SIGNATURE = 49,
-    RBS_TYPE_NAME = 50,
-    RBS_TYPES_ALIAS = 51,
-    RBS_TYPES_BASES_ANY = 52,
-    RBS_TYPES_BASES_BOOL = 53,
-    RBS_TYPES_BASES_BOTTOM = 54,
-    RBS_TYPES_BASES_CLASS = 55,
-    RBS_TYPES_BASES_INSTANCE = 56,
-    RBS_TYPES_BASES_NIL = 57,
-    RBS_TYPES_BASES_SELF = 58,
-    RBS_TYPES_BASES_TOP = 59,
-    RBS_TYPES_BASES_VOID = 60,
-    RBS_TYPES_BLOCK = 61,
-    RBS_TYPES_CLASS_INSTANCE = 62,
-    RBS_TYPES_CLASS_SINGLETON = 63,
-    RBS_TYPES_FUNCTION = 64,
-    RBS_TYPES_FUNCTION_PARAM = 65,
-    RBS_TYPES_INTERFACE = 66,
-    RBS_TYPES_INTERSECTION = 67,
-    RBS_TYPES_LITERAL = 68,
-    RBS_TYPES_OPTIONAL = 69,
-    RBS_TYPES_PROC = 70,
-    RBS_TYPES_RECORD = 71,
-    RBS_TYPES_RECORD_FIELD_TYPE = 72,
-    RBS_TYPES_TUPLE = 73,
-    RBS_TYPES_UNION = 74,
-    RBS_TYPES_UNTYPED_FUNCTION = 75,
-    RBS_TYPES_VARIABLE = 76,
+    RBS_AST_RUBY_ANNOTATIONS_MODULE_SELF_ANNOTATION = 39,
+    RBS_AST_RUBY_ANNOTATIONS_NODE_TYPE_ASSERTION = 40,
+    RBS_AST_RUBY_ANNOTATIONS_PARAM_TYPE_ANNOTATION = 41,
+    RBS_AST_RUBY_ANNOTATIONS_RETURN_TYPE_ANNOTATION = 42,
+    RBS_AST_RUBY_ANNOTATIONS_SKIP_ANNOTATION = 43,
+    RBS_AST_RUBY_ANNOTATIONS_SPLAT_PARAM_TYPE_ANNOTATION = 44,
+    RBS_AST_RUBY_ANNOTATIONS_TYPE_APPLICATION_ANNOTATION = 45,
+    RBS_AST_STRING = 46,
+    RBS_AST_TYPE_PARAM = 47,
+    RBS_METHOD_TYPE = 48,
+    RBS_NAMESPACE = 49,
+    RBS_SIGNATURE = 50,
+    RBS_TYPE_NAME = 51,
+    RBS_TYPES_ALIAS = 52,
+    RBS_TYPES_BASES_ANY = 53,
+    RBS_TYPES_BASES_BOOL = 54,
+    RBS_TYPES_BASES_BOTTOM = 55,
+    RBS_TYPES_BASES_CLASS = 56,
+    RBS_TYPES_BASES_INSTANCE = 57,
+    RBS_TYPES_BASES_NIL = 58,
+    RBS_TYPES_BASES_SELF = 59,
+    RBS_TYPES_BASES_TOP = 60,
+    RBS_TYPES_BASES_VOID = 61,
+    RBS_TYPES_BLOCK = 62,
+    RBS_TYPES_CLASS_INSTANCE = 63,
+    RBS_TYPES_CLASS_SINGLETON = 64,
+    RBS_TYPES_FUNCTION = 65,
+    RBS_TYPES_FUNCTION_PARAM = 66,
+    RBS_TYPES_INTERFACE = 67,
+    RBS_TYPES_INTERSECTION = 68,
+    RBS_TYPES_LITERAL = 69,
+    RBS_TYPES_OPTIONAL = 70,
+    RBS_TYPES_PROC = 71,
+    RBS_TYPES_RECORD = 72,
+    RBS_TYPES_RECORD_FIELD_TYPE = 73,
+    RBS_TYPES_TUPLE = 74,
+    RBS_TYPES_UNION = 75,
+    RBS_TYPES_UNTYPED_FUNCTION = 76,
+    RBS_TYPES_VARIABLE = 77,
     RBS_AST_SYMBOL,
 };
 
@@ -633,6 +634,20 @@ typedef struct rbs_ast_ruby_annotations_module_alias_annotation {
     rbs_location_range type_name_location; /* Optional */
 } rbs_ast_ruby_annotations_module_alias_annotation_t;
 
+typedef struct rbs_ast_ruby_annotations_module_self_annotation {
+    rbs_node_t base;
+
+    rbs_location_range prefix_location;
+    rbs_location_range keyword_location;
+    rbs_location_range colon_location;
+    struct rbs_type_name *name;
+    struct rbs_node_list *args;
+    rbs_location_range open_bracket_location;  /* Optional */
+    rbs_location_range close_bracket_location; /* Optional */
+    rbs_location_range_list_t *args_comma_locations;
+    rbs_location_range comment_location; /* Optional */
+} rbs_ast_ruby_annotations_module_self_annotation_t;
+
 typedef struct rbs_ast_ruby_annotations_node_type_assertion {
     rbs_node_t base;
 
@@ -980,6 +995,7 @@ rbs_ast_ruby_annotations_double_splat_param_type_annotation_t *rbs_ast_ruby_anno
 rbs_ast_ruby_annotations_instance_variable_annotation_t *rbs_ast_ruby_annotations_instance_variable_annotation_new(rbs_allocator_t *allocator, rbs_location_range location, rbs_location_range prefix_location, rbs_ast_symbol_t *ivar_name, rbs_location_range ivar_name_location, rbs_location_range colon_location, rbs_node_t *type, rbs_location_range comment_location);
 rbs_ast_ruby_annotations_method_types_annotation_t *rbs_ast_ruby_annotations_method_types_annotation_new(rbs_allocator_t *allocator, rbs_location_range location, rbs_location_range prefix_location, rbs_node_list_t *overloads, rbs_location_range_list_t *vertical_bar_locations, rbs_location_range dot3_location);
 rbs_ast_ruby_annotations_module_alias_annotation_t *rbs_ast_ruby_annotations_module_alias_annotation_new(rbs_allocator_t *allocator, rbs_location_range location, rbs_location_range prefix_location, rbs_location_range keyword_location, rbs_type_name_t *type_name, rbs_location_range type_name_location);
+rbs_ast_ruby_annotations_module_self_annotation_t *rbs_ast_ruby_annotations_module_self_annotation_new(rbs_allocator_t *allocator, rbs_location_range location, rbs_location_range prefix_location, rbs_location_range keyword_location, rbs_location_range colon_location, rbs_type_name_t *name, rbs_node_list_t *args, rbs_location_range open_bracket_location, rbs_location_range close_bracket_location, rbs_location_range_list_t *args_comma_locations, rbs_location_range comment_location);
 rbs_ast_ruby_annotations_node_type_assertion_t *rbs_ast_ruby_annotations_node_type_assertion_new(rbs_allocator_t *allocator, rbs_location_range location, rbs_location_range prefix_location, rbs_node_t *type);
 rbs_ast_ruby_annotations_param_type_annotation_t *rbs_ast_ruby_annotations_param_type_annotation_new(rbs_allocator_t *allocator, rbs_location_range location, rbs_location_range prefix_location, rbs_location_range name_location, rbs_location_range colon_location, rbs_node_t *param_type, rbs_location_range comment_location);
 rbs_ast_ruby_annotations_return_type_annotation_t *rbs_ast_ruby_annotations_return_type_annotation_new(rbs_allocator_t *allocator, rbs_location_range location, rbs_location_range prefix_location, rbs_location_range return_location, rbs_location_range colon_location, rbs_node_t *return_type, rbs_location_range comment_location);

include/rbs/lexer.h

@@ -70,6 +70,7 @@ enum RBSTokenType {
     kATRBS,        /* @rbs */
     kSKIP,         /* skip */
     kRETURN,       /* return */
+    kMODULESELF,   /* module-self */
 
     tLIDENT,    /* Identifiers starting with lower case */
     tUIDENT,    /* Identifiers starting with upper case */

lib/rbs/ast/ruby/annotations.rb

@@ -353,6 +353,48 @@ def type_fingerprint
           end
         end
 
+        class ModuleSelfAnnotation < Base
+          attr_reader :keyword_location, :colon_location, :name, :args, :open_bracket_location, :close_bracket_location, :args_comma_locations, :comment_location
+
+          def initialize(location:, prefix_location:, keyword_location:, colon_location:, name:, args:, open_bracket_location:, close_bracket_location:, args_comma_locations:, comment_location:)
+            super(location, prefix_location)
+            @keyword_location = keyword_location
+            @colon_location = colon_location
+            @name = name
+            @args = args
+            @open_bracket_location = open_bracket_location
+            @close_bracket_location = close_bracket_location
+            @args_comma_locations = args_comma_locations
+            @comment_location = comment_location
+          end
+
+          def map_type_name
+            mapped_args = args.map { |type| type.map_type_name { yield _1 } }
+
+            self.class.new(
+              location:,
+              prefix_location:,
+              keyword_location:,
+              colon_location:,
+              name: yield(name),
+              args: mapped_args,
+              open_bracket_location:,
+              close_bracket_location:,
+              args_comma_locations:,
+              comment_location:
+            ) #: self
+          end
+
+          def type_fingerprint
+            [
+              "annots/module_self",
+              name.to_s,
+              args.map(&:to_s),
+              comment_location&.source
+            ]
+          end
+        end
+
         class BlockParamTypeAnnotation < Base
           attr_reader :ampersand_location, :name_location, :colon_location, :question_location, :type_location, :type, :comment_location
 

lib/rbs/ast/ruby/declarations.rb

@@ -137,7 +137,17 @@ def each_decl(&block)
 
           def type_params = []
 
-          def self_types = []
+          def self_types
+            members.filter_map do |member|
+              if member.is_a?(Members::ModuleSelfMember)
+                AST::Declarations::Module::Self.new(
+                  name: member.name,
+                  args: member.args,
+                  location: member.location
+                )
+              end
+            end
+          end
 
           def location
             rbs_location(node.location)

lib/rbs/ast/ruby/members.rb

@@ -717,6 +717,34 @@ def type_fingerprint
             ]
           end
         end
+
+        class ModuleSelfMember < Base
+          attr_reader :annotation
+
+          def initialize(buffer, annotation)
+            super(buffer)
+            @annotation = annotation
+          end
+
+          def name
+            annotation.name
+          end
+
+          def args
+            annotation.args
+          end
+
+          def location
+            annotation.location
+          end
+
+          def type_fingerprint
+            [
+              "members/module_self",
+              annotation.type_fingerprint
+            ]
+          end
+        end
       end
     end
   end

lib/rbs/environment.rb

@@ -854,6 +854,12 @@ def resolve_ruby_member(resolver, member, context:)
           member.buffer,
           resolved_annotation
         )
+      when AST::Ruby::Members::ModuleSelfMember
+        resolved_annotation = member.annotation.map_type_name {|name| absolute_type_name(resolver, nil, name, context: context) }
+        AST::Ruby::Members::ModuleSelfMember.new(
+          member.buffer,
+          resolved_annotation
+        )
       else
         raise "Unknown member type: #{member.class}"
       end

lib/rbs/inline_parser.rb

@@ -134,6 +134,32 @@ def visit_module_node(node)
         end
 
         module_decl = AST::Ruby::Declarations::ModuleDecl.new(buffer, module_name, node)
+
+        if leading_ref = comments.leading_block(node)
+          unused_annotations = [] #: Array[AST::Ruby::CommentBlock::AnnotationSyntaxError | AST::Ruby::Annotations::leading_annotation]
+
+          leading_ref.block.each_paragraph([]) do |paragraph|
+            case paragraph
+            when AST::Ruby::Annotations::ModuleSelfAnnotation
+              module_decl.members << AST::Ruby::Members::ModuleSelfMember.new(buffer, paragraph)
+            when Location
+              # Skip
+            when AST::Ruby::CommentBlock::AnnotationSyntaxError
+              unused_annotations << paragraph
+            when AST::Ruby::Annotations::SkipAnnotation
+              # Already handled by skip_node?
+            else
+              unused_annotations << paragraph
+            end
+          end
+
+          unless unused_annotations.empty?
+            report_unused_annotation(*unused_annotations)
+          end
+
+          leading_ref.associate!
+        end
+
         insert_declaration(module_decl)
         visit_class_or_module_body(module_decl, node)
       end