Improve the Memoizable module

spec/lucky/memoize_spec.cr

@@ -6,15 +6,21 @@ private class ObjectWithMemoizedMethods
   include Lucky::Memoizable
   getter times_method_1_called = 0
   getter times_method_2_called = 0
+  getter times_method_3_called = 0
 
   memoize def method_1 : String
     @times_method_1_called += 1
     "method_1"
   end
 
-  memoize def method_2 : Int32
+  memoize def method_2 : Int32?
     @times_method_2_called += 1
-    5
+    nil
+  end
+
+  memoize def method_3(arg_a : String, arg_b : String = "default-arg-b") : String
+    @times_method_3_called += 1
+    arg_a + ", " + arg_b
   end
 end
 
@@ -23,11 +29,55 @@ describe "memoizations" do
     object = ObjectWithMemoizedMethods.new
 
     object.method_1.should eq "method_1"
-    2.times { object.method_1 }
+    2.times { object.method_1.should eq("method_1") }
     object.times_method_1_called.should eq 1
+  end
 
-    object.method_2.should eq 5
-    9.times { object.method_2 }
+  it "can cache a nil result" do
+    object = ObjectWithMemoizedMethods.new
+
+    object.method_2.should be_nil
+    2.times { object.method_2.should be_nil }
     object.times_method_2_called.should eq 1
   end
+
+  it "caches based on argument equality" do
+    object = ObjectWithMemoizedMethods.new
+
+    object.method_3("arg-a", "arg-b").should eq("arg-a, arg-b")
+    2.times { object.method_3("arg-a", "arg-b").should eq("arg-a, arg-b") }
+    object.times_method_3_called.should eq 1
+
+    object.method_3("arg-a", "arg-c").should eq("arg-a, arg-c")
+    2.times { object.method_3("arg-a", "arg-c").should eq("arg-a, arg-c") }
+    object.times_method_3_called.should eq 2
+  end
+
+  it "handles default arguments" do
+    object = ObjectWithMemoizedMethods.new
+
+    object.method_3("arg-a", "default-arg-b").should eq("arg-a, default-arg-b")
+    object.method_3("arg-a", "default-arg-b").should eq("arg-a, default-arg-b")
+    object.method_3("arg-a").should eq("arg-a, default-arg-b")
+    object.times_method_3_called.should eq 1
+  end
+
+  it "handles calling with named arguments" do
+    object = ObjectWithMemoizedMethods.new
+
+    object.method_3("arg-a", "arg-b").should eq("arg-a, arg-b")
+    object.method_3("arg-a", arg_b: "arg-b").should eq("arg-a, arg-b")
+    object.method_3(arg_a: "arg-a", arg_b: "arg-b").should eq("arg-a, arg-b")
+    object.method_3(arg_b: "arg-b", arg_a: "arg-a").should eq("arg-a, arg-b")
+    object.times_method_3_called.should eq 1
+  end
+
+  it "does not hold on to result of previous calls" do
+    object = ObjectWithMemoizedMethods.new
+
+    object.method_3("arg-a", "arg-b").should eq("arg-a, arg-b")
+    object.method_3("arg-a", "arg-c").should eq("arg-a, arg-c")
+    object.method_3("arg-a", "arg-b").should eq("arg-a, arg-b")
+    object.times_method_3_called.should eq 3
+  end
 end

src/lucky/memoizable.cr

@@ -15,25 +15,73 @@ module Lucky::Memoizable
   # then each subsequent call returns the user record.
   #
   # The `memoize` method will raise a compile time exception if you forget to include
-  # a return type for your method, or if your return type is a `Union`.
-  # Arguments are not allowed in memoized methods because these can change the return value.
+  # a return type for your method, or if any arguments are missing a type.
+  # The result of a set of arguments is only kept until the passed arguments change.
+  # Once they change, passing previous arguments will re-run the memoized method.
+  # Equality (==) is used for checking on argument updates.
   macro memoize(method_def)
-    {% raise "Return type of memoize method must not be a Union" if method_def.return_type.is_a?(Union) %}
-    {% raise "You must define a return type for memoize methods" if method_def.return_type.is_a?(Nop) %}
-    {% raise "Memoize methods can not be defined with arguments" if method_def.args.size > 0 %}
+    {% raise "You must define a return type for memoized methods" if method_def.return_type.is_a?(Nop) %}
+    {% raise "All arguments must have an explicit type for memoized methods" if method_def.args.any? &.is_a?(Nop) %}
 
-    @__{{ method_def.name }} : {{ method_def.return_type }}? = nil
+    @__memoized_{{method_def.name}} : Tuple(
+      {{ method_def.return_type }},
+      {% for arg in method_def.args %}
+        {{ arg.restriction }},
+      {% end %}
+    )?
 
     # Returns uncached value
-    def {{ method_def.name }}__uncached : {{ method_def.return_type }}
+    def {{ method_def.name }}__uncached(
+      {% for arg in method_def.args %}
+        {{ arg.name }} : {{ arg.restriction }},
+      {% end %}
+    ) : {{ method_def.return_type }}
       {{ method_def.body }}
     end
 
-    # Returns cached value
-    def {{ method_def.name }} : {{ method_def.return_type }}
-      @__{{ method_def.name }} ||= -> do
-        {{ method_def.name }}__uncached
+    # Checks the passed arguments against the memoized args
+    # and runs the method body if it is the very first call
+    # or the arguments do not match
+    def {{ method_def.name }}__tuple_cached(
+      {% for arg in method_def.args %}
+        {{ arg.name }} : {{ arg.restriction }},
+      {% end %}
+    ) : Tuple(
+      {{ method_def.return_type }},
+      {% for arg in method_def.args %}
+        {{ arg.restriction }},
+      {% end %}
+    )
+      {% for arg, index in method_def.args %}
+        @__memoized_{{ method_def.name }} = nil if {{arg.name}} != @__memoized_{{ method_def.name }}.try &.at({{index}} + 1)
+      {% end %}
+      @__memoized_{{ method_def.name }} ||= -> do
+        result = {{ method_def.name }}__uncached(
+          {% for arg in method_def.args %}
+            {{arg.name}},
+          {% end %}
+        )
+        {
+          result,
+          {% for arg in method_def.args %}
+            {{arg.name}},
+          {% end %}
+        }
       end.call.not_nil!
     end
+
+    # Returns cached value
+    def {{ method_def.name }}(
+      {% for arg in method_def.args %}
+        {% has_default = arg.default_value || arg.default_value == false || arg.default_value == nil %}
+        {{ arg.name }} : {{ arg.restriction }}{% if has_default %} = {{ arg.default_value }}{% end %},
+      {% end %}
+    ) : {{ method_def.return_type }}
+      {{ method_def.name }}__tuple_cached(
+        {% for arg in method_def.args %}
+          {{arg.name}},
+        {% end %}
+      ).first
+    end
   end
 end