Swapped from maps to tries

disco.cabal

@@ -486,6 +486,7 @@ library
                        polysemy-plugin >= 0.4 && < 0.5,
                        reflection >= 2.1.7 && < 2.2,
                        random >= 1.2.1.1 && < 1.3,
+                       list-tries,
                        constraints >= 0.13.4 && < 0.15,
                        text >= 2.0.2 && < 2.2,
                        lens >= 4.14 && < 5.4,

src/Disco/Compile.hs

@@ -199,26 +199,25 @@ compileDTerm term@(DTAbs q _ _) = do
    TyAtom a -> checkAtom a && canMemo xs
    TyCon CArr tys -> arrMemo tys && canMemo tys && canMemo xs
    TyCon c tys -> checkCon c && canMemo tys && canMemo xs
-
-  arrMemo :: [Type] -> Bool 
-  arrMemo [] = True 
-  arrMemo (x : xs) = case x of 
-   TyCon CArr _ -> False 
-   TyCon _ _ -> arrMemo xs
-   TyAtom _ -> arrMemo xs 
-
-  checkCon :: Con -> Bool 
-  checkCon (CUser _) = False 
-  checkCon CGraph = False 
-  checkCon CMap = False 
+
+  arrMemo :: [Type] -> Bool
+  arrMemo [] = True
+  arrMemo (x : xs) = case x of
+   TyCon CArr _ -> False
+   _ -> arrMemo xs
+
+  checkCon :: Con -> Bool
+  checkCon (CUser _) = False
+  checkCon CGraph = False
+  checkCon CMap = False
   checkCon (CContainer a) = checkAtom a
   checkCon _ = True
 
-  checkAtom :: Atom -> Bool 
-  checkAtom (AVar _) = False 
-  checkAtom (ABase b) = checkBase b 
-  
-  checkBase :: BaseTy -> Bool 
+  checkAtom :: Atom -> Bool
+  checkAtom (AVar _) = False
+  checkAtom (ABase b) = checkBase b
+
+  checkBase :: BaseTy -> Bool
   checkBase CtrList = False
   checkBase CtrBag = False
   checkBase CtrSet = False

src/Disco/Interpret/CESK.hs

@@ -31,6 +31,7 @@ import qualified Data.List.Infinite as InfList
 import qualified Data.Map as M
 import Data.Maybe (isJust)
 import Data.Ratio
+import qualified Data.ListTrie.Map as T
 import Disco.AST.Core
 import Disco.AST.Generic (
   Ellipsis (..),
@@ -179,7 +180,8 @@ step cesk = case cesk of
     (xs, body) <- unbind b
     case mem of
       True -> do 
-         cell <- allocateValue (VMap M.empty)
+         cell <- allocateValue (VTrie T.empty)
+         -- cell <- allocateValue (VMap M.empty)
          return $ Out (VClo [] (Just cell) e xs body) k
       False -> return $ Out (VClo [] Nothing e xs body) k
 

src/Disco/Value.hs

@@ -89,6 +89,7 @@ import Data.Char (chr, ord)
 import Data.IntMap (IntMap)
 import qualified Data.IntMap as IM
 import Data.List (foldl')
+import qualified Data.ListTrie.Map as T
 import Data.Map (Map)
 import qualified Data.Map as M
 import Data.Ratio
@@ -166,6 +167,9 @@ data Value where
   --   actually construct the set of entries, while functions only have this
   --   property when the key type is finite.
   VMap :: Map SimpleValue Value -> Value
+
+  VTrie :: T.TrieMap M.Map SimpleValue Value -> Value
+
   VGen :: StdGen -> Value
   deriving (Show)
 
@@ -312,11 +316,11 @@ pattern SMExists :: SearchMotive
 pattern SMExists = SearchMotive (True, True)
 
 -- | A collection of variables that might need to be reported for
---   a toMap, along with their types and user-legible names.
+--   a search, along with their types and user-legible names.
 newtype TestVars = TestVars [(String, Type, Name Core)]
   deriving newtype (Show, Semigroup, Monoid)
 
--- | A variable assignment found during a toMap.
+-- | A variable assignment found during a search.
 newtype TestEnv = TestEnv [(String, Type, Value)]
   deriving newtype (Show, Semigroup, Monoid)
 
@@ -341,16 +345,16 @@ interpLOp LImpl = (==>)
   True ==> False = False
   _ ==> _ = True
 
--- | The possible outcomes of a property toMap, parametrized over
+-- | The possible outcomes of a property search, parametrized over
 --   the type of values. A @TestReason@ explains why a proposition
 --   succeeded or failed.
 data TestReason_ a
   = -- | The prop evaluated to a boolean.
     TestBool
-  | -- | The toMap was an equality toMap. Records the values being
+  | -- | The search was an equality search. Records the values being
     --   compared and also their type (which is needed for printing).
     TestEqual Type a a
-  | -- | The toMap was a less than toMap. Records the values being
+  | -- | The search was a less than search. Records the values being
     --   compared and also their type (which is needed for printing).
     TestLt Type a a
   | -- | The search didn't find any examples/counterexamples.
@@ -370,16 +374,16 @@ type TestReason = TestReason_ Value
 data TestResult = TestResult Bool TestReason TestEnv
   deriving (Show)
 
--- | Whether the property toMap resulted in a runtime error.
+-- | Whether the property search resulted in a runtime error.
 testIsError :: TestResult -> Bool
 testIsError (TestResult _ (TestRuntimeError _) _) = True
 testIsError _ = False
 
--- | Whether the property toMap resulted in success.
+-- | Whether the property search resulted in success.
 testIsOk :: TestResult -> Bool
 testIsOk (TestResult b _ _) = b
 
--- | The reason the property toMap had this result.
+-- | The reason the property search had this result.
 testReason :: TestResult -> TestReason
 testReason (TestResult _ r _) = r
 
@@ -481,18 +485,19 @@ set :: Members '[State Mem] r => Int -> Cell -> Sem r ()
 set n c = modify $ \(Mem nxt m) -> Mem nxt (IM.insert n c m)
 
 memoLookup :: Members '[State Mem] r => Int -> SimpleValue -> Sem r (Maybe Value)
-memoLookup n sv = gets (M.lookup sv . toMap . IM.lookup n . mu)
+memoLookup n sv = gets (search . IM.lookup n . mu)
    where 
-      toMap (Just (Disco.Value.V (VMap vmap))) = vmap 
-      toMap _ = M.empty 
+      search (Just (Disco.Value.V (VMap vmap))) = M.lookup sv vmap 
+      search (Just (Disco.Value.V (VTrie vtrie))) = T.lookup [sv] vtrie
+      search _ = Nothing
 
 memoSet :: Members '[State Mem] r => Int -> SimpleValue -> Value -> Sem r () 
 memoSet n sv v = do 
    mc <- lkup n 
    case mc of 
-      Nothing -> undefined 
       Just (Disco.Value.V (VMap vmap)) -> set n (Disco.Value.V (VMap (M.insert sv v vmap))) 
-      Just _ -> undefined 
+      Just (Disco.Value.V (VTrie trie)) -> set n (Disco.Value.V (VTrie (T.insert [sv] v trie)))
+      _ -> undefined 
 
 ------------------------------------------------------------
 -- Pretty-printing values