Fix trie lookup, added a test to ensure memoized function doesn't get…

bower.json

@@ -24,6 +24,8 @@
     "purescript-generics-rep": "^3.0.0"
   },
   "devDependencies": {
-    "purescript-console": "^2.0.0"
+    "purescript-refs": "^2.0.0",		
+    "purescript-console": "^2.0.0",
+    "purescript-quickcheck": "^3.0.0" 
   }
 }

src/Data/Function/Memoize.purs

@@ -16,8 +16,8 @@ module Data.Function.Memoize
 import Prelude
 import Data.Char (fromCharCode, toCharCode)
 import Data.Either (Either(..))
-import Data.Generic.Rep (class Generic, Argument(..), Constructor(..),
-                         NoArguments(..), Product(..), Sum(..), from, to)
+import Data.Generic.Rep (class Generic, Argument(..), Constructor(..), NoArguments(..), Product(..), Sum(..), from, to)
+import Data.Int.Bits ((.&.), zshr)
 import Data.Lazy (Lazy, force, defer)
 import Data.List (List(..), fromFoldable, toUnfoldable)
 import Data.Maybe (Maybe(..))
@@ -117,18 +117,10 @@ instance tabulateNat :: Tabulate Int where
       tabulateImpl f = go
         where
           go :: Int -> Lazy r
-          go 0 = zer
-          go n = walk (bits (if n > 0 then n else (-n)))
-                      (if n > 0 then pos else neg)
+          go n = walk (bits n) trie
 
-          pos :: NatTrie r
-          pos = build 1
-
-          neg :: NatTrie r
-          neg = build (-1)
-
-          zer :: Lazy r
-          zer = defer \_ -> f 0
+          trie :: NatTrie r
+          trie = build 0
 
           build :: Int -> NatTrie r
           build n = NatTrie (defer \_ -> f n)
@@ -138,8 +130,8 @@ instance tabulateNat :: Tabulate Int where
           bits :: Int -> List Boolean
           bits = bits' Nil
             where
-            bits' acc 1 = acc
-            bits' acc n = bits' (Cons (mod n 2 /= 0) acc) (n / 2)
+            bits' acc 0 = acc
+            bits' acc n = bits' (Cons (n .&. 1 /= 0) acc) (n `zshr` 1)            
 
           walk :: forall a. List Boolean -> NatTrie a -> Lazy a
           walk Nil             (NatTrie a _ _) = a

test/Main.purs

@@ -1,12 +1,18 @@
 module Test.Main where
 
 import Prelude
+import Data.Generic.Rep as G
 import Control.Monad.Eff (Eff)
 import Control.Monad.Eff.Console (CONSOLE, logShow)
+import Control.Monad.Eff.Exception (EXCEPTION)
+import Control.Monad.Eff.Random (RANDOM)
+import Control.Monad.Eff.Ref (REF, newRef, modifyRef, readRef)
+import Control.Monad.Eff.Unsafe (unsafePerformEff)
 import Data.Function.Memoize (class Tabulate, memoize, memoize2, gTabulate)
-import Data.Generic.Rep (class Generic)
 import Data.List ((:), length, singleton)
 import Data.String (take, drop)
+import Test.QuickCheck (quickCheck')
+import Test.QuickCheck.Arbitrary (class Arbitrary, arbitrary)
 
 data Diff a = Add a | Remove a
 
@@ -18,12 +24,25 @@ data Ints
   = Int1 Int
   | Int2 Int
 
-derive instance genericInts :: Generic Ints _
+instance genericInts :: G.Generic Ints
+                        (G.Sum
+                         (G.Constructor "Int1" (G.Argument Int))
+                         (G.Constructor "Int2" (G.Argument Int))) where
+  to (G.Inl (G.Constructor (G.Argument x))) = Int1 x
+  to (G.Inr (G.Constructor (G.Argument x))) = Int2 x
+  from (Int1 x) = G.Inl (G.Constructor (G.Argument x))
+  from (Int2 x) = G.Inr (G.Constructor (G.Argument x))  
 
 instance tabulateInts :: Tabulate Ints where
   tabulate = gTabulate
 
-main :: forall eff. Eff (console :: CONSOLE | eff) Unit
+
+newtype SmallInt = SmallInt Int
+
+instance arbSmallInt :: Arbitrary SmallInt where
+  arbitrary = SmallInt <<< (_ `mod` 1000) <$> arbitrary
+
+main :: forall eff. Eff (ref :: REF, console :: CONSOLE, random :: RANDOM, err :: EXCEPTION | eff) Unit
 main = do
   let fibonacciFast = go 0 1
         where
@@ -62,3 +81,13 @@ main = do
             | otherwise -> smallest (Add (take 1 s2)    : diff s1 (drop 1 s2))
                                     (Remove (take 1 s1) : diff (drop 1 s1) s2)
   logShow $ diff "Hello, PureScript" "ello, PureScript!"
+
+  called <- newRef 0
+  let fn x = 2 * x
+      msin = memoize \n -> unsafePerformEff do
+        modifyRef called (_ + 1) 
+        pure $ fn n
+  quickCheck' 10000 $ \(SmallInt x) -> fn x == msin x
+  ncalled <- readRef called
+  quickCheck' 1 $ ncalled < 2000
+  pure unit