-----------------------------------------------------------------------
-- |
-- Module           : Lang.Crucible.CFG.EarlyMergeLoops
-- Description      : Provides transformations on pre-SSA CFGs
-- Copyright        : (c) Galois, Inc 2020
-- License          : BSD3
-- Maintainer       : 
-- Stability        : experimental
--
-- This modules exposes a transformation that attempts to ensure that loop branches
-- are post-dominated by nodes in the loop.
--
-- The module is organized into 3 main components:
--   1. An analysis that computes the natural loops of a CFG;
--   2. An analysis that inserts postdominators into loops that have
--      "early exits" (and hence have postdominators outside the loop);
--   3. A "fixup" pass that ensures that, in the transformed CFG, all
--      values are well defined along all (new) paths.
------------------------------------------------------------------------
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}

module Lang.Crucible.CFG.EarlyMergeLoops
  ( earlyMergeLoops
  ) where

import           Control.Monad (when, (>=>))
import           Control.Applicative ((<**>))
import qualified Data.Graph.Inductive as G
import qualified Data.Foldable as Fold
import           Data.Kind
import           Data.List (nub, minimumBy)
import qualified Data.Map.Strict as Map
import           Data.Maybe (fromMaybe)
import           Data.Parameterized.Classes
import qualified Data.Parameterized.Context as Ctx
import qualified Data.Parameterized.Map as MapF
import           Data.Parameterized.Nonce
import           Data.Parameterized.Some
import           Data.Parameterized.TraversableFC
import           Data.Set (Set)
import qualified Data.Set as Set
import           Data.Sequence (Seq)
import qualified Data.Sequence as Seq ((<|), fromList)
import           Data.String (fromString)

import           What4.ProgramLoc (Position(..), Posd(..))

import           Lang.Crucible.CFG.Expr
import           Lang.Crucible.CFG.Reg
import           Lang.Crucible.Panic
import           Lang.Crucible.Types

--------------------------
-- | Natural Loop Analysis
--------------------------

-- | This structure identifies natural loops. Natural loops are either
-- disjoint from each other, nested, or they share the same header.
data LoopInfo s = LoopInfo
  { liFooter     :: !(BlockID s)
    -- ^ This is the block with a backedge (to the header)
  , liHeader     :: !(BlockID s)
    -- ^ This is the destination of the backedge.
  , liMembers    :: !(Set (BlockID s))
  -- ^ The loop members, which is the set of nodes that can reach the footer
  -- without going through the header.
  , liEarlyExits :: ![CFGEdge s]
  -- ^ An exiting edge is an edge from a node in the loop to an edge
  -- not in the loop. An early exit is such an edge from a node that
  -- is not the footer node.
  , liFooterIn   :: ![CFGEdge s]
  -- ^ All the edges to the footer
  , liDominators :: ![BlockID s]
  -- ^ The dominators of the loop header.
  } deriving (Eq, Show)
type CFGEdge s = (BlockID s, BlockID s)

-- | Detect all loops in a cfg.
-- The assumption is that two backedges in a cfg will have distinct destination blocks.
-- If this assumption does not hold, then return the empty list.
cfgLoops :: CFG ext s init ret -> [LoopInfo s]
cfgLoops cfg
  | distinct  = lis
  | otherwise = []
  where
    (nm, gr) = blocksGraph (cfgBlocks cfg)
    root     = toNode (blockID (cfgEntryBlock cfg))
    ls       = loops root gr
    lis      = mkLoopInfo <$> ls
    distinct = length (nub (liHeader <$> lis)) == length lis

    mkLoopInfo ((footer, header, _), members) =
      LoopInfo
      { liFooter     = toBlockID footer
      , liHeader     = toBlockID header
      , liMembers    = Set.map toBlockID members
      , liEarlyExits = [ (toBlockID i, toBlockID j) | (i,j) <- exits members, i /= footer ]
      , liFooterIn   = [ (toBlockID j, toBlockID footer) | j <- G.pre gr footer ]
      , liDominators = maybe [] (fmap toBlockID) $ lookup header (G.dom gr root)
      }

    toBlockID n = nm Map.! n

    exits bs = [ (i, j) | i <- Set.toList bs, j <- G.suc gr i, j `Set.notMember` bs ]

-- | Is li1 nested in li2
isNested :: LoopInfo s -> LoopInfo s -> Ordering
isNested li1 li2
  | liHeader li2 `elem` liDominators li1 = LT
  | otherwise = EQ

-- | Return all loops in @g@, which are edges from a node in g to a
-- dominator of that node.
loops :: G.Node -- ^ entry node
      -> G.UGr -- ^ the graph
      -> [(G.LEdge (), Set G.Node)]
loops root g = [ (e, loopMembers g dominators header) | e@(_,header,_) <- edges ]
  where
    edges = loop dominators g =<< G.nodes g
    dominators = Map.fromList $ G.dom g root

-- | Return any edges from @n@ to a dominator of @n@, @n'@. The edge
-- @n@ to @n'@ is a loop.
loop :: Map.Map G.Node [G.Node] -- ^ Dominators
     -> G.UGr -- ^ The graph itself
     -> G.Node -- ^ The root node to inspect for backedges
     -> [G.LEdge ()] -- ^ A loop is an edge to a dominator
loop domMap g n =
  -- A back edge (loop) is an edge from n -> n' where n' dominates n
  [ (n, n', ()) | n' <- G.suc g n, n /= n', n' `elem` Map.findWithDefault [] n domMap ]

-- | The members of a loop are just those nodes dominated by the
-- header that can reach the header again
loopMembers :: G.UGr -> Map.Map G.Node [G.Node] -> G.Node -> Set G.Node
loopMembers g doms header =
  Set.fromList members
  where
    fromHeader = G.reachable header g
    members    = [ x | x <- fromHeader, header `elem` G.reachable x g, headerDominates x ]
    headerDominates n
      | Just ds <- Map.lookup n doms
      = header `elem` ds
      | otherwise
      = False

-- | View a blockID as a node
toNode :: BlockID s -> G.Node
toNode i =
  case i of
    LabelID l  -> fromIntegral $ indexValue $ labelId l
    LambdaID l -> fromIntegral $ indexValue $ lambdaId l

-- | Compute the successor nodes of this block
blockSuccessors :: Block ext s ret -> [G.Node]
blockSuccessors b =
  maybe [] (map toNode) $ termNextLabels (pos_val (blockTerm b))

-- | Returns the edges from this block to its successors
blockEdges :: Block ext s ret -> [G.LEdge ()]
blockEdges b =
  mkEdge (toNode (blockID b)) <$> blockSuccessors b
  where
    mkEdge x y = (x, y, ())

blocksGraph :: [Block ext s ret] -> (Map.Map G.Node (BlockID s), G.UGr)
blocksGraph blocks = (m, G.mkGraph ((,()) <$> nodes) edges)
  where
    nodes = toNode . blockID <$> blocks
    m     = Map.fromList (zip nodes (blockID <$> blocks))
    edges = blockEdges       =<< blocks

-----------------------------------------
-- | Undefined Value Fixup Transformation
-----------------------------------------

-- | A PartialValue of type @t@ closes over a register of type @Maybe t@
type ValueToPartialMap s  = MapF.MapF (Value s) (PartialValue s)
newtype PartialValue s tp = PartialValue { getPartial :: Reg s (MaybeType tp) }

type AtomSubst s = MapF.MapF (Atom s :: CrucibleType -> Type)
                              (Atom s :: CrucibleType -> Type)
type AtomPair s  = MapF.Pair (Atom s :: CrucibleType -> Type)
                              (Atom s :: CrucibleType -> Type)

-- | Undefined Value Fixup pass
-- The merge-block insertion process introduces infeasible paths along which
-- some registers may be undefined: this will later be interpreted as a block
-- input in the SSA transformation. To avoid this, we introduce a pass to
-- replace registers/atoms that may be undefined along some path with a partial
-- register (i.e. of type Maybe t).
--
-- Assuming that the input CFG has no paths along which a value is
-- read before being written, the paths along which the reference is
-- read but never written are a subset of the infeasible paths.
lowerUndefPass :: (Monad m, TraverseExt ext)
               => NonceGenerator m s
               -> Label s
               -> CFG ext s init ret
               -> m (CFG ext s init ret)
lowerUndefPass ng rootLabel cfg =
  do (pvals, refInits) <- mkPartialRegMap ng cfg

     let root' = mkBlock (blockID root) (blockExtraInputs root) (blockStmts root <> refInits) (blockTerm root)
     let lower blk
           | blockID blk == LabelID rootLabel
           = return root'
           | otherwise
           = lowerBlock ng pvals blk

     blks' <- mapM lower (cfgBlocks cfg)
     let cfg' = cfg { cfgBlocks = blks' }
     return cfg'
  where
    root = fromMaybe err $ findBlock cfg rootLabel
    err  = panic "EarlyMergeLoops.lowerUndefPass"
                 [ "Root block not found in cfg" ]

-- | Fixup the reads and writes in a block. This means, for each value in the domain
-- of the @ValueToPartialMap@ argument,
-- 1. If the value is read, then find the associated partial value register and read that instead
-- 2. Dually, if the value is written, then write that value to the associated partial value register.
lowerBlock :: forall s m ext ret
            . (Monad m, TraverseExt ext)
           => NonceGenerator m s
           -> ValueToPartialMap s
           -> Block ext s ret
           -> m (Block ext s ret)
lowerBlock ng pvals blk =
  do -- If this is a lambda block, treat the associated atom as a write to that atom.
     initInputs               <- lowerBlockIDValues ng pvals blk
     -- Dually, any values passed to a successor should be treated as reads
     (preOutput, loweredTerm) <- lowerTermStmtValues ng pvals blk
     -- Fix the reads and writes in the body of this block
     lowered <- concatMapSeqM (lowerReads >=> concatMapSeqM lowerWrites) (blockStmts blk)
     return $ mkBlock (blockID blk) (blockExtraInputs blk) (initInputs <> lowered <> preOutput) loweredTerm
  where
    lowerWrites = lowerValueWrites ng pvals
    lowerReads  = lowerValueReads ng pvals

    concatMapSeqM :: Monad m => (a -> m (Seq b)) -> Seq a -> m (Seq b)
    concatMapSeqM f seq0 =
      Fold.foldrM (\s ss -> f s <**> pure (<> ss)) mempty seq0

-- | The atom in a lambda ID is essentially an 'atom definition', so
-- we need to check if this lambda's atom needs to be 'lowered'.
lowerBlockIDValues :: forall s m ext ret
                     . (Monad m, TraverseExt ext)
                    => NonceGenerator m s
                    -> ValueToPartialMap s
                    -> Block ext s ret
                    -> m (Seq (Posd (Stmt ext s)))
lowerBlockIDValues ng pvals blk = 
  case blockID blk of
    LambdaID (lambdaAtom -> a)
      | Just (getPartial -> pr) <- MapF.lookup (AtomValue a) pvals ->
        do pa <- freshAtom ng (atomPosition a) (MaybeRepr (typeOfAtom a))
           let setPa = Posd (atomPosition a) (DefineAtom pa (EvalApp (JustValue (typeOfAtom a) a)))
               setPr = Posd (atomPosition a) (SetReg pr pa)
           return $ Seq.fromList [ setPa, setPr ]
      | otherwise ->
        -- Not in our list of values to lower
        return mempty
    LabelID {} ->
      -- No atoms defined
      return mempty

-- | Jumping to a block with a value a la Output is akin to
-- 'reading' the atom, so if we already lifted the original value
-- of type T to a (Maybe T), we need to convert it back to a T
-- here.
lowerTermStmtValues :: forall s m ext ret
                     . (Monad m, TraverseExt ext)
                    => NonceGenerator m s
                    -> ValueToPartialMap s
                    -> Block ext s ret
                    -> m (Seq (Posd (Stmt ext s)), Posd (TermStmt s ret))
lowerTermStmtValues ng pvals blk =
  case pos_val (blockTerm blk) of
    Output ll a          -> withLowered a $ Output ll
    MaybeBranch t a ll l -> withLowered a $ \a' -> MaybeBranch t a' ll l
    Return a             -> withLowered a $ Return
    TailCall f c a       -> withLowered f $ \f' -> TailCall f' c a
    ErrorStmt msg        -> withLowered msg $ ErrorStmt
    VariantElim c a ls   -> withLowered a $ \a' -> VariantElim c a' ls
    -- No atoms are output/read
    Jump {}              -> return (mempty, blockTerm blk)
    Br {}                -> return (mempty, blockTerm blk)

  where
    termPos = pos (blockTerm blk)
    withLowered :: forall (tp :: CrucibleType) ty.
                   Atom s tp
                -> (Atom s tp -> TermStmt s ty) -> m (Seq (Posd (Stmt ext s)), Posd (TermStmt s ty))
    withLowered a k =
      do (setAtom, a') <- lowerAtomRead termPos a
         return (setAtom, Posd termPos (k a'))

    lowerAtomRead :: forall (tp :: CrucibleType). Position -> Atom s tp -> m (Seq (Posd (Stmt ext s)), Atom s tp)
    lowerAtomRead p a =
        do (sub, setValue) <- lowerAtom ng pvals (Some a)
           let a' = apSubst (atomSubst sub) a
           return $ (Posd p <$> Seq.fromList setValue, a')
      

-- | Replace each write of a possibly-undef atom/register to a write of the
-- associated partial register by injecting it into a value of Maybe type.
lowerValueWrites :: forall m ext s. (Monad m, TraverseExt ext)
                 => NonceGenerator m s
                 -> ValueToPartialMap s
                 -> Posd (Stmt ext s)
                 -> m (Seq (Posd (Stmt ext s)))
lowerValueWrites ng pvals st =
  case pos_val st of
   -- Replace r := a (morally) with pr := Just a,
   -- where pr is the partial register associated with r
   SetReg r a
     | Just (getPartial -> pr) <- MapF.lookup (RegValue r) pvals ->
         setMaybeAtom pr a (typeOfAtom a)
     | otherwise -> orig

   -- Given a := v, append pr := Just(a) where pr is the
   -- partial register associated with a
   DefineAtom a _
     | Just (getPartial -> pr) <- MapF.lookup (AtomValue a) pvals ->
       do setA' <- setMaybeAtom pr a (typeOfAtom a)
          return (st Seq.<| setA')
     | otherwise -> orig

   -- No registers set or atoms defined:
   WriteGlobal {} -> orig
   WriteRef {}    -> orig
   DropRef {}     -> orig
   Print {}       -> orig
   Assert {}      -> orig
   Assume {}      -> orig
   Breakpoint {}  -> orig

  where
    orig = pure (Seq.fromList [st])

    -- Construct (pa := Just a; pr := pa) where pa is fresh
    setMaybeAtom :: forall (tp :: CrucibleType).
                    Reg s (MaybeType tp) -> Atom s tp -> TypeRepr tp -> m (Seq (Posd (Stmt ext s)))
    setMaybeAtom pr a ty =
       do pa <- freshAtom ng (atomPosition a) (MaybeRepr ty)
          let setPa = Posd (pos st) (DefineAtom pa (EvalApp (JustValue ty a)))
              setPr = Posd (pos st) (SetReg pr pa)
          return $ Seq.fromList [ setPa, setPr ]

-- | Replace each read of a lowered atom/register to a read of the
-- associated register + projection from Maybe
lowerValueReads :: forall m ext s
                . (Monad m, TraverseExt ext)
                => NonceGenerator m s
                -> ValueToPartialMap s
                -> Posd (Stmt ext s)
                -> m (Seq (Posd (Stmt ext s)))
lowerValueReads ng pvals st =
  case pos_val st of
    -- RegReads have only one form
    DefineAtom a (ReadReg r)
     | Just (getPartial -> pr) <- MapF.lookup (RegValue r) pvals ->
       lowerRegRead ng (pos st) a pr
    -- For everything else, we need to check if any of the
    -- referenced atoms need to be lowered
    DefineAtom {}  -> lowerAtomReads ng pvals atomsToLower st
    SetReg {}      -> lowerAtomReads ng pvals atomsToLower st
    WriteGlobal {} -> lowerAtomReads ng pvals atomsToLower st
    WriteRef {}    -> lowerAtomReads ng pvals atomsToLower st
    DropRef {}     -> lowerAtomReads ng pvals atomsToLower st
    Print {}       -> lowerAtomReads ng pvals atomsToLower st
    Assert {}      -> lowerAtomReads ng pvals atomsToLower st
    Assume {}      -> lowerAtomReads ng pvals atomsToLower st
    Breakpoint {}  -> lowerAtomReads ng pvals atomsToLower st
  where
    atomsToLower :: [Some (Atom s)]
    atomsToLower = Set.toList (foldStmtInputs addIfLowered (pos_val st) mempty)

    addIfLowered :: forall tp. Value s tp -> Set.Set (Some (Atom s)) -> Set.Set (Some (Atom s))
    addIfLowered v@(AtomValue a) s
      | MapF.member v pvals = Set.insert (Some a) s
    addIfLowered _ s = s

-- | Replace each read of a lowered atom to a read of the
-- associated register + projection from Maybe
lowerAtomReads :: forall m ext s.
                  (Monad m, TraverseExt ext)
               => NonceGenerator m s
               -> ValueToPartialMap s
               -> [Some (Atom s)]
               -> Posd (Stmt ext s)
               -> m (Seq (Posd (Stmt ext s)))
lowerAtomReads ng pvals atomsToLower st =
  do (substs, readRegs) <- unzip <$> mapM (lowerAtom ng pvals) atomsToLower
     let substMap        = atomSubst (concat substs)
     st'                <- mapStmtAtom (return . apSubst substMap) (pos_val st)
     let stmts           = Seq.fromList (Posd (pos st) <$> (concat readRegs ++ [st']))
     return stmts

apSubst :: AtomSubst s -> (forall (tp :: CrucibleType). Atom s tp -> Atom s tp)
apSubst sub n = fromMaybe n $ MapF.lookup n sub

atomSubst :: [MapF.Pair (Atom s :: CrucibleType -> Type) (Atom s)] -> AtomSubst s
atomSubst substs = MapF.fromList substs

-- | Given an atom @a@ of type @t@ whose definition we've already
-- replaced with a register @r@ of type @Maybe t@, produce the
-- statements to
-- 1. read @r@ into a fresh @a'@
-- 2. set fresh @a''@ to @fromJust a'@
-- returns a mapping from @a@ to @a'@ and the above
lowerAtom :: forall m s ext.
             Monad m
          => NonceGenerator m s
          -> ValueToPartialMap s
          -> Some (Atom s)
          -> m ([AtomPair s], [Stmt ext s])
lowerAtom ng pvals (Some a)
  | Just (getPartial -> r) <- MapF.lookup (AtomValue a) pvals =
      do a'  <- substAtom (const (freshNonce ng)) a
         str <- freshAtom ng (atomPosition a) knownRepr
         a'' <- freshAtom ng (atomPosition a) (MaybeRepr (typeOfAtom a))
         let defs = [ DefineAtom a'' (ReadReg r)
                    , DefineAtom str (EvalApp (StringLit ("Lower Atom Pass: " <> fromString (show (atomId a)))))
                    , DefineAtom a'  (EvalApp (FromJustValue (typeOfAtom a) a'' str))
                    ]
         return ([MapF.Pair a a'], defs)
  | otherwise =
      return ([], [])

-- | @lowerRegRead ng pos a pr@ constructs @a' := pr; a := fromJust a'@.
lowerRegRead :: Monad m
             => NonceGenerator m s
             -> Position
             -- ^ The position we should use for the new statements
             -> Atom s tp
             -- ^ The atom we're defininig
             -> Reg s (MaybeType tp)
             -- ^ The partial register to read from
             -> m (Seq (Posd (Stmt ext s)))
lowerRegRead ng p a pr =
  do a' <- freshAtom ng (atomPosition a) (MaybeRepr (typeOfAtom a))
     str <- freshAtom ng (atomPosition a) knownRepr
     -- insert a new atom to read the reg, then replace with FromJustVal
     let stmts = [ DefineAtom str (EvalApp (StringLit "Lower Register Pass"))
                 , DefineAtom a'  (ReadReg pr)
                 , DefineAtom a   (EvalApp (FromJustValue (typeOfAtom a) a' str))
                 ]
     return $ Seq.fromList (Posd p <$> stmts)
  
-- | Traverse all dfs paths, avoiding backedges, to find values that may be read but not written.
-- Returns:
--  1. a mapping from values of type @t@ to corresponding registers of type @Maybe t@ 
--  2. statements to initialize the registers mentioned in said mapping.
mkPartialRegMap :: (Monad m, TraverseExt ext)
                => NonceGenerator m s
                -> CFG ext s init ret
                -> m ((ValueToPartialMap s, Seq (Posd (Stmt ext s))))
mkPartialRegMap ng cfg =
  traverseCFG gatherPvals (MapF.empty, mempty) (blockExtraInputs entry) entry cfg
  where
    entry = cfgEntryBlock cfg

    gatherPvals pvals env blk =
      do refs'   <- Fold.foldlM addPval pvals (blockUndefVals env blk)
         return (refs', env <> blockAssignedValues blk)
  
    addPval (pvals, inits) val@(Some v)  
      | Just _ <- MapF.lookup v pvals
      = return (pvals, inits)
      | otherwise
      = do (MapF.Pair vundef pval, is) <- makePartialReg ng val
           return (MapF.insert vundef pval pvals, inits <> is)

-- | Given a value @v@ of type @t@, create a new register @r@ of type
--   @Maybe t@.  This function returns 1. the mapping from @v@ to such
--   an @r@, as well as the @stmts@ that will initialize @r@ to
--   @Nothing : Maybe t@.
makePartialReg :: Monad m
               => NonceGenerator m s
               -> Some (Value s)
               -> m (MapF.Pair (Value s) (PartialValue s), Seq (Posd (Stmt ext s)))
makePartialReg ng (Some val) =
  do a <- freshAtom ng p (MaybeRepr ty)
     r <- freshReg ng p (MaybeRepr ty)
     let v = EvalApp (NothingValue ty)
         s = DefineAtom a v
         sr = SetReg r a
         inits = Posd p <$> [s, sr]
     return (MapF.Pair val (PartialValue r), Seq.fromList inits)
  where
    (ty, p) = case val of
                RegValue reg   -> (typeOfReg reg, regPosition reg)
                AtomValue atom -> (typeOfAtom atom, atomPosition atom)

-------------------
-- | Merging Paths
-------------------

-- | This is a record used to construct/manage the variant type that the router block
-- will use to switch on. The important piece is the map that relates blockIDs to an index
-- into the variant type's ctx -- with this map we can associate a _value_ of the variant type
-- with a given blockID
data BlockSwitchInfo s ctx = BlockSwitchInfo
                             { switchRepr :: CtxRepr ctx
                             , switchSize :: Ctx.Size ctx
                             , switchMap  :: Map.Map (BlockID s) (Ctx.Index ctx UnitType)
                             }
                           deriving Show

-- | Used as a substitution between labels.
-- Closes over the type of LambdaLabels
data BlockIDPair s where
   Labels       :: Label s -> Label s -> BlockIDPair s
   LambdaLabels :: LambdaLabel s tp -> LambdaLabel s tp -> BlockIDPair s

instance Show (BlockIDPair s) where
  show (Labels l1 l2)       = show l1 ++ " => " ++ show l2
  show (LambdaLabels l1 l2) = show l1 ++ " =>{lambda} " ++ show l2

-- | This is the main pass that attempts to optimize early exits in
-- the loops in a CFG.  In particular, this transformation ensures that the
-- postdominator of the loop header is a member of the loop.
--
-- Given a natural loop, its members are a set of blocks @bs@. Let the exit edges be 
-- the edges from some block @b@ in @bs@ to either the loop header or a block not in @bs@.
--
-- Let (i, j) be such an exit edge.This transofrmation inserts a new
-- block @r@ such that in the transformed cfg there is an edge (i, r)
-- and an edge (r, j). Moreover, the transformation ensures that [i,
-- r, j'] is not feasible for j' != j. This works by setting a
-- "destination" register @d := j@ in each block i for each exit edge
-- (i, j), and switching on the value of @d@ in the block @r@.
earlyMergeLoops :: ( TraverseExt ext, Monad m, Show (CFG ext s init ret) )
                => NonceGenerator m s
                -> CFG ext s init ret
                -> m (CFG ext s init ret)
earlyMergeLoops ng cfg0 =
  do cfg' <- earlyMergeLoops' ng mempty (cfgLoops cfg0) cfg0
     lowerUndefPass ng (cfgEntryLabel cfg') cfg'
      
-- Merge a loop from loops in cfg.  The loops parameter is passed
-- to earlyMergeLoops (rather than calculated directly from cfg)
-- so that we can use see how the set of loops changes from
-- iteration to iteration: in particular, we want to make sure
-- that the number of loops does not increases
earlyMergeLoops' :: ( TraverseExt ext, Monad m, Show (CFG ext s init ret) )
                 => NonceGenerator m s
                 -> Set (BlockID s)
                 -> [LoopInfo s]
                 -> CFG ext s init ret
                 -> m (CFG ext s init ret)
earlyMergeLoops' ng seen ls cfg
  | Just l <- nextLoop
  = do cfg' <- earlyMergeLoop ng cfg l

       -- Check if we should proceed: in particular, if the new CFG
       -- renamed or produced new loops, then we need to bail as we
       -- might not terminate in that case.
       let ls'      = cfgLoops cfg'
       let seen'    = (Set.insert (liHeader l) seen)
       let nextStep = candidates seen' ls'
       -- The termination transition invariant is that 
       when (length thisStep <= length nextStep) $
         panic "EarlyMergeLoops.earlyMergeLoops'"
               ["Non-decreasing number of loops in earlyMegeLoops'"]
       
       earlyMergeLoops' ng seen' ls' cfg'
  | otherwise
  = return cfg
  where
    thisStep =
      candidates seen ls

    nextLoop 
      | null thisStep  = Nothing
      | otherwise = Just (minimumBy isNested thisStep)

    candidates seenBlocks lis =
      filter (unseen seenBlocks) lis


    unseen s li = liHeader li `Set.notMember` s

-- | Apply the transformation described in @earlyMergeLoops@ to a single loop.
earlyMergeLoop :: ( TraverseExt ext, Monad m )
               => NonceGenerator m s
               -> CFG ext s init ret
               -> LoopInfo s
               -> m (CFG ext s init ret)
earlyMergeLoop ng cfg li
  | not (null exits) =
      do bmap' <- funnelPaths ng bmap exits
         return cfg { cfgBlocks = Map.elems bmap' }
  | otherwise =
      return cfg
  where
    exits         = filterErrors (liEarlyExits li ++ liFooterIn li)
    filterErrors  = filter (not . errorPath bmap)
    bmap          = blockMap cfg

-- | Given a set of edges (i, j) in E,
-- Create a single block that merges all paths before
-- continuing on to the respective j's
--
-- Create a unique block F and multiple blocks  i' j' such
-- that i -> i' -> F -> j' -> j.
-- This function defines a register 'r : () + () + ... + ()'
-- Each 'j' corresponds to one of these tags, so each
-- i' sets r to indicate which 'j' to jump to, and F
-- switches on r. Each j' is a lambda block that jumps to the
-- original j.
--
-- E.g. given i0 -> j0, i1 -> j0, i2 -> j1,
-- then i0' =  r := inj(0, ()); jump F
--      i1' =  r := inj(0, ()); jump F
--      i2' =  r := inj(1, ()); jump F
--      F   =  switch r { 0: j0', 1: j1' }
--      j0' = jump j0
--      j1' = jump j1
funnelPaths :: (TraverseExt ext, Monad m)
             => NonceGenerator m s
             -> Map.Map (BlockID s) (Block ext s ret)
             -> [(BlockID s, BlockID s)]
             -> m (Map.Map (BlockID s) (Block ext s ret))
funnelPaths ng bmap paths =
  case mkBlockSwitchInfo outBlocks of
    Some p ->
      do (rename, newBlocks) <- routePaths ng p outBlocks
         let renamedMap       = Fold.foldl' (updateBlock rename) bmap (fst <$> paths)
             newMap           = Fold.foldl' addNewBlock renamedMap newBlocks

         return newMap
  where
    outBlocks = nub (snd <$> paths)

    sub renaming stmt = Fold.foldl' runRename stmt renaming
    addNewBlock m b        = Map.insert (blockID b) b m
    updateBlock ren m bid  = Map.update (doUpdate ren) bid m
    doUpdate renaming blk  = Just $ mkBlock (blockID blk) (blockExtraInputs blk)
                                      (blockStmts blk) (sub renaming <$> blockTerm blk)

    -- Apply the label substitution in 'renaming'
    -- to the term stmt in tgt (no-op on stmts without labels)
    runRename tgt renaming =
      case (tgt, renaming) of
        (Jump t, Labels from to)
          | t == from -> Jump to
          | otherwise -> tgt
        (Jump _, LambdaLabels {}) -> tgt

        (Br p t1 t2, Labels from to)
          | t1 == from -> Br p to t2
          | t2 == from -> Br p t1 to
          | otherwise  -> tgt
        (Br {}, LambdaLabels {}) -> tgt

        (Output ll a, LambdaLabels from to)
          | Just Refl <- testEquality ll from -> Output to a
          | otherwise -> tgt
        (Output {}, Labels {}) -> tgt

        (VariantElim ctx a assgn, r@(LambdaLabels {})) ->
          VariantElim ctx a (fmapFC (renameLabel r) assgn)
        (VariantElim {}, Labels {}) -> tgt

        (MaybeBranch t a l1 l2, Labels from to)
          | l2 == from -> MaybeBranch t a l1 to
          | otherwise  -> tgt
        (MaybeBranch t a l1 l2, LambdaLabels from to)
          | Just Refl <- testEquality l1 from ->
            MaybeBranch t a to l2
          | otherwise  -> tgt

        {- No labels to rename -}
        (Return _, _)     -> tgt
        (TailCall {}, _)  -> tgt
        (ErrorStmt {}, _) -> tgt

    renameLabel (LambdaLabels from to) ll
      | Just Refl <- testEquality from ll = to
    renameLabel _ l = l


-- | Given a list of blocks, build a record of a variant type such that each injection
-- is associated with a single block.
mkBlockSwitchInfo :: [BlockID s] -> Some (BlockSwitchInfo s)
mkBlockSwitchInfo bs =
  case bs of
    []      ->
      Some (BlockSwitchInfo Ctx.empty Ctx.zeroSize mempty)
    (b:bs') ->
      case mkBlockSwitchInfo bs' of
        Some (BlockSwitchInfo ctxRepr sz indices) ->
          Some $ BlockSwitchInfo { switchRepr = ctxRepr Ctx.:> UnitRepr
                                 , switchSize = Ctx.incSize sz
                                 , switchMap  = Map.insert b (Ctx.nextIndex sz) (Map.map Ctx.skipIndex indices)
                                 }

-- | This function does most of the work for @funnelPaths@.
-- In particular, given a list of blocks @b0...bn@,
-- it constructs a single @r@, a list @b_in0...b_inn@, and a list
-- @b_out0...@b_outn@ such that @b_in_i@ jump to @r@, and @r@
-- jumps to the corresponding @b_out_i@.
--
-- The return value is a pair of
-- (1) A list of BlockIDPairs, which should be understood as a substitution
-- on labels. This is necessary since we are introducing *new* blocks to replace
-- *old* jump targets. This substitution is used to update the old jump targets
-- (2) A list of newly created blocks
routePaths :: forall m ext ctx s ret
            . (TraverseExt ext, Monad m)
           => NonceGenerator m s
           -> BlockSwitchInfo s ctx
           -- ^ the variant info should map each element of @outs@ to an index into `ctx`
           -> [BlockID s]
           -- ^ the blocks that we want to join & then fan out from
           -> m ([BlockIDPair s], [Block ext s ret])
routePaths ng (BlockSwitchInfo ctx sz idxMap) outs =
  do -- 'bsi' has the information we need to associate each input block
     -- with an a particular *index* into a list of outputs. This means
     -- for each destination that we're routing to,
     -- we can set a register `r` to some value `v`
     -- and case split on `r` in a 'router block' to recover the intended destination

     -- this associates each such index with the new 'destination'
     mapping <- mkMapping

     -- This is the block that switches on the destination block.
     -- l is its label, r is the register that we'll case split on,
     -- and 'router' is the actual block definition containing the switch
     (l, r, router) <- routerBlock ng ctx mapping
    
     -- construct the blocks feeding into the router. for each output block,
     -- set 'r' appropriately, by fetching the value from the index map in 'vi'
     (rename, injBlocks) <- unzip <$> traverse (mkInjBlock r l) outs

     -- Finally make the destination blocks that continue to the original
     -- targets
     let outputBlocks = Map.foldrWithKey (mkLambdaBlock rename mapping) [] idxMap

     return (rename, router : outputBlocks ++ injBlocks)
  where
    bidToTerm :: BlockID s -> [BlockIDPair s] -> TermStmt s ret
    bidToTerm origOut rename =
      case (origOut, rename) of
        (LabelID ll, _ ) -> Jump ll
        (LambdaID ll, LambdaLabels l1 l2:_)
          | Just Refl <- testEquality (lambdaId ll) (lambdaId l1) -> Output l1 (lambdaAtom l2)
        (_, _:rest) -> bidToTerm origOut rest
        _ ->
          error "Output blocks mismatched in routePaths"
  
    mkMapping = Ctx.generateM sz $ \idx ->
      do n <- freshNonce ng
         a <- freshAtom ng internal (ctx Ctx.! idx)
         return (LambdaLabel n a)

    mkInjBlock reg rlabel j =
      do let idx = idxMap Map.! j
         (rename, injBlock) <- routerEntryBlock ng ctx reg j idx rlabel
         return (rename, injBlock)

    mkLambdaBlock rename mapping blkId blkIdx blks =
      mkBlock (LambdaID (mapping Ctx.! blkIdx))
              mempty
              mempty
              (Posd internal (bidToTerm blkId rename)) : blks

-- | Create a block that switches on a register. This does the work of
-- allocating the new label and discriminant register
routerBlock :: ( TraverseExt ext, Monad m )
            => NonceGenerator m s
            -> CtxRepr routing
            -> Ctx.Assignment (LambdaLabel s) routing
            -> m (Label s, Reg s (VariantType routing), Block ext s ret)
routerBlock ng ctx mapping =
  do l        <- Label <$> freshNonce ng
     destReg  <- freshReg ng internal (VariantRepr ctx)
     readDest <- freshAtom ng internal (VariantRepr ctx)
     let elim    = VariantElim ctx readDest mapping
         readVar = Posd internal (DefineAtom readDest (ReadReg destReg))
         funnel  = mkBlock (LabelID l) mempty (pure readVar) (Posd internal elim)
     return (l, destReg, funnel)

-- | This creates a block that to be substituted for @origID@.
-- This block will set the given register to the injection
-- given by @destIdx@, and then jump to a router block (described in @routePaths@)
-- that will `switch` on this register. 
routerEntryBlock :: ( TraverseExt ext, Monad m )
                 => NonceGenerator m s
                 -> CtxRepr routing
                 -> Reg s (VariantType routing)
                 -- ^ The register we will switch on
                 -> BlockID s
                 -- ^ The ID of the block we're substituting for
                 -> Ctx.Index routing UnitType
                 -- ^ Which injection in the variant type to use
                 -> Label s
                 -- ^ The label of the router block
                 -> m (BlockIDPair s, Block ext s ret)
routerEntryBlock ng ctx r origID destIdx routerLabel =
  do aUnit <- freshAtom ng internal UnitRepr
     aInj  <- freshAtom ng internal (VariantRepr ctx)
   
     (l, rename) <-
       case origID of
         LabelID l1 ->
           do l2 <- substLabel (\_ -> freshNonce ng) l1
              return (LabelID l2, Labels l1 l2)
         LambdaID l1@(LambdaLabel _ a) ->
           do l' <- freshNonce ng
              a' <- freshNonce ng
              -- knot-tying in the LambdaLabel type results in an infinite loop
              let l2 = LambdaLabel l' a { atomId = a' }
              return (LambdaID l2, LambdaLabels l1 l2)

     let defUnit = Posd internal (DefineAtom aUnit (EvalApp EmptyApp))
         defVar  = Posd internal (DefineAtom aInj (EvalApp (InjectVariant ctx destIdx aUnit)))
         setReg  = Posd internal (SetReg r aInj)
         stmts   = Seq.fromList [defUnit, defVar, setReg]
         blk     = mkBlock l mempty stmts (Posd internal (Jump routerLabel))

     return (rename, blk)

errorPath :: Map.Map (BlockID s) (Block ext s ret) -> CFGEdge s -> Bool
errorPath bmap (_, bid) =
  case Map.lookup bid bmap of
    Just (blockTerm -> Posd _ (ErrorStmt _)) -> True
    _ -> False

-- | Generally useful helpers

freshAtom :: Monad m => NonceGenerator m s -> Position -> TypeRepr tp -> m (Atom s tp)
freshAtom ng p tp =
  do i <- freshNonce ng
     return $ Atom { atomPosition = p
                   , atomId = i
                   , atomSource = Assigned
                   , typeOfAtom = tp
                   }

freshReg :: Monad m => NonceGenerator m s -> Position -> TypeRepr tp -> m (Reg s tp)
freshReg ng p tp =
  do i <- freshNonce ng
     return $ Reg { regPosition = p
                  , regId = i
                  , typeOfReg = tp
                  }
      
findBlock :: CFG ext s init ret -> Label s -> Maybe (Block ext s ret)
findBlock g l =
  Fold.find (\b -> blockID b == LabelID l) (cfgBlocks g)
  
blockUndefVals :: ValueSet s -> Block ext s ret -> ValueSet s
blockUndefVals def blk = blockKnownInputs blk Set.\\ def

blockMap :: CFG ext s init ret -> Map.Map (BlockID s) (Block ext s ret)
blockMap cfg = Map.fromList [ (blockID b, b) | b <- cfgBlocks cfg ]

internal :: Position
internal = InternalPos