content: 
  -- | Bubble casts outwards.
  --   We float casts up and outwards so they are just inside the inner-most
  --   enclosing let. This way the functions still have the same effect and
  --   closure, but the casts don't get in the way of subsequent transforms.
  --
  module DDC.Core.Transform.Bubble
          ( bubbleModule
          , bubbleX)
  where
  import DDC.Core.Collect
  import DDC.Core.Transform.BoundX
  import DDC.Core.Module
  import DDC.Core.Exp.Annot
  import DDC.Type.Env                             (KindEnv, TypeEnv)
  import qualified DDC.Type.Env                   as Env
  import qualified DDC.Type.Sum                   as Sum
  import qualified Data.Set                       as Set
  import Data.Set                                 (Set)
  import Data.List
  
  
  -- | Bubble casts outwards in a `Module`.
  bubbleModule
          :: Ord n
          => Module a n -> Module a n
  
  bubbleModule mm@ModuleCore{}
   = let  kenv    = moduleKindEnv mm
          tenv    = moduleTypeEnv mm
     in   mm { moduleBody = bubbleX kenv tenv (moduleBody mm) }
  
  
  -- | Bubble casts outwards in an `Exp`.
  bubbleX :: Ord n
          => KindEnv n -> TypeEnv n -> Exp a n -> Exp a n
  bubbleX kenv tenv x
   = let
          -- Bubble the expression, yielding any casts that floated out
          -- to top-level.
          (cs, x') = bubble kenv tenv x
          a        = annotOfExp x'
  
          -- Reattach top-level casts.
     in   dropAllCasts kenv tenv a cs x'
  
  
  -------------------------------------------------------------------------------
  -- | Bubble casts outwards in some thing.
  class Bubble (c :: * -> * -> *) where
   bubble :: Ord n
          => KindEnv n
          -> TypeEnv n
          -> c a n
          -> ([FvsCast a n], c a n)
  
  
  instance Bubble Exp where
   bubble kenv tenv xx
    = {-# SCC bubble #-}
      case xx of
          XVar{}  -> ([], xx)
          XAtom{} -> ([], xx)
  
          -- Drop casts before we leave lambda abstractions, because the
          -- function type depends on the effect and closure of the body.
          -- The cast could also reference the bound variable.
          XAbs a (MType b) x
           -> let kenv'           = Env.extend b kenv
                  (cs, x')        = bubble kenv' tenv x
              in  ([], XAbs a (MType b) (dropAllCasts kenv' tenv a cs x'))
  
          XAbs a (MTerm b) x
           -> let tenv'           = Env.extend b tenv
                  (cs, x')        = bubble kenv tenv' x
              in  ([], XAbs a (MTerm b) (dropAllCasts kenv tenv' a cs x'))
  
          XAbs a (MImplicit b) x
           -> let tenv'           = Env.extend b tenv
                  (cs, x')        = bubble kenv tenv' x
              in  ([], XAbs a (MImplicit b) (dropAllCasts kenv tenv' a cs x'))
  
          -- Decend into applications.
          XApp a x1 x2
           -> let (cs1, x1')      = bubble kenv tenv x1
                  (cs2, x2')      = bubble kenv tenv x2
              in  (cs1 ++ cs2, XApp a x1' x2')
  
          -- After decending into let-expressions, make sure to drop
          -- any casts that mention variables bound here.
          XLet a lts x2
           -> let (cs1, lts')     = bubble kenv tenv lts
                  (bs1, bs0)      = bindsOfLets lts
                  kenv'           = Env.extends bs1 kenv
                  tenv'           = Env.extends bs0 tenv
                  (cs2, x2')      = bubble kenv' tenv' x2
                  (cs2', x2'')    = dropCasts kenv' tenv' a bs1 bs0 cs2 x2'
              in  ( cs1 ++ cs2'
                  , XLet a lts' x2'')
  
          -- Decent into case-expressions.
          --  Casts that depend on bound variables are dropped
          --  in the corresponding alternatives.
          XCase a x alts
           -> let (cs, x')        = bubble kenv tenv x
                  (css, alts')    = unzip $ map (bubble kenv tenv) alts
              in  ( cs ++ concat css
                  , XCase a x' alts')
  
          -- Strip of cast and pass it up.
          XCast _ c x
           -> let (cs, x')        = bubble kenv tenv x
                  fvsT            = freeT Env.empty c
                  fvsX            = freeX Env.empty c
                  fc              = FvsCast c fvsT fvsX
              in  (fc : cs, x')
  
  
  instance Bubble Arg where
   bubble kenv tenv aa
    = case aa of
          RType{}
           ->     ([], aa)
  
          RWitness{}
           ->     ([], aa)
  
          RTerm x
           -> let (cs, x') = bubble kenv tenv x
              in  (cs, RTerm x')
  
          RImplicit x
           -> let (cs, x') = bubble kenv tenv x
              in  (cs, RImplicit x')
  
  
  instance Bubble Lets where
   bubble kenv tenv lts
    = case lts of
  
          -- Drop casts that mention the bound variable here,
          -- but we can float the others further outwards.
          LLet b x
           -> let (cs, x')        = bubble kenv tenv x
                  a               = annotOfExp x'
                  (cs', xc')      = dropCasts kenv tenv a [] [b] cs x'
              in  (cs', LLet b xc')
  
          -- ISSUE #299: Bubble casts out of recursive lets.
          LRec bxs
           -> let bs              = map fst bxs
                  tenv'           = Env.extends bs tenv
  
                  bubbleRec (b, x)
                   = let  (cs, x') = bubble kenv tenv' x
                          a        = annotOfExp x'
                     in   (b, dropAllCasts kenv tenv' a cs x')
  
                  bxs'            = map bubbleRec bxs
  
              in  ([], LRec bxs')
  
          LPrivate{}
           -> ([], lts)
  
  
  instance Bubble Alt where
  
   -- Default patterns don't bind variables,
   -- so there is no problem floating casts outwards.
   bubble kenv tenv (AAlt PDefault x)
    = let (cs, x') = bubble kenv tenv x
      in  (cs, AAlt PDefault x')
  
   -- Drop casts before we leave the alt because they could contain
   -- variables bound by the pattern.
   bubble kenv tenv (AAlt p x)
    = let bs              = bindsOfPat p
          a               = annotOfExp x'
          tenv'           = Env.extends bs tenv
          (cs, x')        = bubble kenv tenv' x
          (csUp, xcHere)  = dropCasts kenv tenv' a [] bs cs x'
      in  (csUp, AAlt p xcHere)
  
  
  -- FvsCast --------------------------------------------------------------------
  -- | A Cast along with its free level-1 and level-0 vars.
  --   When we first build a `FvsCast` we record its free variables,
  --   so that we don't have to keep recomputing them.
  data FvsCast a n
          = FvsCast (Cast a n)
                    (Set (Bound n))       -- Free level-1 variables.
                    (Set (Bound n))       -- Free level-0 variables.
  
  instance Ord n => MapBoundX (FvsCast a) n where
   mapBoundAtDepthX f d (FvsCast c fvs1 fvs0)
    = FvsCast (mapBoundAtDepthX f d c)
              fvs1
              (Set.fromList
                  $ map (mapBoundAtDepthX f d)
                  $ Set.toList fvs0)
  
  
  packFvsCasts
          :: Ord n
          => KindEnv n -> TypeEnv n
          -> a -> [FvsCast a n] -> [Cast a n]
  
  packFvsCasts kenv tenv a fvsCasts
          = packCasts kenv tenv a [ c | FvsCast c _ _ <- fvsCasts ]
  
  
  -- | Pack down casts by combining multiple 'weakclo' and 'weakeff' casts
  --   together. We pack casts just before we drop them, so that the resulting
  --   code is easier to read.
  packCasts :: Ord n
            => KindEnv n -> TypeEnv n -> a -> [Cast a n] -> [Cast a n]
  packCasts _kenv _tenv _a vs
   = let
          -- Sort casts into weakeff, weakclo and any others.
          -- We'll collect the weakeff and weakclo casts together.
          collect weakEffs weakClos others cc
           = case cc of
              []
               -> (reverse weakEffs, reverse weakClos, reverse others)
  
              CastWeakenEffect eff : cs
               -> collect (eff : weakEffs) weakClos others cs
  
              c : cs
               -> collect weakEffs        weakClos (c : others) cs
  
  
          (effs, csOthers, _)
                  = collect [] [] [] vs
  
     in   (if null effs
                  then []
                  else [CastWeakenEffect  (TSum $ Sum.fromList kEffect effs)])
       ++ csOthers
  
  
  -- Dropping -------------------------------------------------------------------
  -- | Wrap the provided expression with these casts.
  dropAllCasts
          :: Ord n
          => KindEnv n
          -> TypeEnv n
          -> a
          -> [FvsCast a n] -> Exp a n
          -> Exp a n
  
  dropAllCasts kenv tenv a cs x
   = let  cs'     = packFvsCasts kenv tenv a cs
     in   foldr (XCast a) x cs'
  
  
  -- | Split the provided casts into ones that contain variables
  --   bound by these binders. The casts that do are used to wrap
  --   the provided expression, and the casts that don't are returned
  --   seprately so we can keep bubbling them up the tree.
  dropCasts
          :: Ord n
          => KindEnv n -> TypeEnv n
          -> a
          -> [Bind n]             -- ^ Level-1 binders.
          -> [Bind n]             -- ^ Level-0 binders.
          -> [FvsCast a n]
          -> Exp a n
          -> ([FvsCast a n], Exp a n)
  
  dropCasts kenv tenv a bs1 bs0 cs x
   = let  (csHere1, cs1)    = partition (fvsCastUsesBinds1 bs1) cs
          (csHere0, csUp)   = partition (fvsCastUsesBinds0 bs0) cs1
          csHere            = packFvsCasts kenv tenv a $ csHere1 ++ csHere0
     in   ( map (lowerX 1) csUp
          , foldr (XCast a) x csHere)
  
  
  -- | Check if a `FvsCast` mentions any of these level-0 variables.
  fvsCastUsesBinds0 :: Ord n => [Bind n] -> FvsCast a n -> Bool
  fvsCastUsesBinds0 bb (FvsCast _ _ fvs0)
          = bindsMatchBoundSet bb fvs0
  
  
  -- | Check if a `FvsCast` mentions any of these level-1 variables.
  fvsCastUsesBinds1 :: Ord n => [Bind n] -> FvsCast a n -> Bool
  fvsCastUsesBinds1 bb (FvsCast _ fvs1 _)
          = bindsMatchBoundSet bb fvs1
  
  
  -- | Check if a set of bound variables matches any of the given binders.
  bindsMatchBoundSet :: Ord n => [Bind n] -> Set (Bound n) -> Bool
  bindsMatchBoundSet bb fvs
   = go bb
   where  go []           = False
          go (b : bs)
           | Just u       <- takeSubstBoundOfBind b
           = if Set.member u fvs
                  then True
                  else go bs
  
           | otherwise
           = go bs
  
  