diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/Data/Expr/SharingRecovery.hs | 172 | ||||
| -rw-r--r-- | src/Data/StableName/Extra.hs | 27 |
2 files changed, 157 insertions, 42 deletions
diff --git a/src/Data/Expr/SharingRecovery.hs b/src/Data/Expr/SharingRecovery.hs index cdb64eb..f9d27e6 100644 --- a/src/Data/Expr/SharingRecovery.hs +++ b/src/Data/Expr/SharingRecovery.hs @@ -17,7 +17,7 @@ module Data.Expr.SharingRecovery where import Control.Applicative ((<|>)) import Control.Monad.Trans.State.Strict -import Data.Bifunctor (second) +import Data.Bifunctor (first, second) import Data.Char (chr, ord) import Data.Functor.Const import Data.Functor.Identity @@ -25,12 +25,17 @@ import Data.Functor.Product import Data.Hashable import Data.HashMap.Strict (HashMap) import qualified Data.HashMap.Strict as HM +import Data.List (sortBy, intersperse) +import Data.Maybe (fromMaybe) +import Data.Ord (comparing) import Data.Some import Data.Type.Equality import GHC.StableName import Numeric.Natural import Unsafe.Coerce (unsafeCoerce) +-- import Debug.Trace + import Data.StableName.Extra @@ -41,6 +46,16 @@ import Data.StableName.Extra -- is a good opportunity to try to do better. +withMoreState :: Functor m => b -> StateT (s, b) m a -> StateT s m (a, b) +withMoreState b0 (StateT f) = + StateT $ \s -> (\(x, (s2, b)) -> ((x, b), s2)) <$> f (s, b0) + +withLessState :: Functor m => (s -> (s', b)) -> (s' -> b -> s) -> StateT s' m a -> StateT s m a +withLessState split restore (StateT f) = + StateT $ \s -> let (s', b) = split s + in second (flip restore b) <$> f s' + + class Functor1 f where fmap1 :: (forall b. g b -> h b) -> f g a -> f h a @@ -91,49 +106,85 @@ instance Eq (SomeNameFor typ f) where instance Hashable (SomeNameFor typ f) where hashWithSalt salt (SomeNameFor name) = hashWithSalt salt name --- | The number of times a particular name is visited in a preorder traversal --- of the PHOAS expression, excluding children of nodes upon second or later --- visit. That is to say: only the nodes that are visited in a preorder --- traversal that skips repeated subtrees, are counted. -type OccMap typ f = HashMap (SomeNameFor typ f) Natural +prettyPExpr :: Traversable1 f => Int -> PExpr typ f t -> ShowS +prettyPExpr d = \case + PStub (NameFor name) _ -> showString (showStableName name) + POp (NameFor name) _ args -> + let (argslist, _) = traverse1 (\arg -> ([Some arg], Const ())) args + argslist' = map (\(Some arg) -> prettyPExpr 0 arg) argslist + in showParen (d > 10) $ + showString ("<" ++ showStableName name ++ ">(") + . foldr (.) id (intersperse (showString ", ") argslist') + . showString ")" + PLam (NameFor name) _ _ (Tag tag) body -> + showParen (d > 0) $ + showString ("λ" ++ showStableName name ++ " x" ++ show tag ++ ". ") . prettyPExpr 0 body + PVar _ (Tag tag) -> showString ("x" ++ show tag) + +-- | For each name: +-- +-- 1. The number of times the name is visited in a preorder traversal of the +-- PHOAS expression, excluding children of nodes upon second or later visit. +-- That is to say: only the nodes that are visited in a preorder traversal +-- that skips repeated subtrees, are counted. +-- 2. The height of the expression indicated by the name. +-- +-- Missing names have not been seen yet, and have unknown height. +type OccMap typ f = HashMap (SomeNameFor typ f) (Natural, Natural) pruneExpr :: Traversable1 f => (forall v. PHOASExpr typ v f t) -> (OccMap typ f, PExpr typ f t) pruneExpr term = - let (term', (_, mp)) = runState (pruneExpr' term) (0, mempty) + let ((term', _), (_, mp)) = runState (pruneExpr' term) (0, mempty) in (mp, term') -pruneExpr' :: Traversable1 f => PHOASExpr typ Tag f t -> State (Natural, OccMap typ f) (PExpr typ f t) +-- | Returns pruned expression with its height. +pruneExpr' :: Traversable1 f => PHOASExpr typ Tag f t -> State (Natural, OccMap typ f) (PExpr typ f t, Natural) pruneExpr' = \case orig@(PHOASOp ty args) -> do let name = makeStableName' orig - seenBefore <- checkVisited name - if seenBefore - then pure $ PStub (NameFor name) ty - else POp (NameFor name) ty <$> traverse1 pruneExpr' args + mheight <- gets (fmap snd . HM.lookup (SomeNameFor (NameFor name)) . snd) + case mheight of + -- already visited + Just height -> do + modify (second (HM.adjust (first (+1)) (SomeNameFor (NameFor name)))) + pure (PStub (NameFor name) ty, height) + -- first visit + Nothing -> do + -- Traverse the arguments, collecting the maximum height in an + -- additional piece of state. + (args', maxhei) <- + withMoreState 0 $ + traverse1 (\arg -> do + (arg', hei) <- withLessState id (,) (pruneExpr' arg) + modify (second (hei `max`)) + return arg') + args + -- Record this node + modify (second (HM.insert (SomeNameFor (NameFor name)) (1, 1 + maxhei))) + pure (POp (NameFor name) ty args', 1 + maxhei) orig@(PHOASLam tyf tyarg f) -> do let name = makeStableName' orig - seenBefore <- checkVisited name - if seenBefore - then pure $ PStub (NameFor name) tyf - else do - tag <- state (\(i, mp) -> (Tag i, (i + 1, mp))) + mheight <- gets (fmap snd . HM.lookup (SomeNameFor (NameFor name)) . snd) + case mheight of + -- already visited + Just height -> do + modify (second (HM.adjust (first (+1)) (SomeNameFor (NameFor name)))) + pure (PStub (NameFor name) tyf, height) + -- first visit + Nothing -> do + tag <- Tag <$> gets fst + modify (first (+1)) let body = f tag - PLam (NameFor name) tyf tyarg tag <$> pruneExpr' body + (body', bodyhei) <- pruneExpr' body + modify (second (HM.insert (SomeNameFor (NameFor name)) (1, 1 + bodyhei))) + pure (PLam (NameFor name) tyf tyarg tag body', 1 + bodyhei) - PHOASVar ty tag -> pure $ PVar ty tag - where - checkVisited name = do - occmap <- gets snd - let (seenBefore, occmap') = - HM.alterF (\case Nothing -> (False, Just 1) - Just n -> (True, Just (n + 1))) - (SomeNameFor (NameFor name)) - occmap - modify (second (const occmap')) - return seenBefore + PHOASVar ty tag -> pure (PVar ty tag, 1) +-- TODO: Replace "lift" with "float" + -- | Lifted expression: a bunch of to-be let bound expressions on top of an -- LExpr'. Because LExpr' is really just PExpr with the recursive positions -- replaced by LExpr, LExpr should be seen as PExpr with a bunch of to-be let @@ -145,12 +196,35 @@ data LExpr' typ f t where -- TODO: this could be an instantiation of (a general LLam :: NameFor typ f (a -> b) -> typ (a -> b) -> typ a -> Tag a -> LExpr typ f b -> LExpr' typ f (a -> b) LVar :: typ a -> Tag a -> LExpr' typ f a +prettyLExpr :: Traversable1 f => Int -> LExpr typ f t -> ShowS +prettyLExpr d (LExpr [] e) = prettyLExpr' d e +prettyLExpr d (LExpr lifted e) = + showString "[" + . foldr (.) id (intersperse (showString ", ") (map (\(Some e') -> prettyLExpr 0 e') lifted)) + . showString "] " . prettyLExpr' d e + +prettyLExpr' :: Traversable1 f => Int -> LExpr' typ f t -> ShowS +prettyLExpr' d = \case + LStub (NameFor name) _ -> showString (showStableName name) + LOp (NameFor name) _ args -> + let (argslist, _) = traverse1 (\arg -> ([Some arg], Const ())) args + argslist' = map (\(Some arg) -> prettyLExpr 0 arg) argslist + in showParen (d > 10) $ + showString ("<" ++ showStableName name ++ ">(") + . foldr (.) id (intersperse (showString ", ") argslist') + . showString ")" + LLam (NameFor name) _ _ (Tag tag) body -> + showParen (d > 0) $ + showString ("λ" ++ showStableName name ++ " x" ++ show tag ++ ". ") . prettyLExpr 0 body + LVar _ (Tag tag) -> showString ("x" ++ show tag) + liftExpr :: Traversable1 f => OccMap typ f -> PExpr typ f t -> LExpr typ f t liftExpr totals term = snd (liftExpr' totals term) newtype FoundMap typ f = FoundMap - (HashMap (SomeNameFor typ f) (Natural -- how many times seen - ,Maybe (Some (LExpr typ f)))) -- the lifted subterm (once seen) + (HashMap (SomeNameFor typ f) + (Natural -- how many times seen + ,Maybe (Some (LExpr typ f), Natural))) -- the lifted subterm with its height (once seen) instance Semigroup (FoundMap typ f) where FoundMap m1 <> FoundMap m2 = FoundMap $ @@ -161,10 +235,13 @@ instance Monoid (FoundMap typ f) where liftExpr' :: Traversable1 f => OccMap typ f -> PExpr typ f t -> (FoundMap typ f, LExpr typ f t) liftExpr' _totals (PStub name ty) = - (FoundMap $ HM.singleton (SomeNameFor name) (1, Nothing) -- Just (Some (LExpr [] (LStub name))) + -- trace ("Found stub: " ++ (case name of NameFor n -> showStableName n)) $ + (FoundMap $ HM.singleton (SomeNameFor name) (1, Nothing) ,LExpr [] (LStub name ty)) -liftExpr' _totals (PVar ty tag) = (mempty, LExpr [] (LVar ty tag)) +liftExpr' _totals (PVar ty tag) = + -- trace ("Found var: " ++ show tag) $ + (mempty, LExpr [] (LVar ty tag)) liftExpr' totals term = let (FoundMap foundmap, name, termty, term') = case term of @@ -178,19 +255,24 @@ liftExpr' totals term = -- TODO: perhaps this HM.toList together with the foldr HM.delete can be a single traversal of the HashMap saturated = [case mterm of Just t -> (nm, t) - Nothing -> error "Name saturated but no term found" + Nothing -> case nm of + SomeNameFor (NameFor n) -> + error $ "Name saturated (count=" ++ show count ++ ", totalcount=" ++ show totalcount ++ ") but no term found: " ++ showStableName n | (nm, (count, mterm)) <- HM.toList foundmap - , count == HM.findWithDefault 0 nm totals] + , let totalcount = fromMaybe 0 (fst <$> HM.lookup nm totals) + , count == totalcount] foundmap' = foldr HM.delete foundmap (map fst saturated) - lterm = LExpr (map snd saturated) term' + lterm = LExpr (map fst (sortBy (comparing snd) (map snd saturated))) term' - in case HM.findWithDefault 0 (SomeNameFor name) totals of - 1 -> (FoundMap foundmap', lterm) - tot | tot > 1 -> (FoundMap (HM.insert (SomeNameFor name) (1, Just (Some lterm)) foundmap') - ,LExpr [] (LStub name termty)) - | otherwise -> error "Term does not exist, yet we have it in hand" + in case HM.findWithDefault (0, undefined) (SomeNameFor name) totals of + (1, _) -> (FoundMap foundmap', lterm) + (tot, height) + | tot > 1 -> -- trace ("Inserting " ++ (case name of NameFor n -> showStableName n) ++ " into foundmap") $ + (FoundMap (HM.insert (SomeNameFor name) (1, Just (Some lterm, height)) foundmap') + ,LExpr [] (LStub name termty)) + | otherwise -> error "Term does not exist, yet we have it in hand" -- | Untyped De Bruijn expression. No more names: there are lets now, and @@ -334,4 +416,10 @@ retypeExpr' env (UBVar ty idx) = sharingRecovery :: (Traversable1 f, TestEquality typ) => (forall v. PHOASExpr typ v f t) -> BExpr typ '[] f t -sharingRecovery e = retypeExpr $ lowerExpr $ uncurry liftExpr $ pruneExpr e +sharingRecovery e = + let (occmap, pexpr) = pruneExpr e + lexpr = liftExpr occmap pexpr + ubexpr = lowerExpr lexpr + in -- trace ("PExpr: " ++ prettyPExpr 0 pexpr "") $ + -- trace ("LExpr: " ++ prettyLExpr 0 lexpr "") $ + retypeExpr ubexpr diff --git a/src/Data/StableName/Extra.hs b/src/Data/StableName/Extra.hs index f568740..cf37cfe 100644 --- a/src/Data/StableName/Extra.hs +++ b/src/Data/StableName/Extra.hs @@ -1,10 +1,16 @@ {-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ExistentialQuantification #-} {-# OPTIONS_GHC -fno-full-laziness -fno-cse #-} module Data.StableName.Extra ( StableName, makeStableName', + showStableName, ) where +import Data.Hashable +import Data.HashMap.Strict (HashMap) +import qualified Data.HashMap.Strict as HM +import Data.IORef import GHC.StableName import System.IO.Unsafe @@ -15,3 +21,24 @@ import System.IO.Unsafe {-# NOINLINE makeStableName' #-} makeStableName' :: a -> StableName a makeStableName' !x = unsafePerformIO (makeStableName x) + + +data SomeStableName = forall a. SomeStableName (StableName a) + +instance Eq SomeStableName where + SomeStableName a == SomeStableName b = eqStableName a b + +instance Hashable SomeStableName where + hashWithSalt salt (SomeStableName name) = hashWithSalt salt name + +{-# NOINLINE showStableNameCache #-} +showStableNameCache :: IORef (HashMap SomeStableName Int, Int) +showStableNameCache = unsafePerformIO $ newIORef (mempty, 0) + +{-# NOINLINE showStableName #-} +showStableName :: StableName a -> String +showStableName name = + unsafePerformIO $ atomicModifyIORef' showStableNameCache $ \tup@(mp, nexti) -> + case HM.lookup (SomeStableName name) mp of + Just res -> (tup, '$' : show res) + Nothing -> ((HM.insert (SomeStableName name) nexti mp, nexti + 1), '$' : show nexti) |