CleanupHEADmaster

1 files changed, 90 insertions, 100 deletions

diff --git a/src/Data/Expr/SharingRecovery/Internal.hs b/src/Data/Expr/SharingRecovery/Internal.hs
index 0089454..9f7ca7d 100644
--- a/src/Data/Expr/SharingRecovery/Internal.hs
+++ b/src/Data/Expr/SharingRecovery/Internal.hs
@@ -100,7 +100,7 @@ class Functor1 f => Traversable1 f where
 --
 --     Note furthermore that @Oper@ is /not/ a recursive type. Subexpressions
 --     are again 'PHOASExpr's, and 'sharingRecovery' needs to be able to see
---     them. Hence, you should call back to back to @r@ instead of recursing
+--     them. Hence, you should call back to @r@ instead of recursing
 --     manually.
 --
 -- * @t@ is the result type of this expression.
@@ -130,11 +130,16 @@ instance TestEquality (NameFor typ f) where
 -- Note that variables do not, and will never, have a name: we don't bother
 -- detecting sharing for variable references, because that would only introduce
 -- a redundant variable indirection.
-data PExpr typ f t where
-  PStub :: NameFor typ f t -> typ t -> PExpr typ f t
-  POp :: NameFor typ f t -> typ t -> f (PExpr typ f) t -> PExpr typ f t
-  PLam :: NameFor typ f (a -> b) -> typ (a -> b) -> typ a -> Tag a -> PExpr typ f b -> PExpr typ f (a -> b)
-  PVar :: typ a -> Tag a -> PExpr typ f a
+--
+-- This is defined as a base functor; @r@ is the recursive position.
+data PExpr r typ f t where
+  PStub :: NameFor typ f t -> typ t -> PExpr r typ f t
+  POp :: NameFor typ f t -> typ t -> f (r typ f) t -> PExpr r typ f t
+  PLam :: NameFor typ f (a -> b) -> typ (a -> b) -> typ a -> Tag a -> r typ f b -> PExpr r typ f (a -> b)
+  PVar :: typ a -> Tag a -> PExpr r typ f a
+
+-- | Fixpoint of 'PExpr'
+newtype PExpr0 typ f t = PExpr0 (PExpr PExpr0 typ f t)
 
 data SomeNameFor typ f = forall t. SomeNameFor {-# UNPACK #-} !(NameFor typ f t)
 
@@ -144,19 +149,22 @@ instance Eq (SomeNameFor typ f) where
 instance Hashable (SomeNameFor typ f) where
   hashWithSalt salt (SomeNameFor name) = hashWithSalt salt name
 
-prettyPExpr :: Traversable1 f => Int -> PExpr typ f t -> ShowS
-prettyPExpr d = \case
+prettyPExpr0 :: Traversable1 f => Int -> PExpr0 typ f t -> ShowS
+prettyPExpr0 d (PExpr0 ex) = prettyPExpr prettyPExpr0 d ex
+
+prettyPExpr :: Traversable1 f => (forall a. Int -> r typ f a -> ShowS) -> Int -> PExpr r typ f t -> ShowS
+prettyPExpr recur 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
+        argslist' = map (\(Some arg) -> recur 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
+      showString ("λ" ++ showStableName name ++ " x" ++ show tag ++ ". ") . recur 0 body
   PVar _ (Tag tag) -> showString ("x" ++ show tag)
 
 -- | For each name:
@@ -170,13 +178,14 @@ prettyPExpr d = \case
 -- 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 :: Traversable1 f => (forall v. PHOASExpr typ v f t) -> (OccMap typ f, PExpr0 typ f t)
 pruneExpr term =
   let ((term', _), (_, mp)) = runState (pruneExpr' term) (0, mempty)
   in (mp, term')
 
 -- | Returns pruned expression with its height.
-pruneExpr' :: Traversable1 f => PHOASExpr typ Tag f t -> State (Natural, OccMap typ f) (PExpr typ f t, Natural)
+-- State: (ID generator, occurrence map being accumulated)
+pruneExpr' :: Traversable1 f => PHOASExpr typ Tag f t -> State (Natural, OccMap typ f) (PExpr0 typ f t, Natural)
 pruneExpr' = \case
   orig@(PHOASOp ty args) -> do
     let name = makeStableName' orig
@@ -185,7 +194,7 @@ pruneExpr' = \case
       -- already visited
       Just height -> do
         modify (second (HM.adjust (first (+1)) (SomeNameFor (NameFor name))))
-        pure (PStub (NameFor name) ty, height)
+        pure (PExpr0 (PStub (NameFor name) ty), height)
       -- first visit
       Nothing -> do
         -- Traverse the arguments, collecting the maximum height in an
@@ -193,13 +202,14 @@ pruneExpr' = \case
         (args', maxhei) <-
           withMoreState 0 $
             traverse1 (\arg -> do
+                        -- drop the extra state for the recursive call
                         (arg', hei) <- withLessState id (,) (pruneExpr' arg)
-                        modify (second (hei `max`))
+                        modify (second (hei `max`))  -- modify the extra state
                         return arg')
                       args
         -- Record this node
         modify (second (HM.insert (SomeNameFor (NameFor name)) (1, 1 + maxhei)))
-        pure (POp (NameFor name) ty args', 1 + maxhei)
+        pure (PExpr0 (POp (NameFor name) ty args'), 1 + maxhei)
 
   orig@(PHOASLam tyf tyarg f) -> do
     let name = makeStableName' orig
@@ -208,7 +218,7 @@ pruneExpr' = \case
       -- already visited
       Just height -> do
         modify (second (HM.adjust (first (+1)) (SomeNameFor (NameFor name))))
-        pure (PStub (NameFor name) tyf, height)
+        pure (PExpr0 (PStub (NameFor name) tyf), height)
       -- first visit
       Nothing -> do
         tag <- Tag <$> gets fst
@@ -216,45 +226,24 @@ pruneExpr' = \case
         let body = f tag
         (body', bodyhei) <- pruneExpr' body
         modify (second (HM.insert (SomeNameFor (NameFor name)) (1, 1 + bodyhei)))
-        pure (PLam (NameFor name) tyf tyarg tag body', 1 + bodyhei)
+        pure (PExpr0 (PLam (NameFor name) tyf tyarg tag body'), 1 + bodyhei)
 
-  PHOASVar ty tag -> pure (PVar ty tag, 1)
+  PHOASVar ty tag -> pure (PExpr0 (PVar ty tag), 1)
 
 
--- | Floated 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
--- bound expressions on top of every node.
-data LExpr typ f t = LExpr [Some (LExpr typ f)] (LExpr' typ f t)
-data LExpr' typ f t where  -- TODO: this could be an instantiation of (a generalisation of) PExpr
-  LStub :: NameFor typ f t -> typ t -> LExpr' typ f t
-  LOp :: NameFor typ f t -> typ t -> f (LExpr typ f) t -> LExpr' typ f t
-  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
+-- | Floated expression: again a 'PExpr' (it's a fixpoint over the same base
+-- functor), but now with a bunch of to-be let bound expressions on top of
+-- every node.
+data LExpr typ f t = LExpr [Some (LExpr typ f)] (PExpr LExpr typ f t)
 
 prettyLExpr :: Traversable1 f => Int -> LExpr typ f t -> ShowS
-prettyLExpr d (LExpr [] e) = prettyLExpr' d e
+prettyLExpr d (LExpr [] e) = prettyPExpr prettyLExpr d e
 prettyLExpr d (LExpr floated e) =
   showString "["
   . foldr (.) id (intersperse (showString ", ") (map (\(Some e') -> prettyLExpr 0 e') floated))
-  . 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)
+  . showString "] " . prettyPExpr prettyLExpr d e
 
-floatExpr :: Traversable1 f => OccMap typ f -> PExpr typ f t -> LExpr typ f t
+floatExpr :: Traversable1 f => OccMap typ f -> PExpr0 typ f t -> LExpr typ f t
 floatExpr totals term = snd (floatExpr' totals term)
 
 newtype FoundMap typ f = FoundMap
@@ -269,46 +258,47 @@ instance Semigroup (FoundMap typ f) where
 instance Monoid (FoundMap typ f) where
   mempty = FoundMap HM.empty
 
-floatExpr' :: Traversable1 f => OccMap typ f -> PExpr typ f t -> (FoundMap typ f, LExpr typ f t)
-floatExpr' _totals (PStub name ty) =
-  -- trace ("Found stub: " ++ (case name of NameFor n -> showStableName n)) $
-  (FoundMap $ HM.singleton (SomeNameFor name) (1, Nothing)
-  ,LExpr [] (LStub name ty))
-
-floatExpr' _totals (PVar ty tag) =
-  -- trace ("Found var: " ++ show tag) $
-  (mempty, LExpr [] (LVar ty tag))
-
-floatExpr' totals term =
-  let (FoundMap foundmap, name, termty, term') = case term of
-        POp n ty args ->
-          let (fm, args') = traverse1 (floatExpr' totals) args
-          in (fm, n, ty, LOp n ty args')
-        PLam n tyf tyarg tag body ->
-          let (fm, body') = floatExpr' totals body
-          in (fm, n, tyf, LLam n tyf tyarg tag body')
-
-      -- 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 -> 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
-                  , let totalcount = fromMaybe 0 (fst <$> HM.lookup nm totals)
-                  , count == totalcount]
-
-      foundmap' = foldr HM.delete foundmap (map fst saturated)
-
-      lterm = LExpr (map fst (sortBy (comparing snd) (map snd saturated))) term'
-
-  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"
+floatExpr' :: Traversable1 f => OccMap typ f -> PExpr0 typ f t -> (FoundMap typ f, LExpr typ f t)
+floatExpr' totals (PExpr0 term) = case term of
+  PStub name ty ->
+    -- trace ("Found stub: " ++ (case name of NameFor n -> showStableName n)) $
+    (FoundMap $ HM.singleton (SomeNameFor name) (1, Nothing)
+    ,LExpr [] (PStub name ty))
+
+  PVar ty tag ->
+    -- trace ("Found var: " ++ show tag) $
+    (mempty, LExpr [] (PVar ty tag))
+
+  _ ->
+    let (FoundMap foundmap, name, termty, term') = case term of
+          POp n ty args ->
+            let (fm, args') = traverse1 (floatExpr' totals) args
+            in (fm, n, ty, POp n ty args')
+          PLam n tyf tyarg tag body ->
+            let (fm, body') = floatExpr' totals body
+            in (fm, n, tyf, PLam n tyf tyarg tag body')
+
+        -- 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 -> 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
+                    , let totalcount = fromMaybe 0 (fst <$> HM.lookup nm totals)
+                    , count == totalcount]
+
+        foundmap' = foldr HM.delete foundmap (map fst saturated)
+
+        lterm = LExpr (map fst (sortBy (comparing snd) (map snd saturated))) term'
+
+    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 [] (PStub 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
@@ -342,15 +332,15 @@ lowerExpr' namelvl taglvl curlvl (LExpr floated ex) =
       curlvl' = curlvl + length floated
   in prefix $
        case ex of
-         LStub name ty ->
+         PStub 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 ->
+         POp _ ty args ->
            UBOp ty (fmap1 (lowerExpr' namelvl' taglvl curlvl') args)
-         LLam _ tyf tyarg tag body ->
+         PLam _ tyf tyarg tag body ->
            UBLam tyf tyarg (lowerExpr' namelvl' (HM.insert (SomeTag tag) curlvl' taglvl) (curlvl' + 1) body)
-         LVar ty tag ->
+         PVar ty tag ->
            case HM.lookup (SomeTag tag) taglvl of
              Just lvl -> UBVar ty (curlvl - lvl - 1)
              Nothing -> error "Tag variable out of scope"
@@ -363,17 +353,17 @@ lowerExpr' namelvl taglvl curlvl (LExpr floated ex) =
     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"
+                   PStub n _ -> n
+                   POp n _ _ -> n
+                   PLam n _ _ _ _ -> n
+                   PVar _ _ -> 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
+                 PStub{} -> error "Recovering sharing of a stub is useless"
+                 POp _ t _ -> t
+                 PLam _ t _ _ _ -> t
+                 PVar t _ -> t
           prefix = UBLet ty (lowerExpr' namelvl' taglvl lvl rhs)
-      in (prefix .) <$> buildPrefix (HM.insert (SomeNameFor name) lvl namelvl') (lvl + 1) rhss
+      in second (prefix .) $ buildPrefix (HM.insert (SomeNameFor name) lvl namelvl') (lvl + 1) rhss
 
 
 -- | A typed De Bruijn index.
@@ -399,7 +389,7 @@ envLookupU = go id
     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 resu,t of sharing recovery. It is
+-- | Typed De Bruijn expression. This is the result of sharing recovery. It is
 -- not higher-order any more, and furthermore has explicit let-bindings ('BLet')
 -- that denote the sharing inside the term. This is a normal AST.
 --