Add RBS::Source::RBS class

Author: soutaro

lib/rbs.rb

@@ -24,6 +24,7 @@
 require "rbs/ast/members"
 require "rbs/ast/annotation"
 require "rbs/ast/visitor"
+require "rbs/source"
 require "rbs/environment"
 require "rbs/environment/use_map"
 require "rbs/environment/class_entry"

lib/rbs/cli.rb

@@ -175,10 +175,9 @@ def run_ast(args, options)
 
       env = Environment.from_loader(loader).resolve_type_names
 
-      decls = env.declarations.select do |decl|
-        loc = decl.location or raise
+      decls = env.sources.select do |source|
         # @type var name: String
-        name = loc.buffer.name
+        name = source.buffer.name.to_s
 
         patterns.empty? || patterns.any? do |pat|
           case pat
@@ -188,7 +187,7 @@ def run_ast(args, options)
             name.end_with?(pat) || File.fnmatch(pat, name, File::FNM_EXTGLOB)
           end
         end
-      end
+      end.flat_map { _1.declarations }
 
       stdout.print JSON.generate(decls)
       stdout.flush

lib/rbs/environment.rb

@@ -2,16 +2,14 @@
 
 module RBS
   class Environment
-    attr_reader :declarations
-
     attr_reader :class_decls
     attr_reader :interface_decls
     attr_reader :type_alias_decls
     attr_reader :constant_decls
     attr_reader :global_decls
     attr_reader :class_alias_decls
 
-    attr_reader :signatures
+    attr_reader :sources
 
     class SingleEntry
       attr_reader :name
@@ -44,9 +42,7 @@ class GlobalEntry < SingleEntry
     end
 
     def initialize
-      @signatures = {}
-      @declarations = []
-
+      @sources = []
       @class_decls = {}
       @interface_decls = {}
       @type_alias_decls = {}
@@ -57,9 +53,7 @@ def initialize
     end
 
     def initialize_copy(other)
-      @signatures = other.signatures.dup
-      @declarations = other.declarations.dup
-
+      @sources = other.sources.dup
       @class_decls = other.class_decls.dup
       @interface_decls = other.interface_decls.dup
       @type_alias_decls = other.type_alias_decls.dup
@@ -368,16 +362,21 @@ def insert_decl(decl, context:, namespace:)
       end
     end
 
-    def <<(decl)
-      declarations << decl
-      insert_decl(decl, context: nil, namespace: Namespace.root)
-      self
+    def add_source(source)
+      sources << source
+
+      source.declarations.each do |decl|
+        insert_decl(decl, context: nil, namespace: Namespace.root)
+      end
     end
 
-    def add_signature(buffer:, directives:, decls:)
-      signatures[buffer] = [directives, decls]
-      decls.each do |decl|
-        self << decl
+    def each_rbs_source(&block)
+      if block
+        sources.each do |source|
+          yield source
+        end
+      else
+        enum_for(:each_rbs_source)
       end
     end
 
@@ -420,12 +419,14 @@ def resolve_type_names(only: nil)
       table.known_types.merge(interface_decls.keys)
       table.compute_children
 
-      signatures.each do |buffer, (dirs, decls)|
-        resolve = dirs.find { _1.is_a?(AST::Directives::ResolveTypeNames) } #: AST::Directives::ResolveTypeNames?
+      each_rbs_source do |source|
+        resolve = source.directives.find { _1.is_a?(AST::Directives::ResolveTypeNames) } #: AST::Directives::ResolveTypeNames?
         if !resolve || resolve.value
-          _, decls = resolve_signature(resolver, table, dirs, decls)
+          _, decls = resolve_signature(resolver, table, source.directives, source.declarations)
+        else
+          decls = source.declarations
         end
-        env.add_signature(buffer: buffer, directives: dirs, decls: decls)
+        env.add_source(Source::RBS.new(source.buffer, source.directives, decls))
       end
 
       env
@@ -713,15 +714,16 @@ def inspect
     end
 
     def buffers
-      signatures.keys
+      sources.map(&:buffer)
     end
 
     def unload(buffers)
       env = Environment.new
+      bufs = buffers.to_set
 
-      signatures.each do |buf, (dirs, decls)|
-        next if buffers.include?(buf)
-        env.add_signature(buffer: buf, directives: dirs, decls: decls)
+      each_rbs_source do |source|
+        next if bufs.include?(source.buffer)
+        env.add_source(source)
       end
 
       env

lib/rbs/environment_loader.rb

@@ -122,7 +122,7 @@ def load(env:)
         decls.each do |decl|
           loaded << [decl, path, source]
         end
-        env.add_signature(buffer: buffer, directives: dirs, decls: decls)
+        env.add_source(Source::RBS.new(buffer, dirs, decls))
       end
 
       loaded

lib/rbs/source.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+module RBS
+  module Source
+    class RBS
+      attr_reader :buffer, :directives, :declarations
+
+      def initialize(buffer, directives, decls)
+        @buffer = buffer
+        @directives = directives
+        @declarations = decls
+      end
+
+      def each_type_name(&block)
+        if block
+          set = Set[] #: Set[TypeName]
+          declarations.each do |decl|
+            each_declaration_type_name(set, decl, &block)
+          end
+        else
+          enum_for :each_type_name
+        end
+      end
+
+      def each_declaration_type_name(names, decl, &block)
+        case decl
+        when AST::Declarations::Class
+          decl.each_decl { each_declaration_type_name(names, _1, &block) }
+          type_name = decl.name
+        when AST::Declarations::Module
+          decl.each_decl { each_declaration_type_name(names, _1, &block) }
+          type_name = decl.name
+        when AST::Declarations::Interface
+          type_name = decl.name
+        when AST::Declarations::TypeAlias
+          type_name = decl.name
+        when AST::Declarations::ModuleAlias
+          type_name = decl.new_name
+        when AST::Declarations::ClassAlias
+          type_name = decl.new_name
+        end
+
+        if type_name
+          unless names.include?(type_name)
+            yield type_name
+            names << type_name
+          end
+        end
+      end
+    end
+  end
+end

sig/environment.rbs

@@ -30,8 +30,9 @@ module RBS
     class GlobalEntry < SingleEntry[Symbol, AST::Declarations::Global]
     end
 
-    # Top level declarations
-    attr_reader declarations: Array[AST::Declarations::t]
+    # Array of source objects loaded in the environment
+    #
+    attr_reader sources: Array[Source::RBS]
 
     # Class declarations
     attr_reader class_decls: Hash[TypeName, ModuleEntry | ClassEntry]
@@ -50,10 +51,6 @@ module RBS
     # Global declarations
     attr_reader global_decls: Hash[Symbol, GlobalEntry]
 
-    # A hash from Buffer to it's contents
-    #
-    attr_reader signatures: Hash[Buffer, [Array[AST::Directives::t], Array[AST::Declarations::t]]]
-
     def initialize: () -> void
 
     def initialize_copy: (Environment) -> void
@@ -62,11 +59,10 @@ module RBS
     #
     def self.from_loader: (EnvironmentLoader) -> Environment
 
-    def add_signature: (buffer: Buffer, directives: Array[AST::Directives::t], decls: Array[AST::Declarations::t]) -> void
+    def add_source: (Source::RBS) -> void
 
-    # Insert a toplevel declaration into the environment
-    #
-    def <<: (AST::Declarations::t decl) -> self
+    def each_rbs_source: () { (Source::RBS) -> void } -> void
+                      | () -> Enumerator[Source::RBS]
 
     # Insert a declaration into the environment
     #

sig/source.rbs

@@ -0,0 +1,22 @@
+module RBS
+  module Source
+    class RBS
+      attr_reader buffer: Buffer
+
+      attr_reader directives: Array[AST::Directives::t]
+
+      attr_reader declarations: Array[AST::Declarations::t]
+
+      def initialize: (Buffer, Array[AST::Directives::t], Array[AST::Declarations::t]) -> void
+
+      # Enumerates defined type names in the source
+      #
+      # The order is undefined. Deduplicated per source object.
+      #
+      def each_type_name: () { (TypeName) -> void } -> void
+                        | () -> Enumerator[TypeName]
+
+      private def each_declaration_type_name: (Set[TypeName], AST::Declarations::t) { (TypeName) -> void } -> void
+    end
+  end
+end

test/rbs/definition_builder_test.rb

@@ -2323,7 +2323,7 @@ class Foo
 end
       DEF
       RBS::Parser.parse_signature(rbs).tap do |buf, dirs, decls|
-        env.add_signature(buffer: buf, directives: dirs, decls: decls)
+        env.add_source(RBS::Source::RBS.new(buf, dirs, decls))
       end
       definition_builder = RBS::DefinitionBuilder.new(env: env.resolve_type_names)
       definition_builder.build_instance(RBS::TypeName.parse("::Foo")).tap do |defn|