aboutsummaryrefslogtreecommitdiff
path: root/src/Data
diff options
context:
space:
mode:
authorTom Smeding <tom@tomsmeding.com>2024-08-28 12:01:08 +0200
committerTom Smeding <tom@tomsmeding.com>2024-08-28 12:01:08 +0200
commit869be329dd05eede1dd1adb3c3b6ce2340074818 (patch)
tree492e36267ffa55294525e7c83e1e3adbd7e58892 /src/Data
parenta6f925deaa2044a0fe18a74fc52703e00f111056 (diff)
It does something
Diffstat (limited to 'src/Data')
-rw-r--r--src/Data/Expr/SharingRecovery.hs193
1 files changed, 147 insertions, 46 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
- 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)
- 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
+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 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.
+ 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