Add type arguments support to singleton types

allcre · allcre · commit e20d33627afa · 2025-05-20T14:55:17.000-04:00
diff --git a/config.yml b/config.yml
@@ -363,6 +363,8 @@ nodes:
     fields:
       - name: name
         c_type: rbs_type_name
+      - name: args
+        c_type: rbs_node_list
   - name: RBS::Types::Function
     expose_location: false
     fields:
diff --git a/docs/syntax.md b/docs/syntax.md
@@ -3,17 +3,17 @@
 ## Types
 
 ```markdown
-_type_ ::= _class-name_ _type-arguments_                 (Class instance type)
-         | _interface-name_ _type-arguments_             (Interface type)
-         | _alias-name_ _type-arguments_                 (Alias type)
-         | `singleton(` _class-name_ `)`                 (Class singleton type)
-         | _literal_                                     (Literal type)
-         | _type_ `|` _type_                             (Union type)
-         | _type_ `&` _type_                             (Intersection type)
-         | _type_ `?`                                    (Optional type)
-         | `{` _record-name_ `:` _type_ `,` etc. `}`     (Record type)
-         | `[]` | `[` _type_ `,` etc. `]`                (Tuples)
-         | _type-variable_                               (Type variables)
+_type_ ::= _class-name_ _type-arguments_                     (Class instance type)
+         | _interface-name_ _type-arguments_                 (Interface type)
+         | _alias-name_ _type-arguments_                     (Alias type)
+         | `singleton(` _class-name_ `)` _type-arguments_    (Class singleton type)
+         | _literal_                                         (Literal type)
+         | _type_ `|` _type_                                 (Union type)
+         | _type_ `&` _type_                                 (Intersection type)
+         | _type_ `?`                                        (Optional type)
+         | `{` _record-name_ `:` _type_ `,` etc. `}`         (Record type)
+         | `[]` | `[` _type_ `,` etc. `]`                    (Tuples)
+         | _type-variable_                                   (Type variables)
          | `self`
          | `instance`
          | `class`
@@ -85,7 +85,8 @@ Class singleton type denotes _the type of a singleton object of a class_.
 
 ```rbs
 singleton(String)
-singleton(::Hash)            # Class singleton type cannot be parametrized.
+singleton(::Hash)            # Class singleton type
+singleton(Array)[String]     # Class singleton type with type application
 ```
 
 ### Literal type
diff --git a/ext/rbs_extension/ast_translation.c b/ext/rbs_extension/ast_translation.c
diff --git a/include/rbs/ast.h b/include/rbs/ast.h
@@ -555,6 +555,7 @@ typedef struct rbs_types_class_singleton {
     rbs_node_t base;
 
     struct rbs_type_name *name;
+    struct rbs_node_list *args;
 } rbs_types_class_singleton_t;
 
 typedef struct rbs_types_function {
@@ -729,7 +730,7 @@ rbs_types_bases_top_t *rbs_types_bases_top_new(rbs_allocator_t *allocator, rbs_l
 rbs_types_bases_void_t *rbs_types_bases_void_new(rbs_allocator_t *allocator, rbs_location_t *location);
 rbs_types_block_t *rbs_types_block_new(rbs_allocator_t *allocator, rbs_location_t *location, rbs_node_t *type, bool required, rbs_node_t *self_type);
 rbs_types_class_instance_t *rbs_types_class_instance_new(rbs_allocator_t *allocator, rbs_location_t *location, rbs_type_name_t *name, rbs_node_list_t *args);
-rbs_types_class_singleton_t *rbs_types_class_singleton_new(rbs_allocator_t *allocator, rbs_location_t *location, rbs_type_name_t *name);
+rbs_types_class_singleton_t *rbs_types_class_singleton_new(rbs_allocator_t *allocator, rbs_location_t *location, rbs_type_name_t *name, rbs_node_list_t *args);
 rbs_types_function_t *rbs_types_function_new(rbs_allocator_t *allocator, rbs_location_t *location, rbs_node_list_t *required_positionals, rbs_node_list_t *optional_positionals, rbs_node_t *rest_positionals, rbs_node_list_t *trailing_positionals, rbs_hash_t *required_keywords, rbs_hash_t *optional_keywords, rbs_node_t *rest_keywords, rbs_node_t *return_type);
 rbs_types_function_param_t *rbs_types_function_param_new(rbs_allocator_t *allocator, rbs_location_t *location, rbs_node_t *type, rbs_ast_symbol_t *name);
 rbs_types_interface_t *rbs_types_interface_new(rbs_allocator_t *allocator, rbs_location_t *location, rbs_type_name_t *name, rbs_node_list_t *args);
diff --git a/lib/rbs/types.rb b/lib/rbs/types.rb
@@ -202,38 +202,70 @@ def with_nonreturn_void?
     class ClassSingleton
       attr_reader :name
       attr_reader :location
+      attr_reader :args
 
-      def initialize(name:, location:)
+      def initialize(name:, location:, args: [])
         @name = name
         @location = location
+        @args = args
       end
 
       def ==(other)
-        other.is_a?(ClassSingleton) && other.name == name
+        other.is_a?(ClassSingleton) && other.name == name && other.args == args
       end
 
       alias eql? ==
 
       def hash
-        self.class.hash ^ name.hash
+        self.class.hash ^ name.hash ^ args.hash
       end
 
       include NoFreeVariables
-      include NoSubst
+
+      def sub(s)
+        return self if s.empty?
+
+        self.class.new(name: name,
+                       args: args.map {|ty| ty.sub(s) },
+                       location: location)
+      end
 
       def to_json(state = _ = nil)
-        { class: :class_singleton, name: name, location: location }.to_json(state)
+        { class: :class_singleton, name: name, args: args, location: location }.to_json(state)
       end
 
       def to_s(level = 0)
-        "singleton(#{name})"
+        if args.empty?
+          "singleton(#{name})"
+        else
+          "singleton(#{name})[#{args.join(", ")}]"
+        end
       end
 
-      include EmptyEachType
+      def each_type(&block)
+        if block
+          args.each(&block)
+        else
+          enum_for :each_type
+        end
+      end
 
-      def map_type_name(&)
+      def map_type(&block)
+        if block
+          ClassSingleton.new(
+            name: name,
+            args: args.map {|type| yield type },
+            location: location
+          )
+        else
+          enum_for :map_type
+        end
+      end
+
+      def map_type_name(&block)
         ClassSingleton.new(
           name: yield(name, location, self),
+          args: args.map {|type| type.map_type_name(&block) },
           location: location
         )
       end
diff --git a/lib/rbs/unit_test/type_assertions.rb b/lib/rbs/unit_test/type_assertions.rb
@@ -231,15 +231,17 @@ def method_defs(method)
         type, definition = target
 
         case type
-        when Types::ClassInstance
-          subst = RBS::Substitution.build(definition.type_params, type.args)
-          definition.methods[method].defs.map do |type_def|
-            type_def.update(
-              type: type_def.type.sub(subst)
-            )
+        when Types::ClassInstance, Types::ClassSingleton
+          if type.is_a?(Types::ClassSingleton) && type.args.empty?
+            definition.methods[method].defs
+          else
+            subst = RBS::Substitution.build(definition.type_params, type.args)
+            definition.methods[method].defs.map do |type_def|
+              type_def.update(
+                type: type_def.type.sub(subst)
+              )
+            end
           end
-        when Types::ClassSingleton
-          definition.methods[method].defs
         else
           raise
         end
diff --git a/src/ast.c b/src/ast.c
@@ -1138,7 +1138,7 @@ rbs_types_class_instance_t *rbs_types_class_instance_new(rbs_allocator_t *alloca
     return instance;
 }
 #line 156 "prism/templates/src/ast.c.erb"
-rbs_types_class_singleton_t *rbs_types_class_singleton_new(rbs_allocator_t *allocator, rbs_location_t *location, rbs_type_name_t *name) {
+rbs_types_class_singleton_t *rbs_types_class_singleton_new(rbs_allocator_t *allocator, rbs_location_t *location, rbs_type_name_t *name, rbs_node_list_t *args) {
     rbs_types_class_singleton_t *instance = rbs_allocator_alloc(allocator, rbs_types_class_singleton_t);
 
     *instance = (rbs_types_class_singleton_t) {
@@ -1147,6 +1147,7 @@ rbs_types_class_singleton_t *rbs_types_class_singleton_new(rbs_allocator_t *allo
             .location = location,
         },
         .name = name,
+        .args = args,
     };
 
     return instance;
diff --git a/src/parser.c b/src/parser.c
@@ -208,7 +208,7 @@ static bool parse_type_name(rbs_parser_t *parser, TypeNameKind kind, rbs_range_t
         goto error_handling;
     }
 
-success: {
+success : {
     if (rg) {
         rg->end = parser->current_token.range.end;
     }
@@ -220,7 +220,7 @@ success: {
     return true;
 }
 
-error_handling: {
+error_handling : {
     const char *ids = NULL;
     if (kind & ALIAS_NAME) {
         ids = "alias name";
@@ -1019,7 +1019,7 @@ static bool parse_instance_type(rbs_parser_t *parser, bool parse_alias, rbs_node
 }
 
 /*
-  singleton_type ::= {`singleton`} `(` type_name <`)`>
+  singleton_type ::= {`singleton`} `(` type_name <`)`> type_args?
 */
 NODISCARD
 static bool parse_singleton_type(rbs_parser_t *parser, rbs_types_class_singleton_t **singleton) {
@@ -1035,13 +1035,29 @@ static bool parse_singleton_type(rbs_parser_t *parser, rbs_types_class_singleton
     CHECK_PARSE(parse_type_name(parser, CLASS_NAME, &name_range, &type_name));
 
     ADVANCE_ASSERT(parser, pRPAREN);
-    type_range.end = parser->current_token.range.end;
+
+    rbs_node_list_t *types = rbs_node_list_new(ALLOCATOR());
+
+    rbs_range_t args_range;
+    if (parser->next_token.type == pLBRACKET) {
+        rbs_parser_advance(parser);
+        args_range.start = parser->current_token.range.start;
+        CHECK_PARSE(parse_type_list(parser, pRBRACKET, types));
+        ADVANCE_ASSERT(parser, pRBRACKET);
+        args_range.end = parser->current_token.range.end;
+        type_range.end = parser->current_token.range.end;
+    } else {
+        args_range = NULL_RANGE;
+        type_range.end = parser->current_token.range.end;
+    }
 
     rbs_location_t *loc = rbs_location_new(ALLOCATOR(), type_range);
-    rbs_loc_alloc_children(ALLOCATOR(), loc, 1);
+    rbs_loc_alloc_children(ALLOCATOR(), loc, 2);
     rbs_loc_add_required_child(loc, INTERN("name"), name_range);
+    rbs_loc_add_optional_child(loc, INTERN("args"), args_range);
+
+    *singleton = rbs_types_class_singleton_new(ALLOCATOR(), loc, type_name, types);
 
-    *singleton = rbs_types_class_singleton_new(ALLOCATOR(), loc, type_name);
     return true;
 }
 
diff --git a/test/rbs/singleton_type_test.rb b/test/rbs/singleton_type_test.rb
@@ -0,0 +1,121 @@
+require "test_helper"
+
+class RBS::SingletonTypeTest < Test::Unit::TestCase
+  include TestHelper
+
+  Parser = RBS::Parser
+  Buffer = RBS::Buffer
+  Types = RBS::Types
+  TypeName = RBS::TypeName
+  Namespace = RBS::Namespace
+
+  def test_singleton_type_with_arguments
+    Parser.parse_type("singleton(Array)[String]").yield_self do |type|
+      assert_instance_of Types::ClassSingleton, type
+      assert_equal TypeName.new(namespace: Namespace.empty, name: :Array), type.name
+      assert_equal 1, type.args.size
+      assert_instance_of Types::ClassInstance, type.args[0]
+      assert_equal TypeName.new(namespace: Namespace.empty, name: :String), type.args[0].name
+      assert_equal "singleton(Array)[String]", type.location.source
+    end
+
+    Parser.parse_type("singleton(Hash)[Symbol, Integer]").yield_self do |type|
+      assert_instance_of Types::ClassSingleton, type
+      assert_equal TypeName.new(namespace: Namespace.empty, name: :Hash), type.name
+      assert_equal 2, type.args.size
+      assert_instance_of Types::ClassInstance, type.args[0]
+      assert_instance_of Types::ClassInstance, type.args[1]
+      assert_equal TypeName.new(namespace: Namespace.empty, name: :Symbol), type.args[0].name
+      assert_equal TypeName.new(namespace: Namespace.empty, name: :Integer), type.args[1].name
+      assert_equal "singleton(Hash)[Symbol, Integer]", type.location.source
+    end
+
+    Parser.parse_type("singleton(::Foo::Bar)[Baz]").yield_self do |type|
+      assert_instance_of Types::ClassSingleton, type
+      assert_equal TypeName.new(namespace: Namespace.parse("::Foo"), name: :Bar), type.name
+      assert_equal 1, type.args.size
+      assert_instance_of Types::ClassInstance, type.args[0]
+      assert_equal TypeName.new(namespace: Namespace.empty, name: :Baz), type.args[0].name
+      assert_equal "singleton(::Foo::Bar)[Baz]", type.location.source
+    end
+  end
+
+  def test_singleton_type_equality
+    type1 = parse_type("singleton(Array)[String]")
+    type2 = parse_type("singleton(Array)[String]")
+    type3 = parse_type("singleton(Array)[Integer]")
+    type4 = parse_type("singleton(Hash)[String]")
+
+    assert_equal type1, type2
+    refute_equal type1, type3
+    refute_equal type1, type4
+  end
+
+  def test_singleton_type_hash
+    type1 = parse_type("singleton(Array)[String]")
+    type2 = parse_type("singleton(Array)[String]")
+    type3 = parse_type("singleton(Array)[Integer]")
+
+    assert_equal type1.hash, type2.hash
+    refute_equal type1.hash, type3.hash
+  end
+
+  def test_singleton_type_sub
+    type = parse_type("singleton(Array)[T]", variables: [:T])
+    subst = RBS::Substitution.build([:T], [parse_type("String")])
+
+    result = type.sub(subst)
+    assert_instance_of Types::ClassSingleton, result
+    assert_equal TypeName.new(namespace: Namespace.empty, name: :Array), result.name
+    assert_equal 1, result.args.size
+    assert_instance_of Types::ClassInstance, result.args[0]
+    assert_equal TypeName.new(namespace: Namespace.empty, name: :String), result.args[0].name
+  end
+
+  def test_singleton_type_map_type_name
+    type = parse_type("singleton(Array)[String]")
+
+    mapped = type.map_type_name do |name, _, _|
+      TypeName.new(namespace: Namespace.empty, name: :List)
+    end
+
+    assert_instance_of Types::ClassSingleton, mapped
+    assert_equal TypeName.new(namespace: Namespace.empty, name: :List), mapped.name
+    assert_equal 1, mapped.args.size
+    assert_instance_of Types::ClassInstance, mapped.args[0]
+    assert_equal TypeName.new(namespace: Namespace.empty, name: :List), mapped.args[0].name
+  end
+
+  def test_singleton_type_each_type
+    type = parse_type("singleton(Array)[String, Integer]")
+
+    types = []
+    type.each_type { |t| types << t }
+
+    assert_equal 2, types.size
+    assert_instance_of Types::ClassInstance, types[0]
+    assert_instance_of Types::ClassInstance, types[1]
+    assert_equal TypeName.new(namespace: Namespace.empty, name: :String), types[0].name
+    assert_equal TypeName.new(namespace: Namespace.empty, name: :Integer), types[1].name
+  end
+
+  def test_singleton_type_map_type
+    type = parse_type("singleton(Array)[String, Integer]")
+
+    mapped = type.map_type do |t|
+      if t.is_a?(Types::ClassInstance) && t.name.name == :String
+        parse_type("Symbol")
+      else
+        t
+      end
+    end
+
+    assert_instance_of Types::ClassSingleton, mapped
+    assert_equal TypeName.new(namespace: Namespace.empty, name: :Array), mapped.name
+    assert_equal 2, mapped.args.size
+    assert_instance_of Types::ClassInstance, mapped.args[0]
+    assert_instance_of Types::ClassInstance, mapped.args[1]
+    assert_equal TypeName.new(namespace: Namespace.empty, name: :Symbol), mapped.args[0].name
+    assert_equal TypeName.new(namespace: Namespace.empty, name: :Integer), mapped.args[1].name
+  end
+end
diff --git a/test/rbs/type_parsing_test.rb b/test/rbs/type_parsing_test.rb
@@ -257,6 +257,7 @@ def test_class_singleton
       assert_instance_of Types::ClassSingleton, type
 
       assert_equal TypeName.new(namespace: Namespace.empty, name: :Object), type.name
+      assert_equal [], type.args
 
       assert_equal "singleton(Object)", type.location.source
     end
@@ -265,10 +266,21 @@ def test_class_singleton
       assert_instance_of Types::ClassSingleton, type
 
       assert_equal TypeName.new(namespace: Namespace.root, name: :Object), type.name
-
+      assert_equal [], type.args
       assert_equal "singleton(::Object)", type.location.source
     end
 
+    Parser.parse_type("singleton(Array)[String]").yield_self do |type|
+      assert_instance_of Types::ClassSingleton, type
+
+      assert_equal TypeName.new(namespace: Namespace.empty, name: :Array), type.name
+      assert_equal 1, type.args.size
+      assert_instance_of Types::ClassInstance, type.args[0]
+      assert_equal TypeName.new(namespace: Namespace.empty, name: :String), type.args[0].name
+
+      assert_equal "singleton(Array)[String]", type.location.source
+    end
+
     assert_raises RBS::ParsingError do
       Parser.parse_type("singleton(foo)")
     end
diff --git a/test/rbs/types_test.rb b/test/rbs/types_test.rb
