Skip to content

Commit

Permalink
Restyle Memoize (#395)
Browse files Browse the repository at this point in the history
* Restyled by fourmolu

* Restyled by hlint

---------

Co-authored-by: Restyled.io <[email protected]>
  • Loading branch information
restyled-io[bot] and restyled-commits authored Jul 13, 2024
1 parent c2fd969 commit 3f10e6a
Show file tree
Hide file tree
Showing 3 changed files with 41 additions and 50 deletions.
34 changes: 17 additions & 17 deletions src/Disco/Compile.hs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down
27 changes: 9 additions & 18 deletions src/Disco/Interpret/CESK.hs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.
Expand Down
30 changes: 15 additions & 15 deletions src/Disco/Value.hs
Original file line number Diff line number Diff line change
Expand Up @@ -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@).
Expand Down Expand Up @@ -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.
Expand All @@ -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
Expand Down

0 comments on commit 3f10e6a

Please sign in to comment.