From f1e80f05fe4654eb2ee0456adc09636b6ec1c386 Mon Sep 17 00:00:00 2001
From: Nikita Shilnikov <fg@flashgordon.ru>
Date: Wed, 8 Jan 2025 14:06:35 +0100
Subject: [PATCH] Use cached relations for root relation too

---
 repository/lib/rom/repository/class_interface.rb |  5 +++--
 repository/lib/rom/repository/relation_reader.rb | 10 +++++++---
 repository/lib/rom/repository/root.rb            | 13 ++++---------
 3 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/repository/lib/rom/repository/class_interface.rb b/repository/lib/rom/repository/class_interface.rb
index 866d5e390..f143fab0c 100644
--- a/repository/lib/rom/repository/class_interface.rb
+++ b/repository/lib/rom/repository/class_interface.rb
@@ -56,13 +56,14 @@ def [](name)
       #   @option options [Boolean] :auto_struct Enable/Disable auto-struct mapping
       #
       # @api public
-      def new(container = nil, **options)
+      def new(container = nil, root: Undefined, **options)
         container ||= options.fetch(:container)
 
         unless self < relation_reader
           include relation_reader.new(
             relations: container.relations.elements.keys,
-            cache: container.cache
+            cache: container.cache,
+            root: root
           )
         end
 
diff --git a/repository/lib/rom/repository/relation_reader.rb b/repository/lib/rom/repository/relation_reader.rb
index 0b800f262..7f6fefc42 100644
--- a/repository/lib/rom/repository/relation_reader.rb
+++ b/repository/lib/rom/repository/relation_reader.rb
@@ -59,17 +59,21 @@ def define_readers_for_relation(cache, name)
       end
 
       # @api private
-      def initialize(relations:, cache:)
+      def initialize(relations:, cache:, root: Undefined)
         super()
 
-        add_readers(relations, cache)
+        add_readers(relations, cache, root)
       end
 
       # @api private
-      def add_readers(relations, cache)
+      def add_readers(relations, cache, root)
         include cache.fetch_or_store(:relation_readers) {
           Readers.new(relations)
         }
+
+        unless Undefined.equal?(root)
+          define_method(:root) { |**kwargs| public_send(root, **kwargs) }
+        end
       end
     end
   end
diff --git a/repository/lib/rom/repository/root.rb b/repository/lib/rom/repository/root.rb
index 17d0a57dc..dddb3ea6d 100644
--- a/repository/lib/rom/repository/root.rb
+++ b/repository/lib/rom/repository/root.rb
@@ -43,9 +43,10 @@ class Root < Repository
       #     @return [Symbol]
       defines :root
 
-      # @!attribute [r] root
-      #   @return [Relation] The root relation
-      attr_reader :root
+      # @api public
+      def self.new(*args, **kwargs)
+        super(*args, **kwargs, root: root)
+      end
 
       # Sets descendant root relation
       #
@@ -54,12 +55,6 @@ def self.inherited(klass)
         super
         klass.root(root)
       end
-
-      # @see Repository#initialize
-      def initialize(*, **)
-        super
-        @root = prepare_relation(self.class.root)
-      end
     end
   end
 end