content: -- Copyright 2021 Google LLC
  --
  -- Use of this source code is governed by a BSD-style
  -- license that can be found in the LICENSE file or at
  -- https://developers.google.com/open-source/licenses/bsd
  
  {-# LANGUAGE PartialTypeSignatures #-}
  {-# OPTIONS_GHC -Wno-partial-type-signatures #-}
  {-# OPTIONS_GHC -Wno-orphans #-}
  
  module Imp
    ( toImpFunction, repValAtom, impFunType, getIType, abstractLinktimeObjects
    , repValFromFlatList, addImpTracing
    -- These are just for the benefit of serialization/printing. otherwise we wouldn't need them
    , BufferType (..), IdxNest, IndexStructure, IExprInterpretation (..), typeToTree
    , computeOffset, getIExprInterpretation
    , isSingletonType, singletonTypeVal
    ) where
  
  import Prelude hiding ((.), id)
  import Data.Functor
  import Data.Foldable (toList)
  import Data.Maybe (fromJust, isJust)
  import Data.Text.Prettyprint.Doc
  import Control.Category
  import Control.Monad.Identity
  import Control.Monad.Writer.Strict
  import Control.Monad.State.Strict hiding (State)
  import qualified Control.Monad.State.Strict as MTL
  
  import Algebra
  import Builder
  import CheapReduction
  import CheckType (CheckableE (..))
  import Core
  import Err
  import IRVariants
  import MTL1
  import Name
  import PPrint
  import Subst
  import QueryType
  import Types.Core
  import Types.Imp
  import Types.Primitives
  import Types.Top
  import Util (forMFilter, Tree (..), zipTrees, enumerate)
  
  toImpFunction :: EnvReader m => CallingConvention -> STopLam n -> m n (ImpFunction n)
  toImpFunction cc (TopLam True destTy lam) = do
    LamExpr bsAndRefB body <- return lam
    PairB bs destB <- case popNest bsAndRefB of
      Just bsAndRefB' -> return bsAndRefB'
      Nothing -> error "expected a trailing reference binder"
    let ty = piTypeWithoutDest destTy
    impArgTys <- getNaryLamImpArgTypesWithCC cc ty
    liftImpM $ buildImpFunction cc (zip (repeat noHint) impArgTys) \vs -> do
      case cc of
        EntryFunCC _ -> do
          argAtoms <- interpretImpArgs (sink $ EmptyAbs bs) vs
          extendSubst (bs @@> (SubstVal <$> argAtoms)) do
            dest <- case binderType destB of
              RefTy _ ansTy -> allocDestUnmanaged =<< substM ansTy
              _ -> error "Expected a reference type for body destination"
            extendSubst (destB @> SubstVal (destToAtom dest)) do
              void $ translateExpr body
            resultAtom <- loadAtom dest
            repValToList <$> atomToRepVal resultAtom
        _ -> do
          (argAtoms, resultDest) <- interpretImpArgsWithCC cc (sink ty) vs
          extendSubst (bs @@> (SubstVal <$> argAtoms)) do
            extendSubst (destB @> SubstVal (destToAtom (sink resultDest))) do
              void $ translateExpr body
              return []
  toImpFunction _ (TopLam False _ _) = error "expected a lambda in destination-passing form"
  {-# SCC toImpFunction #-}
  
  getNaryLamImpArgTypesWithCC :: EnvReader m
    => CallingConvention -> PiType SimpIR n -> m n [BaseType]
  getNaryLamImpArgTypesWithCC XLACC _ = return [i8pp, i8pp]
    where i8pp = PtrType (CPU, PtrType (CPU, Scalar Word8Type))
  getNaryLamImpArgTypesWithCC (EntryFunCC _) t =
    concat . fst <$> getNaryLamImpArgTypes t
  getNaryLamImpArgTypesWithCC _ t = do
    (argTyss, destTys) <- getNaryLamImpArgTypes t
    return $ concat argTyss ++ destTys
  
  getImpFunType :: EnvReader m
    => CallingConvention -> PiType SimpIR n -> m n IFunType
  getImpFunType StandardCC piTy = do
    argTys <- getNaryLamImpArgTypesWithCC StandardCC piTy
    return $ IFunType StandardCC argTys []
  getImpFunType cc _ = error $ "unsupported calling convention: " ++ pprint cc
  
  interpretImpArgsWithCC :: Emits n
    => CallingConvention -> PiType SimpIR n
    -> [IExpr n] -> SubstImpM i n ([SAtom n], Dest n)
  interpretImpArgsWithCC XLACC t [outsPtr, insPtr] = do
    (argBaseTys, resultBaseTys) <- getNaryLamImpArgTypes t
    argVals <- forM (enumerate $ concat argBaseTys) \(i, ty) -> case ty of
      PtrType (_, pointeeTy) -> loadPtr insPtr i pointeeTy
      _ -> load =<< loadPtr insPtr i ty
    resultPtrs <- case resultBaseTys of
      -- outsPtr points directly to the buffer when there's one output, not to the pointer array.
      [ty] -> (:[]) <$> cast outsPtr ty
      _ ->
        forM (enumerate resultBaseTys) \(i, ty) -> case ty of
          PtrType (_, pointeeTy) -> loadPtr outsPtr i pointeeTy
          _ -> error "Destination arguments should all have pointer types"
    interpretImpArgsWithDest t (argVals ++ resultPtrs)
    where
      loadPtr base i pointeeTy = do
        ptr <- load =<< impOffset base (IIdxRepVal $ fromIntegral i)
        cast ptr (PtrType (CPU, pointeeTy))
  interpretImpArgsWithCC _ t xsAll = interpretImpArgsWithDest t xsAll
  
  getNaryLamImpArgTypes :: EnvReader m
    => PiType SimpIR n -> m n ([[BaseType]], [BaseType])
  getNaryLamImpArgTypes t = liftEnvReaderM $ go t where
    go :: PiType SimpIR n -> EnvReaderM n ([[BaseType]], [BaseType])
    go (PiType bs effTy) = case bs of
      Nest piB rest -> do
        ts <- getRepBaseTypes $ binderType piB
        refreshAbs (Abs piB (PiType rest effTy)) \_ restPi -> do
          (argTys, resultTys) <- go restPi
          return (ts:argTys, resultTys)
      Empty -> ([],) <$> getDestBaseTypes (etTy effTy)
  
  interpretImpArgsWithDest :: EnvReader m
    => PiType SimpIR n -> [IExpr n] -> m n ([SAtom n], Dest n)
  interpretImpArgsWithDest t xs = do
    (PiType bs (EffTy _ resultTy)) <- return t
    (args, xsLeft) <- _interpretImpArgs (EmptyAbs bs) xs
    resultTy' <- applySubst (bs @@> (SubstVal <$> args)) resultTy
    (destTree, xsRest) <- listToTree resultTy' xsLeft
    case xsRest of
      [] -> return ()
      _  -> error "Shouldn't have any Imp arguments left"
    return (args, Dest resultTy' destTree)
  
  interpretImpArgs :: EnvReader m
    => EmptyAbs (Nest SBinder) n -> [IExpr n] -> m n [SAtom n]
  interpretImpArgs t args = do
    (args', xsLeft) <- _interpretImpArgs t args
    case xsLeft of
      [] -> return args'
      _  -> error "Shouldn't have any Imp arguments left"
  
  _interpretImpArgs :: EnvReader m
    => EmptyAbs (Nest SBinder) n -> [IExpr n] -> m n ([SAtom n], [IExpr n])
  _interpretImpArgs t args =
    liftEnvReaderM $ runSubstReaderT idSubst $ go t args where
      go :: EmptyAbs (Nest SBinder) i -> [IExpr o]
         -> SubstReaderT AtomSubstVal EnvReaderM i o ([SAtom o], [IExpr o])
      go (Abs bs UnitE) xs = case bs of
        Nest (b:>argTy) rest -> do
          argTy' <- substM argTy
          (argTree, xsRest) <- listToTree argTy' xs
          arg <- repValAtom $ RepVal argTy' argTree
          (args', xsLeft) <- extendSubst (b @> SubstVal arg) $ go (EmptyAbs rest) xsRest
          return (arg:args', xsLeft)
        Empty -> return ([], xs)
  
  -- === ImpM monad ===
  
  type ImpBuilderEmissions = RNest ImpDecl
  
  newtype ImpDeclEmission (n::S) (l::S) = ImpDeclEmission (ImpDecl n l)
  instance ExtOutMap Env ImpDeclEmission where
    extendOutMap env (ImpDeclEmission d) = env `extendOutMap` toEnvFrag d
    {-# INLINE extendOutMap #-}
  instance ExtOutFrag ImpBuilderEmissions ImpDeclEmission where
    extendOutFrag ems (ImpDeclEmission d) = RNest ems d
    {-# INLINE extendOutFrag #-}
  
  newtype ImpM (n::S) (a:: *) =
    ImpM { runImpM' :: WriterT1 (ListE IExpr)
                         (InplaceT Env ImpBuilderEmissions HardFailM) n a }
    deriving ( Functor, Applicative, Monad, ScopeReader, Fallible, MonadFail)
  
  type SubstImpM = SubstReaderT AtomSubstVal ImpM :: S -> S -> * -> *
  
  instance ExtOutMap Env ImpBuilderEmissions where
    extendOutMap bindings emissions =
      bindings `extendOutMap` toEnvFrag emissions
  
  class (ImpBuilder2 m, SubstReader AtomSubstVal m, EnvReader2 m, EnvExtender2 m)
        => Imper (m::MonadKind2) where
  
  instance EnvReader ImpM where
    unsafeGetEnv = ImpM $ lift11 $ getOutMapInplaceT
  
  instance EnvExtender ImpM where
    refreshAbs ab cont = ImpM $ lift11 $
      refreshAbs ab \b e -> do
        (result, ptrs) <- runWriterT1 $ runImpM' $ cont b e
        case ptrs of
          ListE [] -> return result
          _ -> error "shouldn't be able to emit pointers without `Mut`"
  
  instance ImpBuilder ImpM where
    emitMultiReturnInstr instr = do
      Distinct <- getDistinct
      tys <- impInstrTypes instr
      -- The three cases below are all equivalent, but 0- and 1-return instructions
      -- are so common that it's worth spelling their cases out explicitly to enable
      -- more GHC optimizations.
      ImpM $ lift11 case tys of
        []   -> do
          extendTrivialSubInplaceT $ ImpDeclEmission $ ImpLet Empty instr
          return NoResults
        [ty] -> do
          OneResult <$> freshExtendSubInplaceT noHint \b ->
            (ImpDeclEmission $ ImpLet (Nest (IBinder b ty) Empty) instr, IVar (binderName b) ty)
        _ -> do
          Abs bs vs <- return $ newNames $ length tys
          let impBs = makeImpBinders bs tys
          let decl = ImpLet impBs instr
          ListE vs' <- extendInplaceT $ Abs (RNest REmpty decl) vs
          return $ MultiResult $ zipWith IVar vs' tys
      where
       makeImpBinders :: Nest (NameBinder ImpNameC) n l -> [IType] -> Nest IBinder n l
       makeImpBinders Empty [] = Empty
       makeImpBinders (Nest b bs) (ty:tys) = Nest (IBinder b ty) $ makeImpBinders bs tys
       makeImpBinders _ _ = error "zip error"
  
    emitDeclsImp (Abs decls result) =
      ImpM $ lift11 $ extendInplaceT (Abs (toRNest decls) result)
    {-# INLINE emitDeclsImp #-}
  
    buildScopedImp cont = ImpM $ WriterT1 \w ->
      liftM (, w) do
        Abs rdecls e <- locallyMutableInplaceT
          (do Emits <- fabricateEmitsEvidenceM
              (result, (ListE ptrs)) <- runWriterT1 $ runImpM' do
                 Distinct <- getDistinct
                 cont
              _ <- runWriterT1 $ runImpM' do
                forM ptrs \ptr -> emitStatement $ Free ptr
              return result)
          (\d e -> return $ Abs d e)
        return $ Abs (unRNest rdecls) e
  
    extendAllocsToFree ptr = ImpM $ tell $ ListE [ptr]
    {-# INLINE extendAllocsToFree #-}
  
  instance ImpBuilder m => ImpBuilder (SubstReaderT AtomSubstVal m i) where
    emitMultiReturnInstr instr = liftSubstReaderT $ emitMultiReturnInstr instr
    {-# INLINE emitMultiReturnInstr #-}
    emitDeclsImp ab = liftSubstReaderT $ emitDeclsImp ab
    {-# INLINE emitDeclsImp #-}
    buildScopedImp cont = SubstReaderT \env ->
      buildScopedImp $ runSubstReaderT (sink env) $ cont
    {-# INLINE buildScopedImp #-}
    extendAllocsToFree ptr = liftSubstReaderT $ extendAllocsToFree ptr
    {-# INLINE extendAllocsToFree #-}
  
  instance ImpBuilder m => Imper (SubstReaderT AtomSubstVal m)
  
  liftImpM :: EnvReader m => SubstImpM n n a -> m n a
  liftImpM cont = do
    env <- unsafeGetEnv
    Distinct <- getDistinct
    case runHardFail $ runInplaceT env $ runWriterT1 $
           runImpM' $ runSubstReaderT idSubst $ cont of
      (REmpty, (result, ListE [])) -> return result
      _ -> error "shouldn't be possible because of `Emits` constraint"
  
  -- === the actual pass ===
  
  translateDeclNest :: Emits o => Nest SDecl i i' -> SubstImpM i' o a -> SubstImpM i o a
  translateDeclNest decls cont = case decls of
    Empty -> cont
    Nest (Let b (DeclBinding _ expr)) rest -> do
      x <- translateExpr expr
      extendSubst (b@>SubstVal x) $ translateDeclNest rest cont
  {-# INLINE translateDeclNest #-}
  
  translateExpr :: forall i o. Emits o => SExpr i -> SubstImpM i o (SAtom o)
  translateExpr expr = confuseGHC >>= \_ -> case expr of
    Block _ (Abs decls result) -> translateDeclNest decls $ translateExpr result
    TopApp (EffTy _ resultTy') f' xs' -> do
      resultTy <- substM resultTy'
      f <- substM f'
      xs <- mapM substM xs'
      lookupTopFun f >>= \case
        DexTopFun _ (TopLam _ piTy _) _ -> emitCall piTy f $ toList xs
        FFITopFun _ _ -> do
          scalarArgs <- liftM toList $ mapM fromScalarAtom xs
          results <- impCall f scalarArgs
          restructureScalarOrPairType resultTy results
    TabApp _ f' x' -> do
      x <- substM x'
      f <- atomToRepVal =<< substM f'
      repValAtom =<< indexRepVal f x
    Atom x -> substM x
    PrimOp op -> toImpOp op
    Case e alts (EffTy _ unitResultTy) -> do
      e' <- substM e
      case unitResultTy of
        UnitTy -> return ()
        _ -> error $ "Unexpected returning Case in Imp " ++ pprint expr
      case e' of
        Con con -> do
          SumCon _ i arg <- return con
          Abs b body <- return $ alts !! i
          extendSubst (b @> SubstVal arg) $ translateExpr body
        Stuck _ _ -> do
          RepVal sumTy (Branch (tag:xss)) <- atomToRepVal e'
          ts <- caseAltsBinderTys sumTy
          tag' <- repValAtom $ RepVal TagRepTy tag
          xss' <- zipWithM (\t x -> repValAtom $ RepVal t x) ts xss
          go tag' xss'
          where
            go tag xss = do
              tag' <- fromScalarAtom tag
              emitSwitch tag' (zip xss alts) $
                \(xs, Abs b body) ->
                   extendSubst (b @> SubstVal (sink xs)) $
                     void $ translateExpr body
              return UnitVal
    TabCon _ _ _ -> error "Unexpected `TabCon` in Imp pass."
    Project _ i x -> reduceProj i =<< substM x
  
  toImpRefOp :: Emits o
    => SAtom i -> RefOp SimpIR i -> SubstImpM i o (SAtom o)
  toImpRefOp refDest' m = do
    refDest <- atomToDest =<< substM refDest'
    substM m >>= \case
      MAsk -> loadAtom refDest
      MExtend (BaseMonoid _ combine) x -> do
        xTy <- return $ getType x
        refVal <- loadAtom refDest
        liftMonoidCombine refDest xTy combine refVal x
        return UnitVal
      MPut x -> storeAtom refDest x >> return UnitVal
      MGet -> do
        Dest resultTy _ <- return refDest
        dest <- allocDest resultTy
        -- It might be more efficient to implement a specialized copy for dests
        -- than to go through a general purpose atom.
        storeAtom dest =<< loadAtom refDest
        loadAtom dest
      IndexRef _ i -> destToAtom <$> indexDest refDest i
      ProjRef  _ ~(ProjectProduct i) -> return $ destToAtom $ projectDest i refDest
    where
      liftMonoidCombine :: Emits o
        => (Dest o) -> SType o -> LamExpr SimpIR o
        -> SAtom o -> SAtom o -> SubstImpM n o ()
      liftMonoidCombine accDest accTy bc x y = do
        LamExpr (Nest (_:>baseTy) _) _ <- return bc
        alphaEq accTy baseTy >>= \case
          -- Immediately beta-reduce, beacuse Imp doesn't reduce non-table applications.
          True -> do
            BinaryLamExpr xb yb body <- return bc
            body' <- applySubst (xb @> SubstVal x <.> yb @> SubstVal y) body
            ans <- liftBuilderImp $ emit (sink body')
            storeAtom accDest ans
          False -> case accTy of
            TyCon (TabPi t) -> do
              let ixTy = tabIxType t
              n <- indexSetSizeImp ixTy
              emitLoop noHint Fwd n \i -> do
                idx <- unsafeFromOrdinalImp (sink ixTy) i
                xElt <- liftBuilderImp $ tabApp (sink x) (sink idx)
                yElt <- liftBuilderImp $ tabApp (sink y) (sink idx)
                eltTy <- instantiate (sink t) [idx]
                ithDest <- indexDest (sink accDest) idx
                liftMonoidCombine ithDest eltTy (sink bc) xElt yElt
            _ -> error $ "Base monoid type mismatch: can't lift " ++
                   pprint baseTy ++ " to " ++ pprint accTy
  
  toImpOp :: forall i o . Emits o => PrimOp SimpIR i -> SubstImpM i o (SAtom o)
  toImpOp op = case op of
    Hof hof -> toImpTypedHof hof
    RefOp refDest eff -> toImpRefOp refDest eff
    DAMOp damOp -> case damOp of
      Seq _ d ixTy' carry f -> do
        UnaryLamExpr b body <- return f
        ixTy <- substM ixTy'
        carry' <- substM carry
        n <- indexSetSizeImp ixTy
        emitLoop (getNameHint b) d n \i -> do
          idx <- unsafeFromOrdinalImp (sink ixTy) i
          void $ extendSubst (b @> SubstVal (PairVal idx (sink carry'))) $
            translateExpr body
        return carry'
      RememberDest _ d f -> do
        UnaryLamExpr b body <- return f
        d' <- substM d
        void $ extendSubst (b @> SubstVal d') $ translateExpr body
        return d'
      Place ref val -> do
        val' <- substM val
        refDest <- atomToDest =<< substM ref
        storeAtom refDest val' >> return UnitVal
      Freeze ref -> loadAtom =<< atomToDest =<< substM ref
      AllocDest ty  -> do
        d <- liftM destToAtom $ allocDest =<< substM ty
        return d
    BinOp binOp x y -> returnIExprVal =<< emitInstr =<< (IBinOp binOp <$> fsa x <*> fsa y)
    UnOp  unOp  x   -> returnIExprVal =<< emitInstr =<< (IUnOp  unOp  <$> fsa x)
    MemOp    op' -> toImpMemOp    =<< substM op'
    MiscOp   op' -> toImpMiscOp   =<< substM op'
    VectorOp op' -> toImpVectorOp =<< substM op'
    where
      fsa x = substM x >>= fromScalarAtom
      returnIExprVal x = return $ toScalarAtom x
  
  toImpVectorOp :: Emits o => VectorOp SimpIR o -> SubstImpM i o (SAtom o)
  toImpVectorOp = \case
    VectorBroadcast val vty -> do
      val' <- fromScalarAtom val
      emitInstr (IVectorBroadcast val' $ toIVectorType vty) >>= returnIExprVal
    VectorIota vty -> emitInstr (IVectorIota $ toIVectorType vty) >>= returnIExprVal
    VectorIdx tbl' i vty -> do
      -- VectorIdx requires that tbl' have a scalar element type, which is
      -- ultimately enforced by `Lower.getVectorType` barfing on non-scalars.
      tbl <- atomToRepVal tbl'
      repValAtom =<< vectorIndexRepVal tbl i vty
    VectorSubref ref i vty -> do
      refDest <- atomToDest ref
      refi <- destToAtom <$> indexDest refDest i
      refi' <- fromScalarAtom refi
      resultVal <- castPtrToVectorType refi' (toIVectorType vty)
      repValAtom $ RepVal (RefTy (Con HeapVal) vty) (Leaf resultVal)
    where
      returnIExprVal x = return $ toScalarAtom x
  
  castPtrToVectorType :: Emits n
    => IExpr n -> IVectorType -> SubstImpM i n (IExpr n)
  castPtrToVectorType ptr vty = do
    let PtrType (addrSpace, _) = getIType ptr
    cast ptr (PtrType (addrSpace, vty))
  
  toImpMiscOp :: forall i o . Emits o => MiscOp SimpIR o -> SubstImpM i o (SAtom o)
  toImpMiscOp op = case op of
    ThrowError resultTy -> do
      emitStatement IThrowError
      buildGarbageVal resultTy
    CastOp destTy x -> do
      BaseTy _  <- return $ getType x
      BaseTy bt <- return destTy
      x' <- fsa x
      returnIExprVal =<< cast x' bt
    BitcastOp destTy x -> do
      BaseTy bt <- return destTy
      returnIExprVal =<< emitInstr =<< (IBitcastOp bt <$> fsa x)
    UnsafeCoerce resultTy x -> do
      srcTy <- return $ getType x
      srcRep  <- getRepBaseTypes srcTy
      destRep <- getRepBaseTypes resultTy
      assertEq srcRep destRep $
        "representation types don't match: " ++ pprint srcRep ++ "  !=  " ++ pprint destRep
      RepVal _ tree <- atomToRepVal x
      repValAtom (RepVal resultTy tree)
    GarbageVal resultTy -> buildGarbageVal resultTy
    Select p x y -> do
      BaseTy _ <- return $ getType x
      returnIExprVal =<< emitInstr =<< (ISelect <$> fsa p <*> fsa x <*> fsa y)
    SumTag con -> case con of
      Con (SumCon _ tag _) -> return $ TagRepVal $ fromIntegral tag
      Stuck _ (RepValAtom dRepVal) -> do
        RepVal _ (Branch (tag:_)) <- return dRepVal
        return $ toAtom $ RepVal (TagRepTy :: SType o) tag
      _ -> error $ "Not a data constructor: " ++ pprint con
    ToEnum ty i -> case ty of
      TyCon (SumType cases) -> do
        i' <- fromScalarAtom i
        return $ toAtom $ RepVal ty $ Branch $ Leaf i' : map (const (Branch [])) cases
      _ -> error $ "Not an enum: " ++ pprint ty
    OutputStream -> returnIExprVal =<< emitInstr IOutputStream
    ThrowException _ -> error "shouldn't have ThrowException left" -- also, should be replaced with user-defined errors
    ShowAny _ -> error "Shouldn't have ShowAny in simplified IR"
    ShowScalar x -> do
      resultTy <- return $ getType $ PrimOp $ MiscOp op
      Dest (PairTy sizeTy tabTy) (Branch [sizeTree, tabTree@(Leaf tabPtr)]) <- allocDest resultTy
      xScalar <- fromScalarAtom x
      size <- emitInstr $ IShowScalar tabPtr xScalar
      let size' = toScalarAtom size
      storeAtom (Dest sizeTy sizeTree) size'
      tab <- loadAtom $ Dest tabTy tabTree
      return $ PairVal size' tab
    where
      fsa = fromScalarAtom
      returnIExprVal x = return $ toScalarAtom x
  
  toImpMemOp :: forall i o . Emits o => MemOp SimpIR o -> SubstImpM i o (SAtom o)
  toImpMemOp op = case op of
    IOAlloc n -> do
      n' <- fsa n
      ptr <- emitInstr $ Alloc CPU (Scalar Word8Type) n'
      returnIExprVal ptr
    IOFree ptr -> do
      ptr' <- fsa ptr
      emitStatement $ Free ptr'
      return UnitVal
    PtrOffset arr (IdxRepVal 0) -> return arr
    PtrOffset arr off -> do
      arr' <- fsa arr
      off' <- fsa off
      buf <- impOffset arr' off'
      returnIExprVal buf
    PtrLoad arr -> do
      result <- load =<< fsa arr
      returnIExprVal result
    PtrStore ptr x -> do
      ptr' <- fsa ptr
      x'   <- fsa x
      store ptr' x'
      return UnitVal
    where
      fsa = fromScalarAtom
      returnIExprVal x = return $ toScalarAtom x
  
  toImpTypedHof :: Emits o => TypedHof SimpIR i -> SubstImpM i o (SAtom o)
  toImpTypedHof (TypedHof (EffTy _ resultTy') hof) = do
    resultTy <- substM resultTy'
    case hof of
      For _ _ _ -> error $ "Unexpected `for` in Imp pass " ++ pprint hof
      While body -> do
        body' <- buildBlockImp do
          ans <- fromScalarAtom =<< translateExpr body
          return [ans]
        emitStatement $ IWhile body'
        return UnitVal
      RunReader r f -> do
        BinaryLamExpr h ref body <- return f
        r' <- substM r
        rDest <- allocDest $ getType r'
        storeAtom rDest r'
        extendSubst (h @> SubstVal (Con HeapVal) <.> ref @> SubstVal (destToAtom rDest)) $
          translateExpr body
      RunWriter d (BaseMonoid e _) f -> do
        BinaryLamExpr h ref body <- return f
        let PairTy ansTy accTy = resultTy
        (aDest, wDest) <- case d of
          Nothing -> destPairUnpack <$> allocDest resultTy
          Just d' -> do
            aDest <- allocDest ansTy
            wDest <- atomToDest =<< substM d'
            return (aDest, wDest)
        e' <- substM e
        PairE accTy' e'' <- sinkM $ PairE accTy e'
        liftMonoidEmpty wDest accTy' e''
        extendSubst (h @> SubstVal (Con HeapVal) <.> ref @> SubstVal (destToAtom wDest)) $
          translateExpr body >>= storeAtom aDest
        PairVal <$> loadAtom aDest <*> loadAtom wDest
      RunState d s f -> do
        BinaryLamExpr h ref body <- return f
        let PairTy ansTy _ = resultTy
        (aDest, sDest) <- case d of
          Nothing -> destPairUnpack <$> allocDest resultTy
          Just d' -> do
            aDest <- allocDest ansTy
            sDest <- atomToDest =<< substM d'
            return (aDest, sDest)
        storeAtom sDest =<< substM s
        extendSubst (h @> SubstVal (Con HeapVal) <.> ref @> SubstVal (destToAtom sDest)) $
          translateExpr body >>= storeAtom aDest
        PairVal <$> loadAtom aDest <*> loadAtom sDest
      RunIO body -> translateExpr body
      RunInit body -> translateExpr body
      where
        liftMonoidEmpty :: Emits n => Dest n -> SType n -> SAtom n -> SubstImpM i n ()
        liftMonoidEmpty accDest accTy x = do
          xTy <- return $ getType x
          alphaEq xTy accTy >>= \case
            True -> storeAtom accDest x
            False -> case accTy of
              TyCon (TabPi t) -> do
                let ixTy = tabIxType t
                n <- indexSetSizeImp ixTy
                emitLoop noHint Fwd n \i -> do
                  idx <- unsafeFromOrdinalImp (sink ixTy) i
                  x' <- sinkM x
                  eltTy <- instantiate (sink t) [idx]
                  ithDest <- indexDest (sink accDest) idx
                  liftMonoidEmpty ithDest eltTy x'
              _ -> error $ "Base monoid type mismatch: can't lift " ++
                    pprint xTy ++ " to " ++ pprint accTy
  
  -- === Runtime representation of values and refs ===
  
  data Dest (n::S) = Dest
       (SType n)        -- type of stored value
       (Tree (IExpr n))  -- underlying scalar values
       deriving (Show)
  
  data LeafType n where
    LeafType :: TypeCtx SimpIR n l -> BaseType -> LeafType n
  
  instance SinkableE LeafType where sinkingProofE = undefined
  
  -- Gives the relevant context with which to interpret the leaves of a type tree.
  type TypeCtx r = Nest (TypeCtxLayer r)
  
  type IndexStructure r = EmptyAbs (IdxNest r) :: E
  pattern Singleton :: IndexStructure r n
  pattern Singleton = EmptyAbs Empty
  
  type IxBinder r = PairB (LiftB (IxDict r)) (Binder r)
  type IdxNest r = Nest (IxBinder r)
  
  data TypeCtxLayer (r::IR) (n::S) (l::S) where
   TabCtx     :: IxBinder r n l -> TypeCtxLayer r n l
   DepPairCtx :: MaybeB (Binder r) n l -> TypeCtxLayer r n l
   RefCtx     ::                   TypeCtxLayer r n n
  
  instance SinkableE Dest where
    sinkingProofE = undefined
  
  -- `ScalarDesc` describes how to interpret an Imp value in terms of the nest of
  -- buffers that it points to
  data BufferElementType = UnboxedValue BaseType | BoxedBuffer BufferElementType deriving (Show)
  data BufferType n = BufferType (IndexStructure SimpIR n) BufferElementType
  data IExprInterpretation n = BufferPtr (BufferType n) | RawValue BaseType
  
  getRefBufferType :: LeafType n -> BufferType n
  getRefBufferType fullLeafTy = case splitLeadingIxs fullLeafTy of
    Abs idxs restLeafTy ->
      BufferType (EmptyAbs idxs) (getElemType restLeafTy)
  
  getIExprInterpretation :: LeafType n -> IExprInterpretation n
  getIExprInterpretation fullLeafTy = case splitLeadingIxs fullLeafTy of
    Abs idxs restLeafTy -> case idxs of
      Empty -> RawValue (elemTypeToBaseType $ getElemType restLeafTy)
      _ -> BufferPtr (BufferType (EmptyAbs idxs) (getElemType restLeafTy))
  
  getElemType :: LeafType n -> BufferElementType
  getElemType leafTy = fst $ getElemTypeAndIdxStructure leafTy
  
  getElemTypeAndIdxStructure :: LeafType n -> (BufferElementType, Maybe (IndexStructure SimpIR n))
  getElemTypeAndIdxStructure (LeafType ctxs baseTy) = case ctxs of
    Empty -> (UnboxedValue baseTy, Nothing)
    Nest b rest -> case b of
      TabCtx _ -> error "leading idxs should have been stripped off already"
      DepPairCtx depBinder -> case splitLeadingDepPairs rest of
        PairB depBinders rest' -> case getIExprInterpretation (LeafType rest' baseTy) of
          RawValue bt -> (UnboxedValue bt, Nothing)
          BufferPtr (BufferType ixs eltTy) -> do
            let ixs' = case allNothingBs (Nest depBinder depBinders) of
                  Just UnitB -> Just ixs
                  Nothing -> Nothing
            (BoxedBuffer eltTy, ixs')
      RefCtx -> (,Nothing) $ UnboxedValue $ hostPtrTy $ elemTypeToBaseType eltTy
        where BufferType _ eltTy = getRefBufferType (LeafType rest baseTy)
      where hostPtrTy ty = PtrType (CPU, ty)
  
  allNothingBs :: Nest (MaybeB b) n l -> Maybe (UnitB n l)
  allNothingBs Empty = Just UnitB
  allNothingBs (Nest (LeftB _) _) = Nothing
  allNothingBs (Nest (RightB UnitB) rest) = allNothingBs rest
  
  splitLeadingDepPairs :: TypeCtx SimpIR n l -> PairB (Nest (MaybeB SBinder)) (TypeCtx SimpIR) n l
  splitLeadingDepPairs = \case
    Empty -> PairB Empty Empty
    Nest (DepPairCtx b) rest -> case splitLeadingDepPairs rest of
      PairB bs rest' -> PairB (Nest b bs) rest'
    ctxs -> PairB Empty ctxs
  
  tryGetBoxIdxStructure :: LeafType n -> Maybe (IndexStructure SimpIR n)
  tryGetBoxIdxStructure leafTy = snd $ getElemTypeAndIdxStructure leafTy
  
  iExprInterpretationToBaseType :: IExprInterpretation n -> BaseType
  iExprInterpretationToBaseType = \case
    RawValue b -> b
    BufferPtr (BufferType _  elemTy) -> hostPtrTy $ elemTypeToBaseType elemTy
    where hostPtrTy ty = PtrType (CPU, ty)
  
  splitLeadingIxs :: LeafType n -> Abs (IdxNest SimpIR) LeafType n
  splitLeadingIxs (LeafType (Nest (TabCtx ix) bs) t) =
    case splitLeadingIxs (LeafType bs t) of
      Abs ixs leafTy -> Abs (Nest ix ixs) leafTy
  splitLeadingIxs (LeafType bs t) = Abs Empty $ LeafType bs t
  
  elemTypeToBaseType :: BufferElementType -> BaseType
  elemTypeToBaseType = \case
    UnboxedValue t -> t
    BoxedBuffer elemTy -> hostPtrTy $ elemTypeToBaseType elemTy
    where hostPtrTy ty = PtrType (CPU, ty)
  
  typeToTree :: EnvReader m => SType n -> m n (Tree (LeafType n))
  typeToTree tyTop = return $ go REmpty tyTop
   where
    go :: RNest (TypeCtxLayer SimpIR) n l -> SType l -> Tree (LeafType n)
    go ctx (TyCon con) = case con of
      BaseType b -> Leaf $ LeafType (unRNest ctx) b
      TabPi (TabPiType d b bodyTy) -> go (RNest ctx (TabCtx (PairB (LiftB d) b))) bodyTy
      RefType _ t -> go (RNest ctx RefCtx) t
      DepPairTy (DepPairType _ (b:>t1) (t2)) -> do
        let tree1 = rec t1
        let tree2 = go (RNest ctx (DepPairCtx (JustB (b:>t1)))) t2
        Branch [tree1, tree2]
      ProdType ts -> Branch $ map rec ts
      SumType ts -> do
        let tag = rec TagRepTy
        let xs = map rec ts
        Branch $ tag:xs
      HeapType -> Branch []
      where rec = go ctx
  
  traverseScalarRepTys :: EnvReader m => SType n -> (LeafType n -> m n a) -> m n (Tree a)
  traverseScalarRepTys ty f = traverse f =<< typeToTree ty
  {-# INLINE traverseScalarRepTys #-}
  
  storeRepVal :: Emits n => Dest n -> RepVal n -> SubstImpM i n ()
  storeRepVal (Dest _ destTree) repVal@(RepVal _ valTree) = do
    leafTys <- valueToTree repVal
    forM_ (zipTrees (zipTrees leafTys destTree) valTree) \((leafTy, ptr), val) -> do
      storeLeaf leafTy ptr val
  {-# INLINE storeRepVal #-}
  
  -- Like `typeToTree`, but when we additionally have the value, we can populate
  -- the existentially-hidden fields.
  valueToTree :: EnvReader m => RepVal n -> m n (Tree (LeafType n))
  valueToTree (RepVal tyTop valTop) = do
    go REmpty tyTop valTop
   where
    go :: EnvReader m => RNest (TypeCtxLayer SimpIR) n l -> SType l -> Tree (IExpr n)
       -> m n (Tree (LeafType n))
    go ctx (TyCon ty) val = case ty of
      BaseType b -> return $ Leaf $ LeafType (unRNest ctx) b
      TabPi (TabPiType d b bodyTy) -> go (RNest ctx (TabCtx (PairB (LiftB d) b))) bodyTy val
      RefType _ t -> go (RNest ctx RefCtx) t val
      DepPairTy (DepPairType _ (b:>t1) (t2)) -> case val of
        Branch [v1, v2] -> do
          case allDepPairCtxs (unRNest ctx) of
            Just UnitB -> do
              tree1 <- rec t1 v1
              x <- repValAtom $ RepVal t1 v1
              t2' <- applySubst (b@>SubstVal x) t2
              tree2 <- go (RNest ctx (DepPairCtx NothingB )) t2' v2
              return $ Branch [tree1, tree2]
            Nothing -> do
              tree1 <- rec t1 v1
              tree2 <- go (RNest ctx (DepPairCtx (JustB (b:>t1)))) t2 v2
              return $ Branch [tree1, tree2]
        _ -> error "expected a branch"
      ProdType ts -> case val of
        Branch vals -> Branch <$> zipWithM rec ts vals
        _ -> error "expected a branch"
      SumType ts -> case val of
        Branch (tagVal:vals) -> do
          tag <- rec TagRepTy tagVal
          results <- zipWithM rec ts vals
          return $ Branch $ tag : results
        _ -> error "expected a branch"
      _ -> error $ "not implemented " ++ pprint ty
      where rec = go ctx
  {-# INLINE valueToTree #-}
  
  allDepPairCtxs :: TypeCtx SimpIR n l -> Maybe (UnitB n l)
  allDepPairCtxs ctx = case splitLeadingDepPairs ctx of
    PairB bs Empty -> allNothingBs bs
    _ -> Nothing
  
  storeLeaf :: Emits n => LeafType n -> IExpr n -> IExpr n -> SubstImpM i n ()
  storeLeaf leafTy dest src = case getRefBufferType leafTy of
    BufferType Singleton (UnboxedValue _) -> store dest src
    BufferType idxStructure (UnboxedValue _) -> do
      numel <- computeElemCountImp idxStructure
      memcopy dest src numel
    BufferType Singleton (BoxedBuffer elemTy) -> do
      load dest >>= freeBox elemTy
      Just boxIxStructure <- return $ tryGetBoxIdxStructure leafTy
      boxSize <- computeElemCountImp boxIxStructure
      cloneBox dest elemTy (Just boxSize) src
    BufferType idxStructure (BoxedBuffer elemTy) -> do
      numelem <- computeElemCountImp idxStructure
      emitLoop "i" Fwd numelem \i -> do
        curDest <- impOffset (sink dest) i
        load curDest >>= freeBox elemTy
        srcBox <- impOffset (sink src) i >>= load
        cloneBox curDest elemTy Nothing srcBox
  {-# INLINE storeLeaf #-}
  
  freeBox :: Emits n => BufferElementType -> IExpr n -> SubstImpM i n ()
  freeBox elemTy ptr = do
    ifNull ptr (return ()) do
      ptr' <- sinkM ptr
      case elemTy of
        UnboxedValue _ -> return ()
        BoxedBuffer innerBoxElemTy -> do
          numElem <- emitInstr $ GetAllocSize ptr'
          mapOffsetPtrs numElem [ptr'] \[offsetPtr] ->
            load offsetPtr  >>= freeBox innerBoxElemTy
      emitStatement $ Free ptr'
  
  -- If the buffer size (in elements) is not provided, then it assumed to be
  -- available by querying the runtime. This means that it must be a pointer
  -- directly obtained by calling `malloc_alloc`, not an offset-derived pointer
  -- thereof.
  cloneBox :: Emits n => IExpr n -> BufferElementType -> Maybe (IExpr n) -> IExpr n -> SubstImpM i n ()
  cloneBox dest elemTy maybeBufferNumElem srcPtr = do
    ifNull srcPtr
      (store (sink dest) (nullPtrIExpr $ elemTypeToBaseType elemTy))
      do
        numElem <- case maybeBufferNumElem of
          Just n -> sinkM n
          Nothing -> emitInstr $ GetAllocSize $ sink srcPtr
        -- It's "Unmanaged" by the scoped memory system because its lifetime is
        -- determined by the array it will appear in instead. (Also it has to be
        -- heap-allocated).
        newPtr <- emitAllocWithContext Unmanaged (elemTypeToBaseType elemTy) numElem
        store (sink dest) newPtr
        case elemTy of
          UnboxedValue _ -> memcopy newPtr (sink srcPtr) numElem
          BoxedBuffer elemTy' -> do
            mapOffsetPtrs numElem [newPtr, sink srcPtr] \[newPtrOffset, srcPtrOffset] -> do
              load srcPtrOffset >>= cloneBox newPtrOffset elemTy' Nothing
  
  ifNull
    :: Emits n
    => IExpr n
    -> (forall l. (Emits l, DExt n l) => SubstImpM i l ())
    -> (forall l. (Emits l, DExt n l) => SubstImpM i l ())
    -> SubstImpM i n ()
  ifNull p trueBranch falseBranch = do
    pIsNull <- isNull p
    trueBlock  <- buildBlockImp $ [] <$ trueBranch
    falseBlock <- buildBlockImp $ [] <$ falseBranch
    emitStatement $ ICond pIsNull trueBlock falseBlock
  
  isNull :: (ImpBuilder m, Emits n) => IExpr n -> m n (IExpr n)
  isNull p = do
    let PtrType (_, baseTy) = getIType p
    emitInstr $ IBinOp (ICmp Equal) p (nullPtrIExpr baseTy)
  {-# INLINE isNull #-}
  
  nullPtrIExpr :: BaseType -> IExpr n
  nullPtrIExpr baseTy = ILit $ PtrLit (CPU, baseTy) NullPtr
  
  loadRepVal :: (ImpBuilder m, Emits n) => Dest n -> m n (RepVal n)
  loadRepVal (Dest valTy destTree) = do
    leafTys <- typeToTree valTy
    RepVal valTy <$> forM (zipTrees leafTys destTree) \(leafTy, ptr) -> do
      BufferType size _ <- return $ getRefBufferType leafTy
      case size of
        Singleton -> load ptr
        _         -> return ptr
  {-# INLINE loadRepVal #-}
  
  atomToRepVal :: Emits n => SAtom n -> SubstImpM i n (RepVal n)
  atomToRepVal x = RepVal (getType x) <$> go x where
    go :: Emits n => SAtom n -> SubstImpM i n (Tree (IExpr n))
    go (Con con) = case con of
      DepPair lhs rhs _ -> do
        lhsTree <- go lhs
        rhsTree <- go rhs
        return $ Branch [lhsTree, rhsTree]
      Lit l -> return $ Leaf $ ILit l
      ProdCon xs -> Branch <$> mapM go xs
      SumCon cases tag payload -> do
        tag' <- go $ TagRepVal $ fromIntegral tag
        xs <- forM (enumerate cases) \(i, t) -> if i == tag
          then go payload
          else buildGarbageVal t <&> \(Stuck _ (RepValAtom (RepVal _ tree))) -> tree
        return $ Branch $ tag':xs
      HeapVal -> return $ Branch []
    go (Stuck _ stuck) = case stuck of
      Var v -> lookupAtomName (atomVarName v) >>= \case
        TopDataBound (RepVal _ tree) -> return tree
        _ -> error "should only have pointer and data atom names left"
      PtrVar ty p -> return $ Leaf $ IPtrVar p ty
      RepValAtom dRepVal -> do
        (RepVal _ tree) <- return dRepVal
        return tree
      -- TODO: I think we want to be able to rule this one out by insisting that
      -- RepValAtom is itself part of Stuck and it can't represent a product.
      StuckProject i val -> do
        Branch ts <- go =<< mkStuck val
        return $ ts !! i
      StuckTabApp f x' -> do
        f' <- atomToRepVal =<< mkStuck f
        RepVal _ t <- indexRepVal f' x'
        return t
  
  -- XXX: We used to have a function called `destToAtom` which loaded the value
  -- from the dest. This version is not that. It just lifts a dest into an atom of
  -- type `Ref _`.
  destToAtom :: Dest n -> SAtom n
  destToAtom (Dest valTy tree) = toAtom $ RepVal (RefTy (Con HeapVal) valTy) tree
  
  atomToDest :: EnvReader m => SAtom n -> m n (Dest n)
  atomToDest (Stuck _ (RepValAtom val)) = do
    (RepVal ~(RefTy _ valTy) valTree) <- return val
    return $ Dest valTy valTree
  atomToDest atom = error $ "Expected a non-var atom of type `RawRef _`, got: " ++ pprint atom
  {-# INLINE atomToDest #-}
  
  repValToList :: RepVal n -> [IExpr n]
  repValToList (RepVal _ tree) = toList tree
  
  -- TODO: augment with device, backend information as needed
  data AllocContext = Managed | Unmanaged deriving (Show, Eq)
  allocDestWithAllocContext :: Emits n => AllocContext -> SType n -> SubstImpM i n (Dest n)
  allocDestWithAllocContext ctx destValTy = do
    descTree <- typeToTree destValTy
    destTree <- forM descTree \leafTy -> allocBuffer ctx $ getRefBufferType leafTy
    return $ Dest destValTy $ destTree
  
  allocBuffer :: Emits n => AllocContext -> BufferType n -> SubstImpM i n (IExpr n)
  allocBuffer ctx (BufferType idxStructure elemTy) = do
    numel <- computeElemCountImp idxStructure
    let baseTy = elemTypeToBaseType elemTy
    ptr <- emitAllocWithContext ctx baseTy numel
    case elemTy of
      UnboxedValue _ -> return ()
      BoxedBuffer boxElemTy ->
        case numel of
          IIdxRepVal 1 -> store ptr (nullPtrIExpr $ elemTypeToBaseType boxElemTy)
          _ -> emitStatement $ InitializeZeros ptr numel
    return ptr
  
  emitAllocWithContext
    :: (Emits n, ImpBuilder m)
    => AllocContext -> BaseType -> IExpr n -> m n (IExpr n)
  emitAllocWithContext ctx ty size = do
    if canUseStack ctx size
      then emitInstr $ StackAlloc ty size
      else do
        ptr <- emitInstr $ Alloc CPU ty size
        case ctx of
          Managed   -> extendAllocsToFree ptr
          Unmanaged -> return ()
        return ptr
  {-# INLINE emitAllocWithContext #-}
  
  canUseStack :: AllocContext -> IExpr n -> Bool
  canUseStack Managed size | isSmall size  = True
  canUseStack _ _ = False
  
  isSmall :: IExpr n -> Bool
  isSmall (ILit l) | getIntLit l <= 256 = True
  isSmall _ = False
  
  getRepBaseTypes :: EnvReader m => SType n -> m n [BaseType]
  getRepBaseTypes ty = do
    liftM snd $ runStreamWriterT1 $ traverseScalarRepTys ty \leafTy -> do
      writeStream $ iExprInterpretationToBaseType $ getIExprInterpretation leafTy
  {-# INLINE getRepBaseTypes #-}
  
  getDestBaseTypes :: EnvReader m => SType n -> m n [BaseType]
  getDestBaseTypes ty = do
    liftM snd $ runStreamWriterT1 $ traverseScalarRepTys ty \leafTy -> do
      writeStream $ iExprInterpretationToBaseType $ BufferPtr $ getRefBufferType leafTy
  
  listToTree :: EnvReader m => SType n -> [a] -> m n (Tree a, [a])
  listToTree ty xs = runStreamReaderT1 xs $ traverseScalarRepTys ty \_ -> fromJust <$> readStream
  
  allocDestUnmanaged :: Emits n => SType n -> SubstImpM i n (Dest n)
  allocDestUnmanaged = allocDestWithAllocContext Unmanaged
  {-# INLINE allocDestUnmanaged #-}
  
  allocDest :: Emits n => SType n -> SubstImpM i n (Dest n)
  allocDest = allocDestWithAllocContext Managed
  
  storeAtom :: Emits n => Dest n -> SAtom n -> SubstImpM i n ()
  storeAtom dest x = storeRepVal dest =<< atomToRepVal x
  
  loadAtom :: Emits n => Dest n -> SubstImpM i n (SAtom n)
  loadAtom d = repValAtom =<< loadRepVal d
  
  repValFromFlatList :: (TopBuilder m, Mut n) => SType n -> [LitVal] -> m n (RepVal n)
  repValFromFlatList ty xs = do
    (litValTree, []) <- runStreamReaderT1 xs $ traverseScalarRepTys ty \_ ->
      fromJust <$> readStream
    iExprTree <- mapM litValToIExpr litValTree
    return $ RepVal ty iExprTree
  
  litValToIExpr :: (TopBuilder m, Mut n) => LitVal -> m n (IExpr n)
  litValToIExpr litval = case litval of
    PtrLit ty ptr -> do
      ptrName <- emitBinding "ptr" $ PtrBinding ty ptr
      return $ IPtrVar ptrName ty
    _ -> return $ ILit litval
  
  buildGarbageVal :: Emits n => SType n -> SubstImpM i n (SAtom n)
  buildGarbageVal ty =
    toAtom <$> RepVal ty <$> traverseScalarRepTys ty \leafTy -> do
      case getIExprInterpretation leafTy of
        BufferPtr bufferTy -> allocBuffer Managed bufferTy
        RawValue b -> return $ ILit $ emptyLit b
  
  -- === Operations on dests ===
  
  indexDest :: Emits n => Dest n -> SAtom n -> SubstImpM i n (Dest n)
  indexDest (Dest (TyCon (TabPi tabTy)) tree) i = do
    eltTy <- instantiate tabTy [i]
    ord <- ordinalImp (tabIxType tabTy) i
    leafTys <- typeToTree $ toType tabTy
    Dest eltTy <$> forM (zipTrees leafTys tree) \(leafTy, ptr) -> do
      BufferType ixStruct _ <- return $ getRefBufferType leafTy
      offset <- computeOffsetImp ixStruct ord
      impOffset ptr offset
  indexDest _ _ = error "expected a reference to a table"
  {-# INLINE indexDest #-}
  
  -- TODO: de-dup with indexDest?
  indexRepValParam :: Emits n
    => RepVal n -> SAtom n -> (SType n -> SType n)
    -> (IExpr n -> SubstImpM i n (IExpr n))
    -> SubstImpM i n (RepVal n)
  indexRepValParam (RepVal (TyCon (TabPi tabTy)) vals) i tyFunc func = do
    eltTy <- instantiate tabTy [i]
    ord <- ordinalImp (tabIxType tabTy) i
    leafTys <- typeToTree (toType tabTy)
    vals' <- forM (zipTrees leafTys vals) \(leafTy, ptr) -> do
      BufferPtr (BufferType ixStruct _) <- return $ getIExprInterpretation leafTy
      offset <- computeOffsetImp ixStruct ord
      ptr' <- impOffset ptr offset
      -- we represent scalars by value, not by reference, so we do a load
      -- if this is the last index in the table nest.
      case ixStruct of
        EmptyAbs (Nest _ Empty) -> func ptr' >>= load
        _                       -> func ptr'
    -- `func` may have changed the types of the `vals'`.  The caller must also
    -- supply `tyFunc` to reflect that change in the SType.
    return $ RepVal (tyFunc eltTy) vals'
  indexRepValParam _ _ _ _ = error "expected table type"
  {-# INLINE indexRepValParam #-}
  
  indexRepVal :: Emits n => RepVal n -> SAtom n -> SubstImpM i n (RepVal n)
  indexRepVal rep i = indexRepValParam rep i id return
  {-# INLINE indexRepVal #-}
  
  vectorIndexRepVal :: Emits n => RepVal n -> SAtom n -> SType n -> SubstImpM i n (RepVal n)
  vectorIndexRepVal rep i vty =
    -- Passing `const vty` here depends on knowing that `vectorIndexRepVal` is
    -- only called on references of scalar base type, so that the give `vty` is,
    -- actually, the type of the result of the indexing operation.
    indexRepValParam rep i (const vty) action where
      action ptr = castPtrToVectorType ptr (toIVectorType vty)
  {-# INLINE vectorIndexRepVal #-}
  
  projectDest :: Int -> Dest n -> Dest n
  projectDest i (Dest (TyCon (ProdType tys)) (Branch ds)) =
    Dest (tys!!i) (ds!!i)
  projectDest _ (Dest ty _) = error $ "Can't project dest: " ++ pprint ty
  
  destPairUnpack :: Dest n -> (Dest n, Dest n)
  destPairUnpack (Dest (PairTy t1 t2) (Branch [d1, d2])) =
    ( Dest t1 d1, Dest t2 d2 )
  destPairUnpack (Dest ty tree) =
    error $ "Can't unpack dest: " ++ pprint ty ++ "\n" ++ show tree
  
  -- === Determining buffer sizes and offsets using polynomials ===
  
  type SBuilderM = BuilderM SimpIR
  
  computeElemCountImp :: Emits n => IndexStructure SimpIR n -> SubstImpM i n (IExpr n)
  computeElemCountImp Singleton = return $ IIdxRepVal 1
  computeElemCountImp idxs = do
    result <- liftBuilderImp do
      idxs' <- sinkM idxs
      computeElemCount idxs'
    fromScalarAtom result
  
  computeOffsetImp
    :: Emits n => IndexStructure SimpIR n -> IExpr n -> SubstImpM i n (IExpr n)
  computeOffsetImp idxs ixOrd = do
    let ixOrd' = toScalarAtom ixOrd
    result <- liftBuilderImp do
      PairE idxs' ixOrd'' <- sinkM $ PairE idxs ixOrd'
      computeOffset idxs' ixOrd''
    fromScalarAtom result
  
  computeElemCount :: Emits n => IndexStructure SimpIR n -> SBuilderM n (Atom SimpIR n)
  computeElemCount (EmptyAbs Empty) =
    -- XXX: this optimization is important because we don't want to emit any decls
    -- in the case that we don't have any indices. The more general path will
    -- still compute `1`, but it might emit decls along the way.
    return $ IdxRepVal 1
  computeElemCount idxNest' = do
    let (idxList, idxNest) = indexStructureSplit idxNest'
    sizes <- forM idxList indexSetSize
    listSize <- foldM imul (IdxRepVal 1) sizes
    nestSize <- elemCountPoly idxNest
    imul listSize nestSize
  
  elemCountPoly :: Emits n => IndexStructure SimpIR n -> SBuilderM n (Atom SimpIR n)
  elemCountPoly (Abs bs UnitE) = case bs of
    Empty -> return $ IdxRepVal 1
    Nest b@(PairB (LiftB d) (_:>t)) rest -> do
     curSize <- indexSetSize $ IxType t d
     restSizes <- computeSizeGivenOrdinal b $ EmptyAbs rest
     sumUsingPolysImp curSize restSizes
  
  computeSizeGivenOrdinal
    :: EnvReader m
    => IxBinder SimpIR n l -> IndexStructure SimpIR l
    -> m n (Abs SBinder SExpr n)
  computeSizeGivenOrdinal (PairB (LiftB d) (b:>t)) idxStruct = liftBuilder do
    withFreshBinder noHint IdxRepTy \bOrdinal ->
      Abs bOrdinal <$> buildBlock do
        i <- unsafeFromOrdinal (sink $ IxType t d) $ toAtom $ sink $ binderVar bOrdinal
        idxStruct' <- applySubst (b@>SubstVal i) idxStruct
        elemCountPoly $ sink idxStruct'
  
  -- Split the index structure into a prefix of non-dependent index types
  -- and a trailing nest of indices that can contain inter-dependencies.
  indexStructureSplit :: IndexStructure SimpIR n -> ([IxType SimpIR n], IndexStructure SimpIR n)
  indexStructureSplit (Abs Empty UnitE) = ([], EmptyAbs Empty)
  indexStructureSplit s@(Abs (Nest (PairB (LiftB d) b) rest) UnitE) =
    case hoist b (EmptyAbs rest) of
      HoistFailure _     -> ([], s)
      HoistSuccess rest' -> (IxType (binderType b) d:ans1, ans2)
        where (ans1, ans2) = indexStructureSplit rest'
  
  computeOffset :: forall n. Emits n
                => IndexStructure SimpIR n -> SAtom n -> SBuilderM n (SAtom n)
  computeOffset (EmptyAbs (Nest _ Empty)) i = return i  -- optimization
  computeOffset (EmptyAbs (Nest b idxs)) idxOrdinal = do
    case hoist b (EmptyAbs idxs) of
      HoistFailure _ -> do
       rhsElemCounts <- computeSizeGivenOrdinal b (EmptyAbs idxs)
       sumUsingPolysImp idxOrdinal rhsElemCounts
      HoistSuccess idxs' -> do
        stride <- computeElemCount idxs'
        idxOrdinal `imul` stride
  computeOffset _ _ = error "Expected a nonempty nest of idx binders"
  
  sumUsingPolysImp
    :: Emits n => SAtom n
    -> Abs SBinder SExpr n -> BuilderM SimpIR n (SAtom n)
  sumUsingPolysImp lim (Abs i body) = do
    ab <- hoistDecls i body
    sumUsingPolys lim ab
  
  hoistDecls
    :: ( Builder SimpIR m, EnvReader m, Emits n
       , BindsNames b, BindsEnv b, RenameB b, SinkableB b)
    => b n l -> SExpr l -> m n (Abs b SExpr n)
  hoistDecls b block = do
    emitDecls =<< liftEnvReaderM do
      refreshAbs (Abs b block) \b' body ->
        hoistDeclsRec b' Empty body
  {-# INLINE hoistDecls #-}
  
  hoistDeclsRec
    :: (BindsNames b, SinkableB b)
    => b n1 n2 -> SDecls n2 n3 -> SExpr n3
    -> EnvReaderM n3 (Abs SDecls (Abs b SExpr) n1)
  hoistDeclsRec = undefined
  -- hoistDeclsRec b declsAbove Empty result =
  --   return $ Abs Empty $ Abs b $ Abs declsAbove result
  -- hoistDeclsRec b declsAbove (Nest decl declsBelow) result  = do
  --   let (Let _ expr) = decl
  --   let exprIsPure = isPure expr
  --   refreshAbs (Abs decl (Abs declsBelow result))
  --     \decl' (Abs declsBelow' result') ->
  --       case exchangeBs (PairB (PairB b declsAbove) decl') of
  --         HoistSuccess (PairB hoistedDecl (PairB b' declsAbove')) | exprIsPure -> do
  --           Abs hoistedDecls fullResult <- hoistDeclsRec b' declsAbove' declsBelow' result'
  --           return $ Abs (Nest hoistedDecl hoistedDecls) fullResult
  --         _ -> hoistDeclsRec b (declsAbove >>> Nest decl' Empty) declsBelow' result'
  -- {-# INLINE hoistDeclsRec #-}
  
  -- === Imp IR builder ===
  
  data ImpInstrResult (n::S) = NoResults | OneResult !(IExpr n) | MultiResult !([IExpr n])
  
  class (EnvReader m, EnvExtender m, Fallible1 m) => ImpBuilder (m::MonadKind1) where
    emitMultiReturnInstr :: Emits n => ImpInstr n -> m n (ImpInstrResult n)
    emitDeclsImp :: (RenameE e, Emits n) => Abs (Nest ImpDecl) e n -> m n (e n)
    buildScopedImp
      :: SinkableE e
      => (forall l. (Emits l, DExt n l) => m l (e l))
      -> m n (Abs (Nest ImpDecl) e n)
    extendAllocsToFree :: Mut n => IExpr n -> m n ()
  
  type ImpBuilder2 (m::MonadKind2) = forall i. ImpBuilder (m i)
  
  withFreshIBinder
    :: ImpBuilder m
    => NameHint -> IType
    -> (forall l. DExt n l => IBinder n l -> m l a)
    -> m n a
  withFreshIBinder hint ty cont = do
    withFreshBinder hint (ImpNameBinding ty) \(b:>_) ->
      cont $ IBinder b ty
  {-# INLINE withFreshIBinder #-}
  
  emitCall
    :: Emits n => PiType SimpIR n
    -> ImpFunName n -> [SAtom n] -> SubstImpM i n (SAtom n)
  emitCall (PiType bs (EffTy _ resultTy)) f xs = do
    resultTy' <- applySubst (bs @@> map SubstVal xs) resultTy
    dest <- allocDest resultTy'
    argsImp <- forM xs \x -> repValToList <$> atomToRepVal x
    destImp <- repValToList <$> atomToRepVal (destToAtom dest)
    let impArgs = concat argsImp ++ destImp
    _ <- impCall f impArgs
    loadAtom dest
  
  buildImpFunction
    :: CallingConvention
    -> [(NameHint, IType)]
    -> (forall l. (Emits l, DExt n l) => [IExpr l] -> SubstImpM i l [IExpr l])
    -> SubstImpM i n (ImpFunction n)
  buildImpFunction cc argHintsTys body = do
    Abs bs (Abs decls (ListE results)) <-
      buildImpNaryAbs argHintsTys \vs -> ListE <$> body (map (uncurry IVar) vs)
    let resultTys = map getIType results
    let impFun = IFunType cc (map snd argHintsTys) resultTys
    return $ ImpFunction impFun $ Abs bs $ ImpBlock decls results
  
  buildImpNaryAbs
    :: HasNamesE e
    => [(NameHint, IType)]
    -> (forall l. (Emits l, DExt n l) => [(Name ImpNameC l, BaseType)] -> SubstImpM i l (e l))
    -> SubstImpM i n (Abs (Nest IBinder) (Abs (Nest ImpDecl) e) n)
  buildImpNaryAbs [] cont = do
    Distinct <- getDistinct
    Abs Empty <$> buildScopedImp (cont [])
  buildImpNaryAbs ((hint,ty) : rest) cont = do
    withFreshIBinder hint ty \b -> do
      ab <- buildImpNaryAbs rest \vs -> do
        v <- sinkM $ binderName b
        cont ((v,ty) : vs)
      Abs bs body <- return ab
      return $ Abs (Nest b bs) body
  
  emitInstr :: (ImpBuilder m, Emits n) => ImpInstr n -> m n (IExpr n)
  emitInstr instr = do
    xs <- emitMultiReturnInstr instr
    case xs of
      OneResult x -> return x
      _   -> error "unexpected numer of return values"
  {-# INLINE emitInstr #-}
  
  emitStatement :: (ImpBuilder m, Emits n) => ImpInstr n -> m n ()
  emitStatement instr = do
    xs <- emitMultiReturnInstr instr
    case xs of
      NoResults -> return ()
      _         -> error "unexpected numer of return values"
  {-# INLINE emitStatement #-}
  
  impCall :: (ImpBuilder m, Emits n) => ImpFunName n -> [IExpr n] -> m n [IExpr n]
  impCall f args = emitMultiReturnInstr (ICall f args) <&> \case
    NoResults      -> []
    OneResult x    -> [x]
    MultiResult xs -> xs
  {-# INLINE impCall #-}
  
  impOffset :: Emits n => IExpr n -> IExpr n -> SubstImpM i n (IExpr n)
  impOffset ref (IIdxRepVal 0) = return ref
  impOffset ref off = emitInstr $ IPtrOffset ref off
  {-# INLINE impOffset #-}
  
  memcopy :: (ImpBuilder m, Emits n) => IExpr n -> IExpr n -> IExpr n -> m n ()
  memcopy new src numElem = emitStatement $ MemCopy new src numElem
  {-# INLINE memcopy #-}
  
  cast :: Emits n => IExpr n -> BaseType -> SubstImpM i n (IExpr n)
  cast x bt = emitInstr $ ICastOp bt x
  
  load :: (ImpBuilder m, Emits n) => IExpr n -> m n (IExpr n)
  load x = emitInstr $ IPtrLoad x
  {-# INLINE load #-}
  
  store :: (ImpBuilder m, Emits n) => IExpr n -> IExpr n -> m n ()
  store dest src = emitStatement $ Store dest src
  {-# INLINE store #-}
  
  -- TODO: Consider targeting LLVM's `switch` instead of chained conditionals.
  emitSwitch :: forall i n a.  Emits n
             => IExpr n
             -> [a]
             -> (forall l. (Emits l, DExt n l) => a -> SubstImpM i l ())
             -> SubstImpM i n ()
  emitSwitch testIdx args cont = do
    Distinct <- getDistinct
    rec 0 args
    where
      rec :: forall l. (Emits l, DExt n l) => Int -> [a] -> SubstImpM i l ()
      rec _ [] = error "Shouldn't have an empty list of alternatives"
      rec _ [arg] = cont arg
      rec curIdx (arg:rest) = do
        curTag     <- fromScalarAtom $ TagRepVal $ fromIntegral curIdx
        cond       <- emitInstr $ IBinOp (ICmp Equal) (sink testIdx) curTag
        thisCase   <- buildBlockImp $ cont arg >> return []
        otherCases <- buildBlockImp $ rec (curIdx + 1) rest >> return []
        emitStatement $ ICond cond thisCase otherCases
  
  mapOffsetPtrs
    :: Emits n
    => IExpr n -> [IExpr n]
    -> (forall l. (DExt n l, Emits l) => [IExpr l] -> SubstImpM i l ())
    -> SubstImpM i n ()
  mapOffsetPtrs numelem basePtrs cont = do
    emitLoop "i" Fwd numelem \i -> do
      offsetPtrs <- forM basePtrs \basePtr -> impOffset (sink basePtr) i
      cont offsetPtrs
  
  emitLoop :: Emits n
           => NameHint -> Direction -> IExpr n
           -> (forall l. (DExt n l, Emits l) => IExpr l -> SubstImpM i l ())
           -> SubstImpM i n ()
  emitLoop _ _ (IIdxRepVal 0) _ = return ()
  emitLoop _ _ (IIdxRepVal 1) cont = do
    Distinct <- getDistinct
    cont (IIdxRepVal 0)
  emitLoop hint d n cont = do
    loopBody <- do
      withFreshIBinder hint (getIType n) \b@(IBinder _ ty)  -> do
        let i = IVar (binderName b) ty
        body <- buildBlockImp do
                  cont =<< sinkM i
                  return []
        return $ Abs b body
    emitStatement $ IFor d n loopBody
  
  _emitDebugPrint :: (ImpBuilder m, Emits n) => String -> IExpr n -> m n ()
  _emitDebugPrint fmt x = do
   x' <- emitInstr $ ICastOp (Scalar Int64Type) x
   emitStatement $ DebugPrint fmt x'
  
  restructureScalarOrPairType :: SType n -> [IExpr n] -> SubstImpM i n (SAtom n)
  restructureScalarOrPairType topTy topXs =
    go topTy topXs >>= \case
      (atom, []) -> return atom
      _ -> error "Wrong number of scalars"
    where
      go (PairTy t1 t2) xs = do
        (atom1, rest1) <- go t1 xs
        (atom2, rest2) <- go t2 rest1
        return (PairVal atom1 atom2, rest2)
      go (BaseTy _) (x:xs) = do
        let x' = toScalarAtom x
        return (x', xs)
      go ty _ = error $ "Not a scalar or pair: " ++ pprint ty
  
  buildBlockImp
    :: (forall l. (Emits l, DExt n l) => SubstImpM i l [IExpr l])
    -> SubstImpM i n (ImpBlock n)
  buildBlockImp cont = do
    Abs decls (ListE results) <- buildScopedImp $ ListE <$> cont
    return $ ImpBlock decls results
  
  -- === Atom <-> IExpr conversions ===
  
  fromScalarAtom :: Emits n => SAtom n -> SubstImpM i n (IExpr n)
  fromScalarAtom atom = atomToRepVal atom >>= \case
    RepVal ty tree -> case tree of
      Leaf x -> return x
      _ -> error $ "Not a scalar atom:" ++ pprint ty
  
  toScalarAtom :: forall n. IExpr n -> SAtom n
  toScalarAtom x = toAtom $ RepVal (BaseTy (getIType x) :: SType n) (Leaf x)
  
  liftBuilderImp :: (Emits n, SubstE AtomSubstVal e, SinkableE e)
                 => (forall l. (Emits l, DExt n l) => BuilderM SimpIR l (e l))
                 -> SubstImpM i n (e n)
  liftBuilderImp cont = do
    Abs decls result <- liftBuilder $ buildScoped cont
    dropSubst $ translateDeclNest decls $ substM result
  {-# INLINE liftBuilderImp #-}
  
  -- === Type classes ===
  
  ordinalImp :: Emits n => IxType SimpIR n -> SAtom n -> SubstImpM i n (IExpr n)
  ordinalImp (IxType _ (DictCon dict)) i = fromScalarAtom =<< case dict of
    IxRawFin _ -> return i
    IxSpecialized d params -> do
      appSpecializedIxMethod d Ordinal (params ++ [i])
  
  unsafeFromOrdinalImp :: Emits n => IxType SimpIR n -> IExpr n -> SubstImpM i n (SAtom n)
  unsafeFromOrdinalImp (IxType _ (DictCon dict)) i = do
    let i' = toScalarAtom i
    case dict of
      IxRawFin _ -> return i'
      IxSpecialized d params ->
        appSpecializedIxMethod d UnsafeFromOrdinal (params ++ [i'])
  
  indexSetSizeImp :: Emits n => IxType SimpIR n -> SubstImpM i n (IExpr n)
  indexSetSizeImp (IxType _ (DictCon dict)) = do
    fromScalarAtom =<< case dict of
      IxRawFin n -> return n
      IxSpecialized d params ->
        appSpecializedIxMethod d Size (params ++ [])
  
  appSpecializedIxMethod :: Emits n => SpecDictName n -> IxMethod -> [SAtom n] -> SubstImpM i n (SAtom n)
  appSpecializedIxMethod d method args = do
    SpecializedDict _ (Just fs) <- lookupSpecDict d
    TopLam _ _ (LamExpr bs body) <- return $ fs !! fromEnum method
    dropSubst $ extendSubst (bs @@> map SubstVal args) $ translateExpr body
  
  -- === Abstracting link-time objects ===
  
  abstractLinktimeObjects :: forall m n. EnvReader m
    => ImpFunction n -> m n (ClosedImpFunction n, [TopFunName n], [PtrName n])
  abstractLinktimeObjects f = do
    let allVars = freeVarsE f
    (funVars, funTys) <- unzip <$> forMFilter (nameSetToList @TopFunNameC allVars) \v ->
      lookupTopFun v >>= \case
        DexTopFun _ (TopLam _ piTy _) _ -> do
          ty' <- getImpFunType StandardCC piTy
          return $ Just (v, ty')
        FFITopFun _ _ -> return Nothing
    (ptrVars, ptrTys) <- unzip <$> forMFilter (nameSetToList @PtrNameC allVars) \v -> do
      (ty, _) <- lookupPtrName v
      return $ Just (v, ty)
    Abs funBs (Abs ptrBs f') <- return $ abstractFreeVarsNoAnn funVars $
                                         abstractFreeVarsNoAnn ptrVars f
    let funBs' =  zipWithNest funBs funTys \b ty -> IFunBinder b ty
    let ptrBs' =  zipWithNest ptrBs ptrTys \b ty -> PtrBinder  b ty
    return (ClosedImpFunction funBs' ptrBs' f', funVars, ptrVars)
  
  -- === singleton types ===
  
  -- The following implementation should be valid:
  --   isSingletonType :: EnvReader m => Type n -> m n Bool
  --   isSingletonType ty =
  --     singletonTypeVal ty >>= \case
  --       Nothing -> return False
  --       Just _  -> return True
  -- But a separate implementation doesn't have to go under binders,
  -- because it doesn't have to reconstruct the singleton value (which
  -- may be type annotated and whose type may refer to names).
  
  isSingletonType :: Type SimpIR n -> Bool
  isSingletonType topTy = isJust $ checkIsSingleton topTy
    where
      checkIsSingleton :: SType n -> Maybe ()
      checkIsSingleton (TyCon ty) = case ty of
        TabPi (TabPiType _ _ body) -> checkIsSingleton body
        ProdType tys -> mapM_ checkIsSingleton tys
        _ -> Nothing
  
  singletonTypeVal :: (EnvReader m)
    => Type SimpIR n -> m n (Maybe (Atom SimpIR n))
  singletonTypeVal ty = do
    tree <- typeToTree ty
    if length tree == 0 then do
      -- The tree has 0 of these if the type is empty
      let tree' = fmap (const $ ILit $ Int32Lit 0) tree
      Just <$> mkStuck (RepValAtom $ RepVal ty tree')
    else
      return Nothing
  {-# INLINE singletonTypeVal #-}
  
  -- === type checking imp programs ===
  
  toIVectorType :: SType n -> IVectorType
  toIVectorType = \case
    BaseTy vty@(Vector _ _) -> vty
    _ -> error "Not a vector type"
  
  impFunType :: ImpFunction n -> IFunType
  impFunType (ImpFunction ty _) = ty
  
  getIType :: IExpr n -> IType
  getIType (ILit l) = litType l
  getIType (IVar _ ty) = ty
  getIType (IPtrVar _ ty) = PtrType ty
  
  impInstrTypes :: EnvReader m => ImpInstr n -> m n [IType]
  impInstrTypes instr = case instr of
    IBinOp op x _   -> return [typeBinOp op (getIType x)]
    IUnOp  op x     -> return [typeUnOp  op (getIType x)]
    ICastOp t _     -> return [t]
    IBitcastOp t _  -> return [t]
    Alloc a ty _    -> return [PtrType (a  , ty)]
    StackAlloc ty _ -> return [PtrType (CPU, ty)]
    Store _ _       -> return []
    Free _          -> return []
    IThrowError     -> return []
    MemCopy _ _ _   -> return []
    InitializeZeros _ _ -> return []
    GetAllocSize _  -> return [IIdxRepTy]
    IFor _ _ _      -> return []
    IWhile _        -> return []
    ICond _ _ _     -> return []
    ILaunch _ _ _   -> return []
    ISyncWorkgroup  -> return []
    IVectorBroadcast _ vty -> return [vty]
    IVectorIota vty        -> return [vty]
    DebugPrint _ _  -> return []
    IQueryParallelism _ _ -> return [IIdxRepTy, IIdxRepTy]
    ICall f _ -> lookupTopFun f >>= \case
      DexTopFun _ (TopLam _ piTy _) _ -> do
        IFunType _ _ resultTys <- getImpFunType StandardCC piTy
        return resultTys
      FFITopFun _ (IFunType _ _ resultTys) -> return resultTys
    ISelect  _ x  _  -> return [getIType x]
    IPtrLoad ref -> return [ty]  where PtrType (_, ty) = getIType ref
    IPtrOffset ref _ -> return [PtrType (addr, ty)]  where PtrType (addr, ty) = getIType ref
    IOutputStream    -> return [hostPtrTy $ Scalar Word8Type]
    IShowScalar _ _  -> return [Scalar Word32Type]
    where hostPtrTy ty = PtrType (CPU, ty)
  
  instance CheckableE SimpIR ImpFunction where
    checkE = renameM -- TODO
  
  -- TODO: Don't use Core Envs for Imp!
  instance BindsEnv ImpDecl where
    toEnvFrag (ImpLet bs _) = toEnvFrag bs
  
  instance BindsEnv IBinder where
    toEnvFrag (IBinder b ty) =  toEnvFrag $ b :> ImpNameBinding ty
  
  instance Pretty (LeafType n) where
    pretty (LeafType ctx base) = pretty ctx <+> pretty base
  
  instance Pretty (TypeCtxLayer SimpIR n l) where
    pretty = \case
      TabCtx (PairB _ b)        -> pretty b
      DepPairCtx (RightB UnitB) -> "dep-pair-instantiated"
      DepPairCtx (LeftB b)      -> "dep-pair" <+> pretty b
      RefCtx               -> "refctx"
  
  -- See Note [Confuse GHC] from Simplify.hs
  confuseGHC :: EnvReader m => m n (DistinctEvidence n)
  confuseGHC = getDistinct
  {-# INLINE confuseGHC #-}
  
  -- === debug instrumentation ===
  
  -- Adds a `DebugPrint XXX` after each Imp instruction. It's a crude tracing
  -- mechanism that's useful for debugging segfaults.
  -- To use: put something like this in an appropriate place in `TopLevel.hs`
  --    let fImp = addImpTracing (pprint fname ++ " %x\n") impFun' fImp'
  addImpTracing :: String -> ImpFunction n -> ImpFunction n
  addImpTracing fmt (ImpFunction ty (Abs bs body)) =
    ImpFunction ty (Abs bs (evalState (go REmpty body) 0))
   where
     go :: RNest ImpDecl n l -> ImpBlock l -> MTL.State Int (ImpBlock n)
     go accum (ImpBlock Empty result) = return $ ImpBlock (unRNest accum) result
     go accum (ImpBlock (Nest decl decls) result) = do
       n <- next
       let traceDecl = ImpLet Empty (DebugPrint fmt (ILit $ Word64Lit $ fromIntegral n))
       decl' <- traverseDeclBlocks (go REmpty) decl
       go (accum `RNest` traceDecl `RNest` decl') (ImpBlock decls result)
  
     next :: MTL.State Int Int
     next = do
       n <- get
       put $ n + 1
       return n
  
  traverseDeclBlocks
    :: Monad m => (forall n'. ImpBlock n' -> m (ImpBlock n'))
    -> ImpDecl n l -> m (ImpDecl n l)
  traverseDeclBlocks f (ImpLet bs instr) = ImpLet bs <$> case instr of
    IFor d n (Abs b block) -> IFor d n <$> (Abs b <$> f block)
    ICond cond block1 block2 -> ICond cond <$> f block1 <*> f block2
    -- TODO: fill in the other cases as needed
    _ -> return instr
  
  -- === notes ===
  
  {-
  
  Note [memory management]
  
  Dex emphasizes flat array-like data structures unlike other functional languages
  which tend to use trees of pointers. Instead of dynamically counting references
  to heap-allocated objects we do all memory management at compile time based on
  static scopes. The generated code uses malloc/free in a style similar to C
  programming. We use destination-passing style to call functions.
  
  Every function (or code block -- loop body, case body, etc.) is passed a
  destination in which to store its final result. The function or block may
  allocate memory for intermediate results but it must free anything it allocated
  before returning. There's an important exception, to do with dependent pairs,
  which we discuss soon.
  
  The main unit of memory management is the fixed-size buffer, a contiguous chunk
  of memory directly storing scalar base values, int32, float64 etc, with no
  wasted space. The main Dex runtime data types -- tables, products, sums, base
  types and mixtures of all these -- are all represented by some number of these.
  We can determine the required size of the buffer based solely on the type, since
  tables encode their size in their type. This lets us allocate buffers for
  results before calling a function or starting a block.
  
  A buffer is represented by a pointer which is guaranteed to have been generated
  by a call to `malloc_dex` which means that it can be freed with a call to
  `free_dex`. `malloc_dex` also stores the size of the buffer in the few bytes
  before the data part, which lets us make copies of buffers without having to
  keep track of their sizes (e.g. back in the Haskell runtime).
  
  The destination-passing style, together with the free-before-returning style,
  lets us generate as many views as we like of a buffer without having to count
  references (dynamically or statically). For example, we use pointers to
  interior memory locations in a buffer to represent array slices and subarrays.
  We don't need to worry about use-after-free errors because we don't free the
  buffer until the function/block returns, at which point those views will no
  longer be available. The only thing that escapes the function is the data
  written into the destination supplied by the caller.
  
  Now back to the important exception, dependent pairs. When we have a dependent
  pair, like `(n:Nat &> Fin n => Float)`, we don't know the size of the `Fin n =>
  Float` table because the `n` is given by the *value* of the (first element of
  the) pair. We use these to encode dynamically-sized lists.
  
  We handle having an array of these by pretending that we can have an array of
  arbitrary-sized boxes. The implementation is just an array of pointers. But
  these pointers behave quite differently from either the pointers that point to
  standard buffers or the pointers used as views of those buffers. We think of
  them as an implementation detail modeling the interface of an array with
  variable-sized elements, stored *as values* in the array. In particular, the
  outer buffer exclusively owns the memory backing its boxes. When we free the
  buffer, we free its boxes. When we write to a buffer at an index, allocate fresh
  memory for the new box and free the old box (unless it's null, as for an
  uninitialized buffer). We can still take read-only views of the data in the
  boxes, but only if we know the buffer itself is immutable/frozen, because
  otherwise the box memory might be freed when it's overwritten by another value,
  which could happen before exiting the scope (in that sense it's no different
  from taking a slice of a buffer to represent a subarray).
  
  Separate from this memory system, we have user-facing references, `Ref h a`,
  from the `State` effect. The memory for these is managed by a separate system,
  the `runState` scope. The references are represented as pointers, and they may
  appear in buffers if we happen to have a table of references in the user
  program. But they are not boxes, and the memory they point to certainly isn't
  owned by the buffer. For one thing, the same pointer can appear in multiple
  tables, or multiple times in one table. To handle these references, we just
  treat them as pointers to foreign memory. Once we've actually loaded one of
  these pointers and we have it as a value, *then* we can interpret it as a buffer
  so that we can read/write from it. Since the allocation and freeing is done by
  `runState`, and the use of the references is guarded by the heap parameter, we
  have the same sort of guarantee about having no use-after-free bugs.
  
  -}
  