diff options
Diffstat (limited to 'src/Data')
| -rw-r--r-- | src/Data/Expr/SharingRecovery.hs | 189 |
1 files changed, 145 insertions, 44 deletions
diff --git a/src/Data/Expr/SharingRecovery.hs b/src/Data/Expr/SharingRecovery.hs index e386f4e..cdb64eb 100644 --- a/src/Data/Expr/SharingRecovery.hs +++ b/src/Data/Expr/SharingRecovery.hs @@ -1,17 +1,27 @@ +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE DataKinds #-} +{-# LANGUAGE DefaultSignatures #-} {-# LANGUAGE DeriveFunctor #-} {-# LANGUAGE DerivingVia #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE QuantifiedConstraints #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} {-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} module Data.Expr.SharingRecovery where import Control.Applicative ((<|>)) import Control.Monad.Trans.State.Strict import Data.Bifunctor (second) -import Data.GADT.Compare +import Data.Char (chr, ord) +import Data.Functor.Const +import Data.Functor.Identity +import Data.Functor.Product import Data.Hashable import Data.HashMap.Strict (HashMap) import qualified Data.HashMap.Strict as HM @@ -24,8 +34,6 @@ import Unsafe.Coerce (unsafeCoerce) import Data.StableName.Extra --- TODO: This is not yet done, see the bottom of this file - -- TODO: This implementation needs extensive documentation. 1. It is written -- quite generically, meaning that the actual algorithm is easily obscured to -- all but the most willing readers; 2. the original paper leaves something to @@ -36,24 +44,28 @@ import Data.StableName.Extra class Functor1 f where fmap1 :: (forall b. g b -> h b) -> f g a -> f h a + default fmap1 :: Traversable1 f => (forall b. g b -> h b) -> f g a -> f h a + fmap1 f x = runIdentity (traverse1 (Identity . f) x) + class Functor1 f => Traversable1 f where traverse1 :: Applicative m => (forall b. g b -> m (h b)) -> f g a -> m (f h a) -- | Expression in parametric higher-order abstract syntax form data PHOASExpr typ v f t where PHOASOp :: typ t -> f (PHOASExpr typ v f) t -> PHOASExpr typ v f t - PHOASLam :: typ (a -> b) -> typ a -> (PHOASExpr typ v f a -> PHOASExpr typ v f b) -> PHOASExpr typ v f (a -> b) + PHOASLam :: typ (a -> b) -> typ a -> (v a -> PHOASExpr typ v f b) -> PHOASExpr typ v f (a -> b) PHOASVar :: typ t -> v t -> PHOASExpr typ v f t newtype Tag t = Tag Natural deriving (Show, Eq) + deriving (Hashable) via Natural newtype NameFor typ f t = NameFor (StableName (PHOASExpr typ Tag f t)) deriving (Eq) deriving (Hashable) via (StableName (PHOASExpr typ Tag f t)) -instance GEq (NameFor typ f) where - geq (NameFor n1) (NameFor n2) +instance TestEquality (NameFor typ f) where + testEquality (NameFor n1) (NameFor n2) | eqStableName n1 n2 = Just unsafeCoerceRefl | otherwise = Nothing where @@ -106,7 +118,7 @@ pruneExpr' = \case then pure $ PStub (NameFor name) tyf else do tag <- state (\(i, mp) -> (Tag i, (i + 1, mp))) - let body = f (PHOASVar tyarg tag) + let body = f tag PLam (NameFor name) tyf tyarg tag <$> pruneExpr' body PHOASVar ty tag -> pure $ PVar ty tag @@ -181,56 +193,145 @@ liftExpr' totals term = | otherwise -> error "Term does not exist, yet we have it in hand" --- | Errors on stubs. -lexprTypeOf :: LExpr typ f t -> typ t -lexprTypeOf (LExpr _ e) = case e of - LStub{} -> error "lexprTypeOf: got a stub" - LOp _ t _ -> t - LLam _ t _ _ _ -> t - LVar t _ -> t - - --- TODO: lower LExpr into a normal expression with let bindings. Every LStub --- should correspond to some let-bound expression higher up in the tree (if it --- does not, that's a bug), and should become a De Bruijn variable reference to --- said let-bound expression. Lambdas should also get proper De Bruijn indices --- instead of tags, and LVar is also a normal variable (referring to a --- lambda-abstracted argument). - -- | Untyped De Bruijn expression. No more names: there are lets now, and -- variable references are De Bruijn indices. These indices are not type-safe -- yet, though. data UBExpr typ f t where UBOp :: typ t -> f (UBExpr typ f) t -> UBExpr typ f t - UBLam :: typ a -> UBExpr typ f b -> UBExpr typ f (a -> b) + UBLam :: typ (a -> b) -> typ a -> UBExpr typ f b -> UBExpr typ f (a -> b) UBLet :: typ a -> UBExpr typ f a -> UBExpr typ f b -> UBExpr typ f b - -- De Bruijn index + -- | De Bruijn index UBVar :: typ t -> Int -> UBExpr typ f t -lowerExpr :: LExpr typ f t -> UBExpr typ f t -lowerExpr = lowerExpr' mempty 0 +lowerExpr :: Functor1 f => LExpr typ f t -> UBExpr typ f t +lowerExpr = lowerExpr' mempty mempty 0 --- 1. name |-> De Bruijn level of the variable defining that name --- 2. Number of variables already in scope -lowerExpr' :: forall typ f t. Traversable1 f => HashMap (SomeNameFor typ f) Int -> Int -> LExpr typ f t -> UBExpr typ f t -lowerExpr' namelvl curlvl (LExpr lifted ex) = - let prefix = buildPrefix curlvl lifted +data SomeTag = forall t. SomeTag (Tag t) + +instance Eq SomeTag where + SomeTag (Tag n) == SomeTag (Tag m) = n == m + +instance Hashable SomeTag where + hashWithSalt salt (SomeTag tag) = hashWithSalt salt tag + +lowerExpr' :: forall typ f t. Functor1 f + => HashMap (SomeNameFor typ f) Int -- ^ node |-> De Bruijn level of defining binding + -> HashMap SomeTag Int -- ^ tag |-> De Bruijn level of defining binding + -> Int -- ^ Number of variables already in scope + -> LExpr typ f t -> UBExpr typ f t +lowerExpr' namelvl taglvl curlvl (LExpr lifted ex) = + let (namelvl', prefix) = buildPrefix namelvl curlvl lifted curlvl' = curlvl + length lifted - in case ex of - LStub name ty -> - case HM.lookup (SomeNameFor name) namelvl of - Just lvl -> UBVar ty (curlvl - lvl - 1) - Nothing -> error "Variable out of scope" - LOp name ty args -> - UBOp ty (_ $ traverse1 _ args) + in prefix $ + case ex of + LStub name ty -> + case HM.lookup (SomeNameFor name) namelvl' of + Just lvl -> UBVar ty (curlvl - lvl - 1) + Nothing -> error "Name variable out of scope" + LOp _ ty args -> + UBOp ty (fmap1 (lowerExpr' namelvl' taglvl curlvl') args) + LLam _ tyf tyarg tag body -> + UBLam tyf tyarg (lowerExpr' namelvl' (HM.insert (SomeTag tag) curlvl' taglvl) (curlvl' + 1) body) + LVar ty tag -> + case HM.lookup (SomeTag tag) taglvl of + Just lvl -> UBVar ty (curlvl - lvl - 1) + Nothing -> error "Tag variable out of scope" where - buildPrefix :: forall b. Int -> [Some (LExpr typ f)] -> UBExpr typ f b -> UBExpr typ f b - buildPrefix _ [] = id - buildPrefix lvl (Some rhs : rhss) = - UBLet (lexprTypeOf rhs) (lowerExpr' namelvl lvl rhs) - . buildPrefix (lvl + 1) rhss + buildPrefix :: forall b. + HashMap (SomeNameFor typ f) Int + -> Int + -> [Some (LExpr typ f)] + -> (HashMap (SomeNameFor typ f) Int, UBExpr typ f b -> UBExpr typ f b) + buildPrefix namelvl' _ [] = (namelvl', id) + buildPrefix namelvl' lvl (Some rhs@(LExpr _ rhs') : rhss) = + let name = case rhs' of + LStub n _ -> n + LOp n _ _ -> n + LLam n _ _ _ _ -> n + LVar _ _ -> error "Recovering sharing of a tag is useless" + ty = case rhs' of + LStub{} -> error "Recovering sharing of a stub is useless" + LOp _ t _ -> t + LLam _ t _ _ _ -> t + LVar t _ -> t + prefix = UBLet ty (lowerExpr' namelvl' taglvl lvl rhs) + in (prefix .) <$> buildPrefix (HM.insert (SomeNameFor name) lvl namelvl') (lvl + 1) rhss +-- | A typed De Bruijn index. data Idx env t where IZ :: Idx (t : env) t IS :: Idx env t -> Idx (s : env) t +deriving instance Show (Idx env t) + +data Env env f where + ETop :: Env '[] f + EPush :: Env env f -> f t -> Env (t : env) f + +envLookup :: Idx env t -> Env env f -> f t +envLookup IZ (EPush _ x) = x +envLookup (IS i) (EPush e _) = envLookup i e + +-- | Untyped lookup in an 'Env'. +envLookupU :: Int -> Env env f -> Maybe (Some (Product f (Idx env))) +envLookupU = go id + where + go :: (forall a. Idx env a -> Idx env' a) -> Int -> Env env f -> Maybe (Some (Product f (Idx env'))) + go !_ !_ ETop = Nothing + go f 0 (EPush _ t) = Just (Some (Pair t (f IZ))) + go f i (EPush e _) = go (f . IS) (i - 1) e + +-- | Typed De Bruijn expression. This is the final result of sharing recovery. +data BExpr typ env f t where + BOp :: typ t -> f (BExpr typ env f) t -> BExpr typ env f t + BLam :: typ (a -> b) -> typ a -> BExpr typ (a : env) f b -> BExpr typ env f (a -> b) + BLet :: typ a -> BExpr typ env f a -> BExpr typ (a : env) f b -> BExpr typ env f b + BVar :: typ t -> Idx env t -> BExpr typ env f t +deriving instance (forall a. Show (typ a), forall a r. (forall b. Show (r b)) => Show (f r a)) + => Show (BExpr typ env f t) + +prettyBExpr :: (forall m env' a. Monad m => (forall b. Int -> BExpr typ env' f b -> m ShowS) + -> Int -> f (BExpr typ env' f) a -> m ShowS) + -> BExpr typ '[] f t -> String +prettyBExpr prettyOp e = evalState (prettyBExpr' prettyOp ETop 0 e) 0 "" + +prettyBExpr' :: (forall m env' a. Monad m => (forall b. Int -> BExpr typ env' f b -> m ShowS) + -> Int -> f (BExpr typ env' f) a -> m ShowS) + -> Env env (Const String) -> Int -> BExpr typ env f t -> State Int ShowS +prettyBExpr' prettyOp env d = \case + BOp _ args -> + prettyOp (prettyBExpr' prettyOp env) d args + BLam _ _ body -> do + name <- genName + body' <- prettyBExpr' prettyOp (EPush env (Const name)) 0 body + return $ showParen (d > 0) $ showString ("λ" ++ name ++ ". ") . body' + BLet _ rhs body -> do + name <- genName + rhs' <- prettyBExpr' prettyOp env 0 rhs + body' <- prettyBExpr' prettyOp (EPush env (Const name)) 0 body + return $ showParen (d > 0) $ showString ("let " ++ name ++ " = ") . rhs' . showString " in " . body' + BVar _ idx -> + return $ showString (getConst (envLookup idx env)) + where + genName = do + i <- state (\i -> (i, i + 1)) + return $ if i < 26 then [chr (ord 'a' + i)] else 'x' : show i + +retypeExpr :: (Functor1 f, TestEquality typ) => UBExpr typ f t -> BExpr typ '[] f t +retypeExpr = retypeExpr' ETop + +retypeExpr' :: (Functor1 f, TestEquality typ) => Env env typ -> UBExpr typ f t -> BExpr typ env f t +retypeExpr' env (UBOp ty args) = BOp ty (fmap1 (retypeExpr' env) args) +retypeExpr' env (UBLam tyf tyarg body) = BLam tyf tyarg (retypeExpr' (EPush env tyarg) body) +retypeExpr' env (UBLet ty rhs body) = BLet ty (retypeExpr' env rhs) (retypeExpr' (EPush env ty) body) +retypeExpr' env (UBVar ty idx) = + case envLookupU idx env of + Just (Some (Pair defty tidx)) -> + case testEquality ty defty of + Just Refl -> BVar ty tidx + Nothing -> error "Type mismatch in untyped De Bruijn expression" + Nothing -> error "Untyped De Bruijn index out of range" + + +sharingRecovery :: (Traversable1 f, TestEquality typ) => (forall v. PHOASExpr typ v f t) -> BExpr typ '[] f t +sharingRecovery e = retypeExpr $ lowerExpr $ uncurry liftExpr $ pruneExpr e |