Added memoization of has_closure_tree_root

Author: Tom Smyth

lib/closure_tree/has_closure_tree_root.rb

@@ -15,7 +15,23 @@ def has_closure_tree_root(assoc_name, options = {})
       has_one assoc_name, -> { where(parent: nil) }, options
 
       # Fetches the association, eager loading all children and given associations
-      define_method("#{assoc_name}_including_tree") do |assoc_map = nil|
+      define_method("#{assoc_name}_including_tree") do |assoc_map_or_reload = nil, assoc_map = nil|
+        reload = false
+        if assoc_map_or_reload.is_a?(::Hash)
+          assoc_map = assoc_map_or_reload
+        else
+          reload = assoc_map_or_reload
+        end
+
+        unless reload
+          # Memoize
+          @closure_tree_roots ||= {}
+          @closure_tree_roots[assoc_name] ||= {}
+          if @closure_tree_roots[assoc_name].has_key?(assoc_map)
+            return @closure_tree_roots[assoc_name][assoc_map]
+          end
+        end
+
         roots = options[:class_name].constantize.where(parent: nil, options[:foreign_key] => id).to_a
 
         return nil if roots.empty?
@@ -65,7 +81,7 @@ def has_closure_tree_root(assoc_name, options = {})
           end
         end
 
-        root
+        @closure_tree_roots[assoc_name][assoc_map] = root
       end
     end
   end

spec/has_closure_tree_root_spec.rb

@@ -50,6 +50,22 @@
       end.to_not exceed_query_limit(4) # Without this feature, this is 15, and scales with number of nodes.
     end
 
+    it "memoizes by assoc_map" do
+      group_reloaded.root_user_including_tree.email = "x"
+      expect(group_reloaded.root_user_including_tree.email).to eq "x"
+      group_reloaded.root_user_including_tree(contracts: :contract_type).email = "y"
+      expect(group_reloaded.root_user_including_tree(contracts: :contract_type).email).to eq "y"
+      expect(group_reloaded.root_user_including_tree.email).to eq "x"
+    end
+
+    it "doesn't memoize if true argument passed" do
+      group_reloaded.root_user_including_tree.email = "x"
+      expect(group_reloaded.root_user_including_tree(true).email).to eq "1@example.com"
+      group_reloaded.root_user_including_tree(contracts: :contract_type).email = "y"
+      expect(group_reloaded.root_user_including_tree(true, contracts: :contract_type).email).
+        to eq "1@example.com"
+    end
+
     it "eager loads inverse association to group" do
       expect do
         root = group_reloaded.root_user_including_tree