Support for Const = Struct.new(…) syntax

lib/rbs/inline/ast/declarations.rb

@@ -28,7 +28,7 @@ def value_node(node)
         end
 
         # @rbs!
-        #   type t = ClassDecl | ModuleDecl | ConstantDecl | SingletonClassDecl | BlockDecl | DataAssignDecl
+        #   type t = ClassDecl | ModuleDecl | ConstantDecl | SingletonClassDecl | BlockDecl | DataAssignDecl | StructAssignDecl
         #
         #  interface _WithComments
         #    def comments: () -> AnnotationParser::ParsingResult?
@@ -279,9 +279,53 @@ def module_class_annotation #: Annotations::ModuleDecl | Annotations::ClassDecl
           end
         end
 
+        # @rbs module-self _WithTypeDecls
+        module DataStructUtil
+          # @rbs!
+          #   interface _WithTypeDecls
+          #     def type_decls: () -> Hash[Integer, Annotations::TypeAssertion]
+          #
+          #     def each_attribute_argument: () { (Prism::Node) -> void } -> void
+          #
+          #     def comments: %a{pure} () -> AnnotationParser::ParsingResult?
+          #   end
+
+          # @rbs %a{pure}
+          # @rbs () { ([Symbol, Annotations::TypeAssertion?]) -> void } -> void
+          #    | () -> Enumerator[[Symbol, Annotations::TypeAssertion?], void]
+          def each_attribute(&block)
+            if block
+              each_attribute_argument do |arg|
+                if arg.is_a?(Prism::SymbolNode)
+                  if name = arg.value
+                    type = type_decls.fetch(arg.location.start_line, nil)
+                    yield [name.to_sym, type]
+                  end
+                end
+              end
+            else
+              enum_for :each_attribute
+            end
+          end
+
+          def class_annotations #: Array[RBS::AST::Annotation]
+            annotations = [] #: Array[RBS::AST::Annotation]
+
+            comments&.each_annotation do |annotation|
+              if annotation.is_a?(Annotations::RBSAnnotation)
+                annotations.concat annotation.annotations
+              end
+            end
+
+            annotations
+          end
+        end
+
         class DataAssignDecl < Base
           extend ConstantUtil
 
+          include DataStructUtil
+
           attr_reader :node #: Prism::ConstantWriteNode
 
           attr_reader :comments #: AnnotationParser::ParsingResult?
@@ -323,36 +367,140 @@ def self.data_define?(node)
             end
           end
 
+          # @rbs () { (Prism::Node) -> void } -> void
+          def each_attribute_argument(&block)
+            if args = data_define_node.arguments
+              args.arguments.each(&block)
+            end
+          end
+        end
+
+        class StructAssignDecl < Base
+          extend ConstantUtil
+
+          include DataStructUtil
+
+          attr_reader :node #: Prism::ConstantWriteNode
+
+          attr_reader :comments #: AnnotationParser::ParsingResult?
+
+          attr_reader :type_decls #: Hash[Integer, Annotations::TypeAssertion]
+
+          attr_reader :struct_new_node #: Prism::CallNode
+
+          # @rbs (Prism::ConstantWriteNode, Prism::CallNode, AnnotationParser::ParsingResult?, Hash[Integer, Annotations::TypeAssertion]) -> void
+          def initialize(node, struct_new_node, comments, type_decls)
+            @node = node
+            @comments = comments
+            @type_decls = type_decls
+            @struct_new_node = struct_new_node
+          end
+
+          def start_line #: Integer
+            node.location.start_line
+          end
+
           # @rbs %a{pure}
-          # @rbs () { ([Symbol, Annotations::TypeAssertion?]) -> void } -> void
-          #    | () -> Enumerator[[Symbol, Annotations::TypeAssertion?], void]
-          def each_attribute(&block)
-            if block
-              if args = data_define_node.arguments
-                args.arguments.each do |arg|
-                  if arg.is_a?(Prism::SymbolNode)
-                    if name = arg.value
-                      type = type_decls.fetch(arg.location.start_line, nil)
-                      yield [name.to_sym, type]
-                    end
+          # @rbs () -> TypeName?
+          def constant_name
+            TypeName.new(name: node.name, namespace: Namespace.empty)
+          end
+
+          # @rbs () { (Prism::Node) -> void } -> void
+          def each_attribute_argument(&block)
+            if args = struct_new_node.arguments
+              args.arguments.each do |arg|
+                next if arg.is_a?(Prism::KeywordHashNode)
+                next if arg.is_a?(Prism::StringNode)
+
+                yield arg
+              end
+            end
+          end
+
+          # @rbs (Prism::ConstantWriteNode) -> Prism::CallNode?
+          def self.struct_new?(node)
+            value = value_node(node)
+
+            if value.is_a?(Prism::CallNode)
+              if value.receiver.is_a?(Prism::ConstantReadNode)
+                if value.receiver.full_name.delete_prefix("::") == "Struct"
+                  if value.name == :new
+                    return value
                   end
                 end
               end
-            else
-              enum_for :each_attribute
             end
           end
 
-          def class_annotations #: Array[RBS::AST::Annotation]
-            annotations = [] #: Array[RBS::AST::Annotation]
+          # @rbs %a{pure}
+          def keyword_init? #: bool
+            if args = struct_new_node.arguments
+              args.arguments.each do |arg|
+                if arg.is_a?(Prism::KeywordHashNode)
+                  arg.elements.each do |assoc|
+                    if assoc.is_a?(Prism::AssocNode)
+                      if (key = assoc.key).is_a?(Prism::SymbolNode)
+                        if key.value == "keyword_init"
+                          value = assoc.value
+                          if value.is_a?(Prism::FalseNode)
+                            return false
+                          end
+                        end
+                      end
+                    end
+                  end
+                end
+              end
+            end
 
-            comments&.each_annotation do |annotation|
-              if annotation.is_a?(Annotations::RBSAnnotation)
-                annotations.concat annotation.annotations
+            true
+          end
+
+          # @rbs %a{pure}
+          def positional_init? #: bool
+            if args = struct_new_node.arguments
+              args.arguments.each do |arg|
+                if arg.is_a?(Prism::KeywordHashNode)
+                  arg.elements.each do |assoc|
+                    if assoc.is_a?(Prism::AssocNode)
+                      if (key = assoc.key).is_a?(Prism::SymbolNode)
+                        if key.value == "keyword_init"
+                          value = assoc.value
+                          if value.is_a?(Prism::TrueNode)
+                            return false
+                          end
+                        end
+                      end
+                    end
+                  end
+                end
               end
             end
 
-            annotations
+            true
+          end
+
+          # Returns `true` is annotation is given to make all attributes *readonly*
+          #
+          # Add `# @rbs %a{rbs-inline:readonly-attributes=true}` to the class to make all attributes `attr_reader`, instead of `attr_accessor`.
+          #
+          # @rbs %a{pure}
+          def readonly_attributes? #: bool
+            class_annotations.any? do |annotation|
+              annotation.string == "rbs-inline:readonly-attributes=true"
+            end
+          end
+
+          # Returns `true` if annotation is given to make all `.new` arguments required
+          #
+          # Add `# @rbs %a{rbs-inline:new-args=required}` to the class to make all of the parameters required.
+          #
+          # @rbs %a{pure}
+          def required_new_args? #: bool
+            class_annotations.any? do |annotation|
+              annotation.string == "rbs-inline:new-args=required"
+            end
           end
         end
       end

lib/rbs/inline/parser.rb

@@ -432,6 +432,19 @@ def visit_constant_write_node(node)
           end
 
           decl = AST::Declarations::DataAssignDecl.new(node, data_node, comment, type_decls)
+        when struct_node = AST::Declarations::StructAssignDecl.struct_new?(node)
+          type_decls = {} #: Hash[Integer, AST::Annotations::TypeAssertion]
+
+          inner_annotations(node.location.start_line, node.location.end_line).flat_map do |comment|
+            comment.each_annotation do |annotation|
+              if annotation.is_a?(AST::Annotations::TypeAssertion)
+                start_line = annotation.source.comments[0].location.start_line
+                type_decls[start_line] = annotation
+              end
+            end
+          end
+
+          decl = AST::Declarations::StructAssignDecl.new(node, struct_node, comment, type_decls)
         else
           assertion = assertion_annotation(node)
           decl = AST::Declarations::ConstantDecl.new(node, comment, assertion)

lib/rbs/inline/writer.rb

@@ -79,6 +79,8 @@ def translate_decl(decl, rbs)
           translate_constant_decl(decl, rbs)
         when AST::Declarations::DataAssignDecl
           translate_data_assign_decl(decl, rbs)
+        when AST::Declarations::StructAssignDecl
+          translate_struct_assign_decl(decl, rbs)
         when AST::Declarations::BlockDecl
           if decl.module_class_annotation
             case decl.module_class_annotation
@@ -133,7 +135,7 @@ def translate_members(members, decl, rbs)
             else
               translate_members(member.members, decl, rbs)
             end
-          when AST::Declarations::ClassDecl, AST::Declarations::ModuleDecl, AST::Declarations::ConstantDecl, AST::Declarations::DataAssignDecl
+          when AST::Declarations::ClassDecl, AST::Declarations::ModuleDecl, AST::Declarations::ConstantDecl, AST::Declarations::DataAssignDecl, AST::Declarations::StructAssignDecl
             translate_decl(member, rbs)
           end
         end
@@ -270,6 +272,116 @@ def translate_data_assign_decl(decl, rbs) #: void
         )
       end
 
+      # @rbs decl: AST::Declarations::StructAssignDecl
+      # @rbs rbs: _Content
+      def translate_struct_assign_decl(decl, rbs) #: void
+        return unless decl.constant_name
+
+        if decl.comments
+          comment = RBS::AST::Comment.new(string: decl.comments.content(trim: true), location: nil)
+        end
+
+        attributes = decl.each_attribute.map do |name, type|
+          if decl.readonly_attributes?
+            RBS::AST::Members::AttrReader.new(
+              name: name,
+              type: type&.type || Types::Bases::Any.new(location: nil),
+              ivar_name: false,
+              comment: nil,
+              kind: :instance,
+              annotations: [],
+              visibility: nil,
+              location: nil
+            )
+          else
+            RBS::AST::Members::AttrAccessor.new(
+              name: name,
+              type: type&.type || Types::Bases::Any.new(location: nil),
+              ivar_name: false,
+              comment: nil,
+              kind: :instance,
+              annotations: [],
+              visibility: nil,
+              location: nil
+            )
+          end
+        end
+
+        init = RBS::AST::Members::MethodDefinition.new(
+          name: :initialize,
+          kind: :instance,
+          overloads: [],
+          annotations: [],
+          location: nil,
+          comment: nil,
+          overloading: false,
+          visibility: nil
+        )
+
+        if decl.positional_init?
+          attr_params = decl.each_attribute.map do |name, attr|
+            RBS::Types::Function::Param.new(
+              type: attr&.type || Types::Bases::Any.new(location: nil),
+              name: name,
+              location: nil
+            )
+          end
+
+          method_type = Types::Function.empty(Types::Bases::Void.new(location: nil))
+          if decl.required_new_args?
+            method_type = method_type.update(required_positionals: attr_params)
+          else
+            method_type = method_type.update(optional_positionals: attr_params)
+          end
+
+          init.overloads <<
+            RBS::AST::Members::MethodDefinition::Overload.new(
+              method_type: RBS::MethodType.new(type_params: [], type: method_type, block: nil, location: nil),
+              annotations: []
+            )
+        end
+
+        if decl.keyword_init?
+          attr_keywords = decl.each_attribute.map do |name, attr|
+            [
+              name,
+              RBS::Types::Function::Param.new(
+                type: attr&.type || Types::Bases::Any.new(location: nil),
+                name: nil,
+                location: nil
+              )
+            ]
+          end.to_h #: Hash[Symbol, RBS::Types::Function::Param]
+
+          method_type = Types::Function.empty(Types::Bases::Void.new(location: nil))
+          if decl.required_new_args?
+            method_type = method_type.update(required_keywords: attr_keywords)
+          else
+            method_type = method_type.update(optional_keywords: attr_keywords)
+          end
+
+          init.overloads <<
+            RBS::AST::Members::MethodDefinition::Overload.new(
+              method_type: RBS::MethodType.new(type_params: [], type: method_type, block: nil, location: nil),
+              annotations: []
+            )
+        end
+
+        rbs << RBS::AST::Declarations::Class.new(
+          name: decl.constant_name,
+          type_params: [],
+          members: [*attributes, init],
+          super_class: RBS::AST::Declarations::Class::Super.new(
+            name: RBS::TypeName.new(name: :Struct, namespace: RBS::Namespace.empty),
+            args: [RBS::Types::Bases::Any.new(location: nil)],
+            location: nil
+          ),
+          annotations: decl.class_annotations,
+          location: nil,
+          comment: comment
+        )
+      end
+
       # @rbs decl: AST::Declarations::SingletonClassDecl
       # @rbs rbs: _Content
       # @rbs return: void