Support both identity and equality comparison in Loader#for (#53)

Author: Copilot

CHANGELOG.md

@@ -1,3 +1,7 @@
+## [Unreleased]
+
+- Support both identity and equality comparison in `Loader#for`. Identity lookup (via `object_id`) is tried first, with equality lookup as a fallback. Thanks [Alfonso Uceda](https://github.com/AlfonsoUceda) for reporting the issue!
+
 ## [2.1.0] - 2026/03/07
 
 - Add `n1_bind_to` to support context sharing for plain Ruby objects without ActiveRecord, enabling N+1-free batch loading. Nested loading through N1Loader automatically propagates the shared context. Thanks [Pawel Pacana](https://github.com/paneq) for the feature request!

lib/n1_loader/ar_lazy_preload/loader_patch.rb

@@ -7,16 +7,16 @@ module LoaderPatch
       attr_accessor :context_setup
 
       def loaded
-        return @loaded if @already_loaded && @already_context
+        return @loaded_by_identity if @already_loaded && @already_context
 
         super
 
         synchronize do
-          context_setup&.call(@loaded.values.flatten) unless @already_context
+          context_setup&.call(@loaded_by_identity.values.flatten) unless @already_context
         end
 
         @already_context = true
-        @loaded
+        @loaded_by_identity
       end
     end
   end

lib/n1_loader/core/loader.rb

@@ -47,12 +47,16 @@ def initialize(elements, **args)
     end
 
     def for(element)
-      if loaded.empty? && elements.any?
+      identity_loaded = loaded
+
+      if identity_loaded.empty? && elements.any?
         raise NotFilled, "Nothing was preloaded, perhaps you forgot to use fulfill method"
       end
-      raise NotLoaded, "The data was not preloaded for the given element" unless loaded.key?(element)
 
-      loaded[element]
+      return identity_loaded[element] if identity_loaded.key?(element)
+      return @loaded_by_value[element] if @loaded_by_value.key?(element)
+
+      raise NotLoaded, "The data was not preloaded for the given element"
     end
 
     def cache_key
@@ -103,25 +107,28 @@ def perform(_elements)
     end
 
     def fulfill(element, value)
-      @loaded[element] = value
+      @loaded_by_identity[element] = value
+      @loaded_by_value[element] = value
     end
 
-    def loaded
-      return @loaded if @already_loaded
+    def ensure_loaded
+      return if @already_loaded
 
-      synchronize do
-        non_thread_safe_loaded unless @already_loaded
-      end
+      synchronize { non_thread_safe_loaded unless @already_loaded }
+    end
 
-      @loaded
+    def loaded
+      ensure_loaded
+      @loaded_by_identity
     end
 
     def non_thread_safe_loaded # rubocop:disable Metrics/AbcSize, Metrics/MethodLength
-      return @loaded if @already_loaded
+      return if @already_loaded
 
       check_arguments!
 
-      @loaded = {}
+      @loaded_by_identity = {}.compare_by_identity
+      @loaded_by_value = {}
 
       if respond_to?(:single) && elements.size == 1
         fulfill(elements.first, single(elements.first))
@@ -131,7 +138,6 @@ def non_thread_safe_loaded # rubocop:disable Metrics/AbcSize, Metrics/MethodLeng
       end
 
       @already_loaded = true
-      @loaded
     end
   end
 end

spec/n1_loader_spec.rb

@@ -182,13 +182,33 @@ def perform(elements)
   end
 
   describe "loaded comparison" do
-    it "compares by value" do
+    it "compares by identity first" do
       instance = loader.new(objects)
 
+      expect(objects.first).to equal(objects.first)
       expect(instance.for(objects.first)).to eq([objects.first])
 
       expect { instance.for(object) }.to raise_error(N1Loader::NotLoaded)
     end
+
+    it "falls back to equality comparison when no identity match" do
+      equal_klass = Struct.new(:id)
+
+      original = equal_klass.new(1)
+      equal_copy = equal_klass.new(1)
+
+      custom_loader = Class.new(N1Loader::Loader) do
+        def perform(elements)
+          elements.each { |element| fulfill(element, [element]) }
+        end
+      end
+
+      instance = custom_loader.new([original])
+
+      expect(original).not_to equal(equal_copy)
+      expect(instance.for(original)).to eq([original])
+      expect(instance.for(equal_copy)).to eq([original])
+    end
   end
 
   context "when fulfill was not used" do