Merge pull request #2957 from tk0miya/claude/memoize-hash-and-resolver-cache

Flyweight TypeName / Namespace cached in a shared trie

Author: soutaro

ext/rbs_extension/ast_translation.c

@@ -177,6 +177,30 @@ VALUE rbs_type_param_variance_to_ruby(enum rbs_type_param_variance value) {
     rb_class_new_instance(argc, argv, receiver)
 #endif
 
+// Route Namespace / TypeName construction through the Ruby-side
+// flyweight cache (`RBS::Namespace.[]` / `RBS::TypeName.[]`) so that
+// structurally equal values produced by the parser share canonical
+// instances. An in-C trie walk was tried but did not beat
+// `rb_funcallv` on Ruby 4.0+, where method dispatch is well optimized.
+static ID id_intern_brackets;
+
+static inline ID intern_brackets(void) {
+    if (!id_intern_brackets) id_intern_brackets = rb_intern("[]");
+    return id_intern_brackets;
+}
+
+static VALUE rbs_intern_namespace(rbs_translation_context_t ctx, rbs_namespace_t *node) {
+    VALUE args[2];
+    args[0] = rbs_node_list_to_ruby_array(ctx, node->path);
+    args[1] = node->absolute ? Qtrue : Qfalse;
+    return rb_funcallv(RBS_Namespace, intern_brackets(), 2, args);
+}
+
+static VALUE rbs_intern_type_name(VALUE namespace, VALUE name) {
+    VALUE args[2] = { namespace, name };
+    return rb_funcallv(RBS_TypeName, intern_brackets(), 2, args);
+}
+
 VALUE rbs_struct_to_ruby_value(rbs_translation_context_t ctx, rbs_node_t *instance) {
     if (instance == NULL) return Qnil;
 
@@ -1378,19 +1402,7 @@ VALUE rbs_struct_to_ruby_value(rbs_translation_context_t ctx, rbs_node_t *instan
         return CLASS_NEW_INSTANCE(RBS_MethodType, 1, &h);
     }
     case RBS_NAMESPACE: {
-        rbs_namespace_t *node = (rbs_namespace_t *) instance;
-
-        // Compute child VALUEs into locals variables first, before any recursion into `rbs_struct_to_ruby_value()`.
-        VALUE arg_path = rbs_node_list_to_ruby_array(ctx, node->path);
-        VALUE arg_absolute = node->absolute ? Qtrue : Qfalse;
-
-        // Claim the shared kwargs hash, clear it, fill it, and hand it to `.new`.
-        // Must not recurse between `rb_hash_clear()` and `CLASS_NEW_INSTANCE()`.
-        VALUE h = ctx.reusable_kwargs_hash;
-        rb_hash_clear(h);
-        rb_hash_aset(h, ID2SYM(rb_intern("path")), arg_path);
-        rb_hash_aset(h, ID2SYM(rb_intern("absolute")), arg_absolute);
-        return CLASS_NEW_INSTANCE(RBS_Namespace, 1, &h);
+        return rbs_intern_namespace(ctx, (rbs_namespace_t *) instance);
     }
     case RBS_SIGNATURE: {
         rbs_signature_t *signature = (rbs_signature_t *) instance;
@@ -1402,18 +1414,9 @@ VALUE rbs_struct_to_ruby_value(rbs_translation_context_t ctx, rbs_node_t *instan
     }
     case RBS_TYPE_NAME: {
         rbs_type_name_t *node = (rbs_type_name_t *) instance;
-
-        // Compute child VALUEs into locals variables first, before any recursion into `rbs_struct_to_ruby_value()`.
-        VALUE arg_namespace = rbs_struct_to_ruby_value(ctx, (rbs_node_t *) node->rbs_namespace); // rbs_namespace
-        VALUE arg_name = rbs_struct_to_ruby_value(ctx, (rbs_node_t *) node->name);               // rbs_ast_symbol
-
-        // Claim the shared kwargs hash, clear it, fill it, and hand it to `.new`.
-        // Must not recurse between `rb_hash_clear()` and `CLASS_NEW_INSTANCE()`.
-        VALUE h = ctx.reusable_kwargs_hash;
-        rb_hash_clear(h);
-        rb_hash_aset(h, ID2SYM(rb_intern("namespace")), arg_namespace);
-        rb_hash_aset(h, ID2SYM(rb_intern("name")), arg_name);
-        return CLASS_NEW_INSTANCE(RBS_TypeName, 1, &h);
+        VALUE ns = rbs_struct_to_ruby_value(ctx, (rbs_node_t *) node->rbs_namespace);
+        VALUE name = rbs_struct_to_ruby_value(ctx, (rbs_node_t *) node->name);
+        return rbs_intern_type_name(ns, name);
     }
     case RBS_TYPES_ALIAS: {
         rbs_types_alias_t *node = (rbs_types_alias_t *) instance;

lib/rbs/environment.rb

@@ -173,7 +173,7 @@ def constant_entry(type_name, normalized: false)
           unless type_name.namespace.empty?
             parent = type_name.namespace.to_type_name
             normalized_parent = normalize_module_name?(parent) or return
-            constant_name = TypeName.new(name: type_name.name, namespace: normalized_parent.to_namespace)
+            constant_name = TypeName[normalized_parent.to_namespace, type_name.name]
             constant_decls.fetch(constant_name, nil)
           end
         end
@@ -193,7 +193,7 @@ def normalize_type_name?(name)
           parent = normalize_module_name?(parent)
           return parent unless parent
 
-          TypeName.new(namespace: parent.to_namespace, name: name.name)
+          TypeName[parent.to_namespace, name.name]
         else
           name
         end

lib/rbs/namespace.rb

@@ -9,30 +9,65 @@ def initialize(path:, absolute:)
       @absolute = absolute ? true : false
     end
 
+    # Process-wide flyweight cache. Two tries (one per `absolute` flag)
+    # keyed on path Symbols, with the cached Namespace stored under the
+    # `INTERN_LEAF` sentinel at each path's terminal node.
+    @intern_mutex = Mutex.new
+    @intern_trie_absolute = {}
+    @intern_trie_relative = {}
+    INTERN_LEAF = Module.new
+
+    # Returns a canonical `Namespace` instance for the given `path` /
+    # `absolute` pair. Repeated calls with structurally equal arguments
+    # return the same object, so callers can rely on `equal?` for fast
+    # equality. The path Array is duplicated and frozen on insert.
+    def self.[](path, absolute)
+      absolute = absolute ? true : false
+
+      # Lock-free fast path.
+      node = absolute ? @intern_trie_absolute : @intern_trie_relative
+      path.each do |sym|
+        node = node[sym]
+        break unless node
+      end
+      if node && (cached = node[INTERN_LEAF])
+        return cached
+      end
+
+      @intern_mutex.synchronize do
+        node = absolute ? @intern_trie_absolute : @intern_trie_relative
+        path.each { |sym| node = (node[sym] ||= {}) }
+        node[INTERN_LEAF] ||= begin
+          frozen_path = path.frozen? ? path : path.dup.freeze
+          new(path: frozen_path, absolute: absolute)
+        end
+      end
+    end
+
     def self.empty
-      @empty ||= new(path: [], absolute: false)
+      @empty ||= self[[], false]
     end
 
     def self.root
-      @root ||= new(path: [], absolute: true)
+      @root ||= self[[], true]
     end
 
     def +(other)
       if other.absolute?
         other
       else
-        self.class.new(path: path + other.path, absolute: absolute?)
+        Namespace[path + other.path, absolute?]
       end
     end
 
     def append(component)
-      self.class.new(path: path + [component], absolute: absolute?)
+      Namespace[path + [component], absolute?]
     end
 
     def parent
       @parent ||= begin
         raise "Parent with empty namespace" if empty?
-        self.class.new(path: path.take(path.size - 1), absolute: absolute?)
+        Namespace[path.take(path.size - 1), absolute?]
       end
     end
 
@@ -45,25 +80,26 @@ def relative?
     end
 
     def absolute!
-      self.class.new(path: path, absolute: true)
+      Namespace[path, true]
     end
 
     def relative!
-      self.class.new(path: path, absolute: false)
+      Namespace[path, false]
     end
 
     def empty?
       path.empty?
     end
 
     def ==(other)
+      return true if equal?(other)
       other.is_a?(Namespace) && other.path == path && other.absolute? == absolute?
     end
 
     alias eql? ==
 
     def hash
-      path.hash ^ absolute?.hash
+      @hash ||= path.hash ^ absolute?.hash
     end
 
     def split
@@ -87,14 +123,14 @@ def to_type_name
       raise unless name
       raise unless parent
 
-      TypeName.new(name: name, namespace: parent)
+      TypeName[parent, name]
     end
 
     def self.parse(string)
       if string.start_with?("::")
-        new(path: string.split("::").drop(1).map(&:to_sym), absolute: true)
+        self[string.split("::").drop(1).map(&:to_sym), true]
       else
-        new(path: string.split("::").map(&:to_sym), absolute: false)
+        self[string.split("::").map(&:to_sym), false]
       end
     end
 

lib/rbs/resolver/type_name_resolver.rb

@@ -29,19 +29,17 @@ def self.build(env)
         new(all_names, aliases)
       end
 
-      def try_cache(query)
-        cache.fetch(query) do
-          result = yield
-          cache[query] = result
-        end
+      def try_cache(type_name, context)
+        inner = cache[context] ||= {}
+        inner.fetch(type_name) { inner[type_name] = yield }
       end
 
       def resolve(type_name, context:)
         if type_name.absolute? && has_type_name?(type_name)
           return type_name
         end
 
-        try_cache([type_name, context]) do
+        try_cache(type_name, context) do
           if type_name.class?
             resolve_namespace0(type_name, context, Set.new) || nil
           else
@@ -51,7 +49,7 @@ def resolve(type_name, context:)
               resolve_type_name(type_name.name, context)
             else
               if namespace = resolve_namespace0(namespace.to_type_name, context, Set.new)
-                type_name = TypeName.new(name: type_name.name, namespace: namespace.to_namespace)
+                type_name = TypeName[namespace.to_namespace, type_name.name]
                 has_type_name?(type_name)
               end
             end
@@ -68,7 +66,7 @@ def resolve_namespace(type_name, context:)
           raise "Type name must be a class name: #{type_name}"
         end
 
-        try_cache([type_name, context]) do
+        try_cache(type_name, context) do
           ns = resolve_namespace0(type_name, context, Set.new) or return ns
         end
       end
@@ -93,10 +91,10 @@ def resolve_type_name(type_name, context)
             resolve_type_name(type_name, outer)
           else
             has_type_name?(inner) or raise "Context must be normalized: #{inner.inspect}"
-            has_type_name?(TypeName.new(name: type_name, namespace: inner.to_namespace)) || resolve_type_name(type_name, outer)
+            has_type_name?(TypeName[inner.to_namespace, type_name]) || resolve_type_name(type_name, outer)
           end
         else
-          type_name = TypeName.new(name: type_name, namespace: Namespace.root)
+          type_name = TypeName[Namespace.root, type_name]
           has_type_name?(type_name)
         end
       end
@@ -109,11 +107,11 @@ def resolve_head_namespace(head, context)
             resolve_head_namespace(head, outer)
           when TypeName
             has_type_name?(inner) or raise "Context must be normalized: #{inner.inspect}"
-            type_name = TypeName.new(name: head, namespace: inner.to_namespace)
+            type_name = TypeName[inner.to_namespace, head]
             has_type_name?(type_name) || aliased_name?(type_name) || resolve_head_namespace(head, outer)
           end
         else
-          type_name = TypeName.new(name: head, namespace: Namespace.root)
+          type_name = TypeName[Namespace.root, head]
           has_type_name?(type_name) || aliased_name?(type_name)
         end
       end
@@ -140,7 +138,7 @@ def resolve_namespace0(type_name, context, visited)
 
         head =
           if type_name.absolute?
-            root_name = TypeName.new(name: head, namespace: Namespace.root)
+            root_name = TypeName[Namespace.root, head]
             has_type_name?(root_name) || aliased_name?(root_name)
           else
             resolve_head_namespace(head, context)
@@ -152,7 +150,7 @@ def resolve_namespace0(type_name, context, visited)
           end
 
           tail.inject(head) do |namespace, name|
-            type_name = TypeName.new(name: name, namespace: namespace.to_namespace)
+            type_name = TypeName[namespace.to_namespace, name]
             case
             when has_type_name?(type_name)
               type_name

lib/rbs/type_name.rb

@@ -22,14 +22,40 @@ def initialize(namespace:, name:)
               end
     end
 
+    # Process-wide flyweight cache. Two-level Hash keyed by canonical
+    # Namespace identity (outer uses `compare_by_identity`) and name
+    # Symbol.
+    @intern_mutex = Mutex.new
+    @intern_cache = {}  #: Hash[Namespace, Hash[Symbol, TypeName]]
+    @intern_cache.compare_by_identity
+
+    # Returns a canonical `TypeName` instance for the given `namespace` /
+    # `name` pair. The namespace is canonicalized through `Namespace.[]`
+    # so identity-based lookup works regardless of the caller passing a
+    # fresh `Namespace.new` or an already-interned instance.
+    def self.[](namespace, name)
+      ns = Namespace[namespace.path, namespace.absolute?]
+
+      inner = @intern_cache[ns]
+      if inner && (cached = inner[name])
+        return cached
+      end
+
+      @intern_mutex.synchronize do
+        inner = (@intern_cache[ns] ||= {})
+        inner[name] ||= new(namespace: ns, name: name)
+      end
+    end
+
     def ==(other)
+      return true if equal?(other)
       other.is_a?(self.class) && other.namespace == namespace && other.name == name
     end
 
     alias eql? ==
 
     def hash
-      namespace.hash ^ name.hash
+      @hash ||= namespace.hash ^ name.hash
     end
 
     def to_s
@@ -53,23 +79,23 @@ def alias?
     end
 
     def absolute!
-      self.class.new(namespace: namespace.absolute!, name: name)
+      TypeName[namespace.absolute!, name]
     end
 
     def absolute?
       namespace.absolute?
     end
 
     def relative!
-      self.class.new(namespace: namespace.relative!, name: name)
+      TypeName[namespace.relative!, name]
     end
 
     def interface?
       kind == :interface
     end
 
     def with_prefix(namespace)
-      self.class.new(namespace: namespace + self.namespace, name: name)
+      TypeName[namespace + self.namespace, name]
     end
 
     def split
@@ -80,23 +106,17 @@ def +(other)
       if other.absolute?
         other
       else
-        TypeName.new(
-          namespace: self.to_namespace + other.namespace,
-          name: other.name
-        )
+        TypeName[self.to_namespace + other.namespace, other.name]
       end
     end
-    
+
     def self.parse(string)
       absolute = string.start_with?("::")
 
       *path, name = string.delete_prefix("::").split("::").map(&:to_sym)
       raise unless name
 
-      TypeName.new(
-        name: name,
-        namespace: RBS::Namespace.new(path: path, absolute: absolute)
-      )
+      TypeName[Namespace[path, absolute], name]
     end
   end
 end