Restyle Memoize

src/Disco/Compile.hs

@@ -186,42 +186,41 @@ compileDTerm term@(DTAbs q _ _) = do
 
   quantify :: Op -> Core -> Core
   quantify op = CApp (CConst op)
-  
+
   -- Given a function's arguments, determine if it is memoizable.
   -- A function is memoizable if its arguments can be converted into
   -- a simple value.
   canMemo :: [Type] -> ShouldMemo
   canMemo tys
-   | all canMemoTy tys = Memo
-   | otherwise = NoMemo
+    | all canMemoTy tys = Memo
+    | otherwise = NoMemo
 
   canMemoTy :: Type -> Bool
   canMemoTy (TyAtom a) = canMemoAtom a
   -- Anti-higher order while allowing for multiple arrows.
-  canMemoTy (TyCon CArr tys@(t:_)) = case t of
-   TyCon CArr _ -> False
-   _ -> all canMemoTy tys
+  canMemoTy (TyCon CArr tys@(t : _)) = case t of
+    TyCon CArr _ -> False
+    _ -> all canMemoTy tys
   canMemoTy (TyCon c tys) = canMemoCon c && all canMemoTy tys
 
   canMemoCon :: Con -> Bool
   canMemoCon = \case
-   CArr -> False
-   CUser _ -> False
-   CGraph -> False
-   CMap -> False
-   CContainer a -> canMemoAtom a
-   _ -> True
+    CArr -> False
+    CUser _ -> False
+    CGraph -> False
+    CMap -> False
+    CContainer a -> canMemoAtom a
+    _ -> True
 
   canMemoAtom :: Atom -> Bool
   canMemoAtom (AVar _) = False
   canMemoAtom (ABase b) = canMemoBase b
 
   canMemoBase :: BaseTy -> Bool
   canMemoBase = \case
-   Gen -> False
-   P -> False
-   _ -> True
-
+    Gen -> False
+    P -> False
+    _ -> True
 
 -- Special case for Cons, which compiles to a constructor application
 -- rather than a function application.
@@ -433,7 +432,8 @@ compileMatch (DPPair _ x1 x2) s _ e = do
   -- {? e when s is (x1,x2) ?}   ==>   (\y. (\x1.\x2. e) (fst y) (snd y)) s
   return $
     CApp
-      ( CAbs NoMemo
+      ( CAbs
+          NoMemo
           ( bind
               [y]
               ( CApp

src/Disco/Interpret/CESK.hs

@@ -175,15 +175,13 @@ step cesk = case cesk of
   (In CUnit _ k) -> return $ Out VUnit k
   (In (CPair c1 c2) e k) -> return $ In c1 e (FPairR e c2 : k)
   (In (CProj s c) e k) -> return $ In c e (FProj s : k)
-
   (In (CAbs mem b) e k) -> do
     (xs, body) <- unbind b
     case mem of
-      Memo -> do 
-         cell <- allocateValue (VMap M.empty)
-         return $ Out (VClo (Just (cell,[])) e xs body) k
+      Memo -> do
+        cell <- allocateValue (VMap M.empty)
+        return $ Out (VClo (Just (cell, [])) e xs body) k
       NoMemo -> return $ Out (VClo Nothing e xs body) k
-
   (In (CApp c1 c2) e k) -> return $ In c1 e (FArg e c2 : k)
   (In (CType ty) _ k) -> return $ Out (VType ty) k
   (In (CDelay b) e k) -> do
@@ -203,23 +201,16 @@ step cesk = case cesk of
   (Out v2 (FPairL v1 : k)) -> return $ Out (VPair v1 v2) k
   (Out (VPair v1 v2) (FProj s : k)) -> return $ Out (selectSide s v1 v2) k
   (Out v (FArg e c2 : k)) -> return $ In c2 e (FApp v : k)
-
-
-
-  (Out v (FMemo n sv : k)) -> memoSet n sv v *> (return $ Out v k)
-
+  (Out v (FMemo n sv : k)) -> memoSet n sv v Data.Functor.$> Out v k
   (Out v (FApp (VClo mi e [x] b) : k)) -> case mi of
-   Nothing -> return $ In b (Ctx.insert (localName x) v e) k
-   Just (n,mem) -> do
+    Nothing -> return $ In b (Ctx.insert (localName x) v e) k
+    Just (n, mem) -> do
       let sv = toSimpleValue $ foldr VPair VUnit (v : mem)
       mv <- memoLookup n sv
       case mv of
-         Nothing -> return $ In b (Ctx.insert (localName x) v e) (FMemo n sv : k)
-         Just v' -> return $ Out v' k
-
-  (Out v (FApp (VClo mi e (x : xs) b) : k)) -> return $ Out (VClo (second (v:) <$> mi) (Ctx.insert (localName x) v e) xs b) k
-
-
+        Nothing -> return $ In b (Ctx.insert (localName x) v e) (FMemo n sv : k)
+        Just v' -> return $ Out v' k
+  (Out v (FApp (VClo mi e (x : xs) b) : k)) -> return $ Out (VClo (second (v :) <$> mi) (Ctx.insert (localName x) v e) xs b) k
   (Out v2 (FApp (VConst op) : k)) -> appConst k op v2
   (Out v2 (FApp (VFun f) : k)) -> return $ Out (f v2) k
   -- Annoying to repeat this code, not sure of a better way.

src/Disco/Value.hs

@@ -133,7 +133,7 @@ data Value where
   VPair :: Value -> Value -> Value
   -- | A closure, i.e. a function body together with its
   --   environment.
-  VClo :: Maybe (Int,[Value]) -> Env -> [Name Core] -> Core -> Value
+  VClo :: Maybe (Int, [Value]) -> Env -> [Name Core] -> Core -> Value
   -- | A disco type can be a value.  For now, there are only a very
   --   limited number of places this could ever show up (in
   --   particular, as an argument to @enumerate@ or @count@).
@@ -455,10 +455,10 @@ allocate e t = do
   return n
 
 allocateValue :: Members '[State Mem] r => Value -> Sem r Int
-allocateValue v = do 
-   Mem n m <- get 
-   put $ Mem (n + 1) (IM.insert n (Disco.Value.V v) m)
-   return n
+allocateValue v = do
+  Mem n m <- get
+  put $ Mem (n + 1) (IM.insert n (Disco.Value.V v) m)
+  return n
 
 -- | Allocate new memory cells for a group of mutually recursive
 --   bindings, and return the indices of the allocate cells.
@@ -482,16 +482,16 @@ 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 (mLookup . IM.lookup n . mu)
-   where 
-      mLookup (Just (Disco.Value.V (VMap vmap))) = M.lookup sv vmap 
-      mLookup _ = Nothing
-
-memoSet :: Members '[State Mem] r => Int -> SimpleValue -> Value -> Sem r () 
-memoSet n sv v = do 
-   mc <- lkup n 
-   case mc of 
-      Just (Disco.Value.V (VMap vmap)) -> set n (Disco.Value.V (VMap (M.insert sv v vmap))) 
-      _ -> return ()
+ where
+  mLookup (Just (Disco.Value.V (VMap vmap))) = M.lookup sv vmap
+  mLookup _ = Nothing
+
+memoSet :: Members '[State Mem] r => Int -> SimpleValue -> Value -> Sem r ()
+memoSet n sv v = do
+  mc <- lkup n
+  case mc of
+    Just (Disco.Value.V (VMap vmap)) -> set n (Disco.Value.V (VMap (M.insert sv v vmap)))
+    _ -> return ()
 
 ------------------------------------------------------------
 -- Pretty-printing values