@RBS!: annotation parsing

ruby · soutaro · Jan 16, 2025 · Jan 16, 2025 · Jan 16, 2025 · Jan 16, 2025
commit e5ee033207d04713dac4edcc233aa8a8d2999bf8
diff --git a/config.yml b/config.yml
@@ -463,3 +463,8 @@ nodes:
       - name: colon_location
       - name: type
       - name: comment
+  - name: RBS::AST::Ruby::Annotation::EmbeddedRBSAnnotation
+    fields:
+      - name: location
+      - name: prefix_location
+      - name: members
diff --git a/ext/rbs_extension/parser.c b/ext/rbs_extension/parser.c
@@ -2456,10 +2456,10 @@ static VALUE parse_nested_decl(parserstate *state, const char *nested_in, positi
                   | `public`
                   | `private`
 */
-static VALUE parse_module_members(parserstate *state) {
+static VALUE parse_module_members(parserstate *state, enum TokenType end_token) {
   VALUE members = EMPTY_ARRAY;
 
-  while (state->next_token.type != kEND) {
+  while (state->next_token.type != end_token) {
     VALUE annotations = EMPTY_ARRAY;
     position annot_pos = NullPosition;
     parse_annotations(state, &annotations, &annot_pos);
@@ -2562,7 +2562,7 @@ static VALUE parse_module_decl0(parserstate *state, range keyword_range, VALUE m
     self_types_range = NULL_RANGE;
   }
 
-  VALUE members = parse_module_members(state);
+  VALUE members = parse_module_members(state, kEND);
 
   parser_advance_assert(state, kEND);
   range end_range = state->current_token.range;
@@ -2682,7 +2682,7 @@ static VALUE parse_class_decl0(parserstate *state, range keyword_range, VALUE na
   range lt_range;
   VALUE super = parse_class_decl_super(state, &lt_range);
 
-  VALUE members = parse_module_members(state);
+  VALUE members = parse_module_members(state, kEND);
 
   parser_advance_assert(state, kEND);
 
@@ -3519,6 +3519,16 @@ VALUE parse_inline_annotation(parserstate *state) {
     }
     case kATRBSB: {
       // @rbs!
+      range prefix_range = state->next_token.range;
+      parser_advance_assert(state, kATRBSB);
+
+      VALUE members = parse_module_members(state, pEOF);
+
+      return rbs_ast_ruby_annotation_embedded_rbs_annotation(
+        rbs_new_location(state->buffer, (range) { .start = prefix_range.start, .end = state->current_token.range.end }),
+        rbs_new_location(state->buffer, prefix_range),
+        members
+      );
     }
     case pCOLON: {
       // :

diff --git a/include/rbs/constants.h b/include/rbs/constants.h
@@ -55,6 +55,7 @@ extern VALUE RBS_AST_Ruby_Annotation_ClassIvarTypeAnnotation;
 extern VALUE RBS_AST_Ruby_Annotation_ClassVarTypeAnnotation;
 extern VALUE RBS_AST_Ruby_Annotation_ColonMethodTypeAnnotation;
 extern VALUE RBS_AST_Ruby_Annotation_DoubleSplatParamTypeAnnotation;
+extern VALUE RBS_AST_Ruby_Annotation_EmbeddedRBSAnnotation;
 extern VALUE RBS_AST_Ruby_Annotation_GenericAnnotation;
 extern VALUE RBS_AST_Ruby_Annotation_InheritsAnnotation;
 extern VALUE RBS_AST_Ruby_Annotation_IvarTypeAnnotation;

diff --git a/include/rbs/ruby_objs.h b/include/rbs/ruby_objs.h
@@ -44,6 +44,7 @@ VALUE rbs_ast_ruby_annotation_class_ivar_type_annotation(VALUE location, VALUE p
 VALUE rbs_ast_ruby_annotation_class_var_type_annotation(VALUE location, VALUE prefix_location, VALUE var_name_location, VALUE colon_location, VALUE type, VALUE comment);
 VALUE rbs_ast_ruby_annotation_colon_method_type_annotation(VALUE location, VALUE prefix_location, VALUE annotations, VALUE method_type);
 VALUE rbs_ast_ruby_annotation_double_splat_param_type_annotation(VALUE location, VALUE prefix_location, VALUE operator_location, VALUE param_name_location, VALUE colon_location, VALUE type, VALUE comment);
+VALUE rbs_ast_ruby_annotation_embedded_rbs_annotation(VALUE location, VALUE prefix_location, VALUE members);
 VALUE rbs_ast_ruby_annotation_generic_annotation(VALUE location, VALUE prefix_location, VALUE generic_location, VALUE unchecked_location, VALUE variance_location, VALUE name_location, VALUE upper_bound_operator_location, VALUE upper_bound, VALUE default_type_operator_location, VALUE default_type, VALUE comment);
 VALUE rbs_ast_ruby_annotation_inherits_annotation(VALUE location, VALUE prefix_location, VALUE inherits_location, VALUE type_name, VALUE type_name_location, VALUE open_paren_location, VALUE type_args, VALUE close_paren_location, VALUE comment);
 VALUE rbs_ast_ruby_annotation_ivar_type_annotation(VALUE location, VALUE prefix_location, VALUE var_name_location, VALUE colon_location, VALUE type, VALUE comment);

diff --git a/lib/rbs/ast/ruby/annotation.rb b/lib/rbs/ast/ruby/annotation.rb
@@ -395,6 +395,15 @@ def initialize(location:, prefix_location:, self_location:, dot_location:, var_n
             @comment = comment
           end
         end
+
+        class EmbeddedRBSAnnotation < Base
+          attr_reader :members
+
+          def initialize(location:, prefix_location:, members:)
+            super(location, prefix_location)
+            @members = members
+          end
+        end
       end
     end
   end

diff --git a/sig/ast/ruby/annotation.rbs b/sig/ast/ruby/annotation.rbs
@@ -483,6 +483,14 @@ module RBS
             comment: Location?
           ) -> void
         end
+
+        class EmbeddedRBSAnnotation < Base
+          type member = AST::Declarations::t | AST::Members::t
+
+          attr_reader members: Array[member]
+
+          def initialize: (location: Location, prefix_location: Location, members: Array[member]) -> void
+        end
       end
     end
   end

diff --git a/src/constants.c b/src/constants.c
@@ -55,6 +55,7 @@ VALUE RBS_AST_Ruby_Annotation_ClassIvarTypeAnnotation;
 VALUE RBS_AST_Ruby_Annotation_ClassVarTypeAnnotation;
 VALUE RBS_AST_Ruby_Annotation_ColonMethodTypeAnnotation;
 VALUE RBS_AST_Ruby_Annotation_DoubleSplatParamTypeAnnotation;
+VALUE RBS_AST_Ruby_Annotation_EmbeddedRBSAnnotation;
 VALUE RBS_AST_Ruby_Annotation_GenericAnnotation;
 VALUE RBS_AST_Ruby_Annotation_InheritsAnnotation;
 VALUE RBS_AST_Ruby_Annotation_IvarTypeAnnotation;
@@ -151,6 +152,7 @@ void rbs__init_constants(void) {
   IMPORT_CONSTANT(RBS_AST_Ruby_Annotation_ClassVarTypeAnnotation, RBS_AST_Ruby_Annotation, "ClassVarTypeAnnotation");
   IMPORT_CONSTANT(RBS_AST_Ruby_Annotation_ColonMethodTypeAnnotation, RBS_AST_Ruby_Annotation, "ColonMethodTypeAnnotation");
   IMPORT_CONSTANT(RBS_AST_Ruby_Annotation_DoubleSplatParamTypeAnnotation, RBS_AST_Ruby_Annotation, "DoubleSplatParamTypeAnnotation");
+  IMPORT_CONSTANT(RBS_AST_Ruby_Annotation_EmbeddedRBSAnnotation, RBS_AST_Ruby_Annotation, "EmbeddedRBSAnnotation");
   IMPORT_CONSTANT(RBS_AST_Ruby_Annotation_GenericAnnotation, RBS_AST_Ruby_Annotation, "GenericAnnotation");
   IMPORT_CONSTANT(RBS_AST_Ruby_Annotation_InheritsAnnotation, RBS_AST_Ruby_Annotation, "InheritsAnnotation");
   IMPORT_CONSTANT(RBS_AST_Ruby_Annotation_IvarTypeAnnotation, RBS_AST_Ruby_Annotation, "IvarTypeAnnotation");

diff --git a/src/ruby_objs.c b/src/ruby_objs.c
@@ -522,6 +522,19 @@ VALUE rbs_ast_ruby_annotation_double_splat_param_type_annotation(VALUE location,
   );
 }
 
+VALUE rbs_ast_ruby_annotation_embedded_rbs_annotation(VALUE location, VALUE prefix_location, VALUE members) {
+  VALUE _init_kwargs = rb_hash_new();
+  rb_hash_aset(_init_kwargs, ID2SYM(rb_intern("location")), location);
+  rb_hash_aset(_init_kwargs, ID2SYM(rb_intern("prefix_location")), prefix_location);
+  rb_hash_aset(_init_kwargs, ID2SYM(rb_intern("members")), members);
+
+  return CLASS_NEW_INSTANCE(
+    RBS_AST_Ruby_Annotation_EmbeddedRBSAnnotation,
+    1,
+    &_init_kwargs
+  );
+}
+
 VALUE rbs_ast_ruby_annotation_generic_annotation(VALUE location, VALUE prefix_location, VALUE generic_location, VALUE unchecked_location, VALUE variance_location, VALUE name_location, VALUE upper_bound_operator_location, VALUE upper_bound, VALUE default_type_operator_location, VALUE default_type, VALUE comment) {
   VALUE _init_kwargs = rb_hash_new();
   rb_hash_aset(_init_kwargs, ID2SYM(rb_intern("location")), location);

diff --git a/test/rbs/inline_annotation_parsing_test.rb b/test/rbs/inline_annotation_parsing_test.rb
@@ -386,7 +386,6 @@ def test_instance_variable_annotation
   end
 
   def test_embbed_syntax
-    omit("Implement `@rbs!` annotation")
     Parser.parse_inline(<<~RBS, 0...).tap do |annot|
         @rbs!
           interface _Hello
@@ -395,6 +394,20 @@ def foo: () -> void
 
           type world = top
       RBS
+
+      assert_instance_of AST::Ruby::Annotation::EmbeddedRBSAnnotation, annot
+
+      assert_equal <<~RBS.chomp, annot.location.source
+        @rbs!
+          interface _Hello
+            def foo: () -> void
+          end
+
+          type world = top
+      RBS
+
+      assert_equal "@rbs!", annot.prefix_location.source
+      assert_equal [AST::Declarations::Interface, AST::Declarations::TypeAlias], annot.members.map(&:class)
     end
   end