memoize: fix a few more zero input cases

mrjbq7 · Apr 5, 2023 · a937766 · a937766
1 parent d362b56
commit a937766
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 9 deletions.
diff --git a/core/memoize/memoize-tests.factor b/core/memoize/memoize-tests.factor
@@ -59,3 +59,18 @@ MEMO: bar ( -- x ) bar-counter counter ;
     bar
     bar
 ] unit-test
+
+SYMBOL: baz-counter
+0 baz-counter set-global
+
+MEMO: baz ( -- x ) baz-counter counter drop f ;
+
+{ 0 f 1 f 1 f 1 } [
+    baz-counter get-global
+    baz
+    baz-counter get-global
+    baz
+    baz-counter get-global
+    baz
+    baz-counter get-global
+] unit-test
diff --git a/core/memoize/memoize.factor b/core/memoize/memoize.factor
@@ -23,7 +23,7 @@ IN: memoize
 
 : packer ( seq -- quot )
     length dup 4 <=
-    [ { [ t ] [ ] [ 2array ] [ 3array ] [ 4array ] } nth ]
+    [ { [ f ] [ ] [ 2array ] [ 3array ] [ 4array ] } nth ]
     [ { } [nsequence] ] if ;
 
 : unpacker ( seq -- quot )
@@ -44,9 +44,9 @@ IN: memoize
 : make/0 ( table quot effect -- quot )
     out>> [
         packer '[
-            _
-            [ first-unsafe ]
-            [ @ @ [ 0 rot set-nth-unsafe ] keep ] ?unless
+            _ dup first-unsafe dup null eq? [
+                drop @ @ [ 0 rot set-nth-unsafe ] keep
+            ] [ nip ] if
         ]
     ] keep unpacker compose ;
 
@@ -62,11 +62,11 @@ PRIVATE>
     3tri ;
 
 : define-memoized ( word quot effect -- )
-    dup in>> length zero? [ f 1array ] [ H{ } clone ] if
+    dup in>> length zero? [ null 1array ] [ H{ } clone ] if
     (define-memoized) ;
 
 : define-identity-memoized ( word quot effect -- )
-    dup in>> length zero? [ f 1array ] [ IH{ } clone ] if
+    dup in>> length zero? [ null 1array ] [ IH{ } clone ] if
     (define-memoized) ;
 
 PREDICATE: memoized < word "memoize" word-prop >boolean ;
@@ -83,18 +83,18 @@ M: memoized reset-word
     bi ;
 
 : memoize-quot ( quot effect -- memo-quot )
-    dup in>> length zero? [ f 1array ] [ H{ } clone ] if
+    dup in>> length zero? [ null 1array ] [ H{ } clone ] if
     -rot make-memoizer ;
 
 : reset-memoized ( word -- )
     "memoize" word-prop dup sequence?
-    [ f swap set-first ] [ clear-assoc ] if ;
+    [ null swap set-first ] [ clear-assoc ] if ;
 
 : invalidate-memoized ( inputs... word -- )
     [ stack-effect in>> packer call ]
     [
         "memoize" word-prop dup sequence?
-        [ f swap set-first ] [ delete-at ] if
+        [ null swap set-first ] [ delete-at ] if
     ]
     bi ;