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 UndecidableInstances #-}
  
  module Simplify
    ( simplifyTopBlock, simplifyTopFunction, ReconstructAtom (..), applyReconTop,
      linearizeTopFun, SimplifiedTopLam (..)) where
  
  import Control.Category ((>>>))
  import Control.Monad
  import Control.Monad.Reader
  import Data.Maybe
  
  import Builder
  import CheapReduction
  import CheckType
  import Core
  import Err
  import Generalize
  import IRVariants
  import Linearize
  import Name
  import Subst
  import PPrint
  import QueryType
  import RuntimePrint
  import Transpose
  import Types.Core
  import Types.Source
  import Types.Top
  import Types.Primitives
  import Util (enumerate)
  
  -- === Simplification ===
  
  -- The purpose of simplification is that we want first-class functions
  -- in the surface language, but we don't want to deal with them when
  -- generating code.  To that end, simplification inlines almost all
  -- functions everywhere they appear.
  
  -- In particular, simplification also carries out AD by discharging
  -- the `Linearize` and `Transpose` constructors of `PrimHof`.
  
  -- The exception is functions explicitly tagged @noinline by the
  -- programmer.  Those, however, are second-class: they are all
  -- toplevel, and get specialized until they are first order.
  
  -- Currently, simplification also discharges `CatchException` by
  -- elaborating the expression into a Maybe-style monad.  Note: the
  -- plan is for `CatchException` to become a user-defined effect, and
  -- for simplification to discharge all of them.
  
  -- Simplification also opportunistically does peep-hole optimizations:
  -- some constant folding, case-of-known-constructor, projecting known
  -- elements from products, etc; but is not guaranteed to find all such
  -- opportunities.
  
  -- === Conversions between CoreIR, CoreToSimpIR, SimpIR ===
  
  tryAsDataAtom :: Emits n => CAtom n -> SimplifyM i n (Maybe (SAtom n, Type CoreIR n))
  tryAsDataAtom atom = do
    let ty = getType atom
    isData ty >>= \case
      False -> return Nothing
      True -> Just <$> do
        repAtom <- dropSubst $ toDataAtom atom
        return (repAtom, ty)
  
  data WithSubst (e::E) (o::S) where
    WithSubst :: Subst AtomSubstVal i o -> e i -> WithSubst e o
  
  type ACase = SStuck `PairE` ListE (Abs SBinder CAtom) `PairE` CType
  
  data ConcreteCAtom (n::S) =
     CCCon (WithSubst (Con CoreIR) n)
   | CCLiftSimp (CType n) (Stuck SimpIR n)
   | CCFun   (ConcreteCFun n)
   | CCTabLam (WithSubst TabLamExpr n)
   | CCACase (WithSubst ACase n)
  
  data ConcreteCFun (n::S) =
     CCLiftSimpFun (CorePiType n) (LamExpr SimpIR n)
   | CCNoInlineFun (CAtomVar n) (CType n) (CAtom n)
   | CCFFIFun      (CorePiType n) (TopFunName n)
  
  forceConstructor :: CAtom i -> SimplifyM i o (ConcreteCAtom o)
  forceConstructor atom = withDistinct case atom of
    Stuck _ stuck -> forceStuck stuck
    Con con -> do
      subst <- getSubst
      return $ CCCon $ WithSubst subst con
  
  forceStuck :: forall i o . CStuck i -> SimplifyM i o (ConcreteCAtom o)
  forceStuck stuck = withDistinct case stuck of
    Var v -> lookupSubstM (atomVarName v) >>= \case
      SubstVal x -> dropSubst $ forceConstructor x
      Rename v' -> lookupAtomName v' >>= \case
        LetBound (DeclBinding _ (Atom x)) -> dropSubst $ forceConstructor x
        NoinlineFun t f -> do
          v'' <- toAtomVar v'
          return $ CCFun $ CCNoInlineFun v'' t f
        FFIFunBound t f -> return $ CCFun $ CCFFIFun t f
        _ -> error "shouldn't have other CVars left"
    LiftSimp _ x -> do
      -- the subst should be rename-only for `x`. We should make subst IR-specific
      s <- getSubst
      let s' = newSubst \v -> case s ! v of
            SubstVal _ -> error "subst should be rename-only for SimpIR vars" -- TODO: make subst IR-specific
            Rename v' -> v'
      x' <- runSubstReaderT s' $ renameM x
      returnLifted x'
    -- We "thunk" ACase rather than forcing it because different use-cases require different ways to force it
    ACase e alts resultTy -> do
      subst <- getSubst
      return $ CCACase $ WithSubst subst $ e `PairE` ListE alts `PairE` resultTy
    TabLam e -> do
      subst <- getSubst
      return $ CCTabLam $ WithSubst subst e
    StuckProject i x -> forceStuck x >>= \case
      CCLiftSimp _ x' -> returnLifted $ StuckProject i x'
      CCCon (WithSubst s con) -> withSubst s case con of
        ProdCon xs    -> forceConstructor (xs!!i)
        DepPair l r _ -> forceConstructor ([l, r]!!i)
        _ -> error "not a product"
      CCACase x' -> pushUnderACase x' \x'' -> reduceProj i x''
      CCFun _    -> error "not a product"
      CCTabLam _ -> error "not a product"
    StuckTabApp f x -> forceStuck f >>= \case
      CCLiftSimp _ f' -> do
        x' <- toDataAtom x
        returnLifted $ StuckTabApp f' x'
      CCTabLam (WithSubst s (PairE _ (Abs b body))) -> do
        x' <- toDataAtom x
        result <- withSubst s $ extendSubst (b@>SubstVal x') $ substM body
        dropSubst $ forceConstructor result
      CCACase f' -> pushUnderACase f' \f'' -> reduceTabApp f'' =<< substM x
      CCCon _ -> error "not a table"
      CCFun _ -> error "not a table"
    StuckUnwrap x -> forceStuck x >>= \case
      CCCon (WithSubst s con) -> case con of
        NewtypeCon _ x' -> withSubst s $ forceConstructor x'
        _ -> error "not a newtype"
      CCLiftSimp _ x' -> returnLifted x'
      CCACase x' -> pushUnderACase x' \x'' -> reduceUnwrap x''
      CCFun    _ -> error "not a newtype"
      CCTabLam _ -> error "not a newtype"
    InstantiatedGiven _ _ -> error "shouldn't have this left"
    SuperclassProj _ _    -> error "shouldn't have this left"
    PtrVar ty p -> do
      p' <- substM p
      returnLifted $ PtrVar ty p'
    LiftSimpFun t f -> CCFun <$> (CCLiftSimpFun <$> substM t <*> substM f)
    where
      returnLifted :: SStuck o -> SimplifyM i o (ConcreteCAtom o)
      returnLifted s = do
        resultTy <- getType <$> substMStuck stuck
        return $ CCLiftSimp resultTy s
  
  pushUnderACase
    :: WithSubst ACase o
    -> (forall o'. DExt o o' => CAtom o' -> SimplifyM i o' (CAtom o'))
    -> SimplifyM i o (ConcreteCAtom o)
  pushUnderACase _ _ = undefined
  -- pushUnderACase (WithSubst s (scrut `PairE` ListE alts `PairE` resultTy)) cont = undefined
  -- TODO: make a buildACase to use here and elsewhere in Simplify. Maybe in CheapReduce too?
  
  
  forceACase
    :: Emits o => WithSubst ACase o
    -> (forall o'. (Emits o', DExt o o') => ConcreteCAtom o' -> SimplifyM i o' (CAtom o'))
    -> SimplifyM i o (CAtom o)
  forceACase (WithSubst subst (scrut `PairE` ListE alts `PairE` resultTy)) cont = do
      resultTy' <- withSubst subst $ substM resultTy
      scrut'    <- withSubst subst $ substMStuck scrut
      defuncCase scrut' resultTy' \i x -> do
        Abs b body <- return $ alts !! i
        body' <- withSubst (sink subst) $ extendSubst (b@>SubstVal x) $ forceConstructor body
        cont body'
  
  tryGetRepType :: Type CoreIR n -> SimplifyM i n (Maybe (SType n))
  tryGetRepType t = isData t >>= \case
    False -> return Nothing
    True  -> Just <$> dropSubst (getRepType t)
  
  getRepType :: Type CoreIR i -> SimplifyM i o (SType o)
  getRepType (StuckTy _ stuck) =
    substMStuck stuck >>= \case
      Stuck _ _ -> error "shouldn't have stuck CType after substitution"
      Con (TyConAtom tyCon) -> dropSubst $ getRepType (TyCon tyCon)
      Con _ -> error "not a type"
  getRepType (TyCon con) = case con of
    BaseType b  -> return $ toType $ BaseType b
    ProdType ts -> toType . ProdType <$> mapM getRepType ts
    SumType  ts -> toType . SumType  <$> mapM getRepType ts
    RefType h a -> toType <$> (RefType <$> toDataAtomAssumeNoDecls h <*> getRepType a)
    HeapType    -> return $ toType HeapType
    DepPairTy (DepPairType expl b r) -> do
      withSimplifiedBinder b \b' -> do
        r' <- getRepType r
        return $ toType $ DepPairType expl b' r'
    TabPi (TabPiType ixDict b r) -> do
      ixDict' <- simplifyIxDict ixDict
      withSimplifiedBinder b \b' -> do
        r' <- getRepType r
        return $ toType $ TabPi $ TabPiType ixDict' b' r'
    NewtypeTyCon con' -> do
      (_, ty') <- unwrapNewtypeType =<< substM con'
      dropSubst $ getRepType ty'
    Pi _     -> error notDataType
    DictTy _ -> error notDataType
    TypeKind -> error notDataType
    where notDataType = "Not a type of runtime-representable data"
  
  toDataAtom :: CAtom i -> SimplifyM i o (SAtom o)
  toDataAtom (Con con) = case con of
    Lit v -> return $ toAtom $ Lit v
    ProdCon xs -> toAtom . ProdCon <$> mapM rec xs
    SumCon  tys tag x -> toAtom <$> (SumCon <$> mapM getRepType tys <*> pure tag <*> rec x)
    HeapVal           -> return $ toAtom HeapVal
    DepPair x y ty -> do
      TyCon (DepPairTy ty') <- getRepType $ TyCon $ DepPairTy ty
      toAtom <$> (DepPair <$> rec x <*> rec y <*> pure ty')
    NewtypeCon _ x  -> rec x
    Lam _         -> notData
    DictConAtom _ -> notData
    Eff _         -> notData
    TyConAtom _   -> notData
    where
      rec = toDataAtom
      notData = error $ "Not runtime-representable data"
  toDataAtom (Stuck _ stuck) = forceStuck stuck >>= \case
    CCCon (WithSubst s con) -> withSubst s $ toDataAtom (Con con)
    CCLiftSimp _ e -> mkStuck e
    CCFun    _   -> notData
    CCACase  _   -> notData -- TODO: make sure we observe this invariant"
    CCTabLam _   -> notData -- TODO: make sure we observe this invariant"
    where notData = error $ "Not runtime-representable data"
  
  toDataAtomAssumeNoDecls :: CAtom i -> SimplifyM i o (SAtom o)
  toDataAtomAssumeNoDecls x = do
    Abs decls result <- buildScoped $ toDataAtom x
    case decls of
      Empty -> return result
      _ -> error "unexpected decls"
  
  withSimplifiedBinder
   :: CBinder i i'
   -> (forall o'. DExt o o' => Binder SimpIR o o' -> SimplifyM i' o' a)
   -> SimplifyM i o a
  withSimplifiedBinder (b:>ty) cont = do
    tySimp <- getRepType ty
    tyCore <- substM ty
    withFreshBinder (getNameHint b) tySimp \b' -> do
      x <- liftSimpAtom (sink tyCore) (toAtom $ binderVar b')
      extendSubst (b@>SubstVal x) $ cont b'
  
  withSimplifiedBinders
   :: Nest (Binder CoreIR) o any
   -> (forall o'. DExt o o' => Nest (Binder SimpIR) o o' -> [CAtom o'] -> SimplifyM i o' a)
   -> SimplifyM i o a
  withSimplifiedBinders Empty cont = getDistinct >>= \Distinct -> cont Empty []
  withSimplifiedBinders (Nest (bCore:>ty) bsCore) cont = do
    simpTy <- dropSubst $ getRepType ty
    withFreshBinder (getNameHint bCore) simpTy \bSimp -> do
      x <- liftSimpAtom (sink ty) (toAtom $ binderVar bSimp)
      -- TODO: carry a substitution instead of doing N^2 work like this
      Abs bsCore' UnitE <- applySubst (bCore@>SubstVal x) (EmptyAbs bsCore)
      withSimplifiedBinders bsCore' \bsSimp xs ->
        cont (Nest bSimp bsSimp) (sink x:xs)
  
  -- === Reconstructions ===
  
  data ReconstructAtom (n::S) =
     CoerceRecon (Type CoreIR n)
   | LamRecon (ReconAbs SimpIR CAtom n)
  
  applyRecon :: (EnvReader m, Fallible1 m) => ReconstructAtom n -> SAtom n -> m n (CAtom n)
  applyRecon (CoerceRecon ty) x = liftSimpAtom ty x
  applyRecon (LamRecon ab) x = applyReconAbs ab x
  
  -- === Simplification monad ===
  
  class (ScopableBuilder2 SimpIR m, SubstReader AtomSubstVal m) => Simplifier m
  
  newtype SimplifyM (i::S) (o::S) (a:: *) = SimplifyM
    { runSimplifyM'
      :: SubstReaderT AtomSubstVal (DoubleBuilderT SimpIR TopEnvFrag  HardFailM) i o a }
    deriving ( Functor, Applicative, Monad, ScopeReader, EnvExtender, Fallible
             , EnvReader, SubstReader AtomSubstVal, MonadFail
             , Builder SimpIR, HoistingTopBuilder TopEnvFrag)
  
  liftSimplifyM
    :: (SinkableE e, RenameE e, TopBuilder m, Mut n)
    => (forall l. DExt n l => SimplifyM n l (e l))
    -> m n (e n)
  liftSimplifyM cont = do
    Abs envFrag e <- liftM runHardFail $
      liftDoubleBuilderT $ runSubstReaderT (sink idSubst) $ runSimplifyM' cont
    emitEnv $ Abs envFrag e
  {-# INLINE liftSimplifyM #-}
  
  liftDoubleBuilderToSimplifyM :: DoubleBuilder SimpIR o a -> SimplifyM i o a
  liftDoubleBuilderToSimplifyM cont = SimplifyM $ liftSubstReaderT cont
  
  instance Simplifier SimplifyM
  deriving instance ScopableBuilder SimpIR (SimplifyM i)
  
  -- === Top-level API ===
  
  data SimplifiedTopLam n = SimplifiedTopLam (STopLam n) (ReconstructAtom n)
  data SimplifiedBlock n = SimplifiedBlock (SExpr n) (ReconstructAtom n)
  
  simplifyTopBlock
    :: (TopBuilder m, Mut n) => TopBlock CoreIR n -> m n (SimplifiedTopLam n)
  simplifyTopBlock (TopLam _ _ (LamExpr Empty body)) = do
    SimplifiedBlock block recon <- liftSimplifyM do
      {-# SCC "Simplify" #-} buildSimplifiedBlock $ simplifyExpr body
    topLam <- asTopLam $ LamExpr Empty block
    return $ SimplifiedTopLam topLam recon
  simplifyTopBlock _ = error "not a block (nullary lambda)"
  
  simplifyTopFunction :: (TopBuilder m, Mut n) => CTopLam n -> m n (STopLam n)
  simplifyTopFunction (TopLam False _ f) = do
    asTopLam =<< liftSimplifyM do
      (lam, CoerceReconAbs) <- {-# SCC "Simplify" #-} simplifyLam f
      return lam
  simplifyTopFunction _ = error "shouldn't be in destination-passing style already"
  
  applyReconTop :: (EnvReader m, Fallible1 m) => ReconstructAtom n -> SAtom n -> m n (CAtom n)
  applyReconTop = applyRecon
  
  instance GenericE SimplifiedBlock where
    type RepE SimplifiedBlock = PairE SExpr ReconstructAtom
    fromE (SimplifiedBlock block recon) = PairE block recon
    {-# INLINE fromE #-}
    toE   (PairE block recon) = SimplifiedBlock block recon
    {-# INLINE toE #-}
  
  instance SinkableE SimplifiedBlock
  instance RenameE SimplifiedBlock
  instance HoistableE SimplifiedBlock
  
  instance Pretty (SimplifiedBlock n) where
    pretty (SimplifiedBlock block recon) =
      pretty block <> hardline <> pretty recon
  
  instance SinkableE SimplifiedTopLam where
    sinkingProofE = todoSinkableProof
  
  instance CheckableE SimpIR SimplifiedTopLam where
    checkE (SimplifiedTopLam lam recon) =
      -- TODO: CheckableE instance for the recon too
      SimplifiedTopLam <$> checkE lam <*> renameM recon
  
  instance Pretty (SimplifiedTopLam n) where
    pretty (SimplifiedTopLam lam recon) = pretty lam <> hardline <> pretty recon
  
  -- === All the bits of IR  ===
  
  simplifyDecls :: Emits o => Nest (Decl CoreIR) i i' -> SimplifyM i' o a -> SimplifyM i o a
  simplifyDecls topDecls cont = do
    s  <- getSubst
    s' <- simpDeclsSubst s topDecls
    withSubst s' cont
  {-# INLINE simplifyDecls #-}
  
  simpDeclsSubst
    :: Emits o => Subst AtomSubstVal l o -> Nest (Decl CoreIR) l i'
    -> SimplifyM i o (Subst AtomSubstVal i' o)
  simpDeclsSubst !s = \case
    Empty -> return s
    Nest (Let b (DeclBinding _ expr)) rest -> do
      x <- withSubst s $ simplifyExpr expr
      simpDeclsSubst (s <>> (b@>SubstVal x)) rest
  
  simplifyExpr :: Emits o => Expr CoreIR i -> SimplifyM i o (CAtom o)
  simplifyExpr expr = confuseGHC >>= \_ -> case expr of
    Block _ (Abs decls body) -> simplifyDecls decls $ simplifyExpr body
    App (EffTy _ ty) f xs -> do
      ty' <- substM ty
      f' <- forceConstructor f
      xs' <- mapM simplifyAtom xs
      simplifyApp ty' f' xs'
    TabApp _ f x -> withDistinct do
      x' <- simplifyAtom x
      f' <- forceConstructor f
      simplifyTabApp f' x'
    Atom x -> simplifyAtom x
    PrimOp  op  -> simplifyOp op
    ApplyMethod (EffTy _ ty) dict i xs -> do
      ty' <- substM ty
      xs' <- mapM simplifyAtom xs
      Just dict' <- toMaybeDict <$> simplifyAtom dict
      applyDictMethod ty' dict' i xs'
    TabCon _ ty xs -> do
      ty' <- substM ty
      tySimp <- getRepType ty
      xs' <- forM xs \x -> toDataAtom x
      liftSimpAtom ty' =<< emit (TabCon Nothing tySimp xs')
    Case scrut alts (EffTy _ resultTy) -> do
      scrut' <- simplifyAtom scrut
      resultTy' <- substM resultTy
      defuncCaseCore scrut' resultTy' \i x -> do
        Abs b body <- return $ alts !! i
        extendSubst (b@>SubstVal x) $ simplifyExpr body
    Project ty i x -> do
      ty' <- substM ty
      x'  <- substM x
      tryAsDataAtom x' >>= \case
        Just (x'', _) -> liftSimpAtom ty' =<< proj i x''
        Nothing -> requireReduced $ Project ty' i x'
    Unwrap _ _ -> requireReduced =<< substM expr
  
  requireReduced :: CExpr o -> SimplifyM i o (CAtom o)
  requireReduced expr = reduceExpr expr >>= \case
    Just x -> return x
    Nothing -> error "couldn't reduce expression"
  
  simplifyRefOp :: Emits o => RefOp CoreIR i -> SAtom o -> SimplifyM i o (SAtom o)
  simplifyRefOp op ref = case op of
    MExtend (BaseMonoid em cb) x -> do
      em'  <- toDataAtom em
      x'   <- toDataAtom x
      (cb', CoerceReconAbs) <- simplifyLam cb
      emitRefOp $ MExtend (BaseMonoid em' cb') x'
    MGet   -> emit $ RefOp ref MGet
    MPut x -> do
      x' <- toDataAtom x
      emitRefOp $ MPut x'
    MAsk   -> emitRefOp MAsk
    IndexRef _ x -> do
      x' <- toDataAtom x
      emit =<< mkIndexRef ref x'
    ProjRef _ (ProjectProduct i) -> emit =<< mkProjRef ref (ProjectProduct i)
    ProjRef _ UnwrapNewtype -> return ref
    where emitRefOp op' = emit $ RefOp ref op'
  
  defuncCaseCore :: Emits o
    => Atom CoreIR o -> Type CoreIR o
    -> (forall o'. (Emits o', DExt o o') => Int -> CAtom o' -> SimplifyM i o' (CAtom o'))
    -> SimplifyM i o (CAtom o)
  defuncCaseCore scrut resultTy cont = do
    tryAsDataAtom scrut >>= \case
      Just (scrutSimp, _) -> do
        altBinderTys <- caseAltsBinderTys $ getType scrut
        defuncCase scrutSimp resultTy \i x -> do
          let xCoreTy = altBinderTys !! i
          x' <- liftSimpAtom (sink xCoreTy) x
          cont i x'
      Nothing -> case scrut of
        Con (SumCon _ i arg) -> getDistinct >>= \Distinct -> cont i arg
        _ -> error $ "Don't know how to scrutinize non-data " ++ pprint scrut
  
  defuncCase :: Emits o
    => Atom SimpIR o -> Type CoreIR o
    -> (forall o'. (Emits o', DExt o o') => Int -> SAtom o' -> SimplifyM i o' (CAtom o'))
    -> SimplifyM i o (CAtom o)
  defuncCase scrut resultTy cont = do
    case scrut of
      Con (SumCon _ i arg) -> getDistinct >>= \Distinct -> cont i arg
      Con _ -> error "scrutinee must be a sum type"
      Stuck _ _ -> do
        altBinderTys <- caseAltsBinderTys (getType scrut)
        tryGetRepType resultTy >>= \case
          Just resultTyData -> do
            alts' <- forM (enumerate altBinderTys) \(i, bTy) -> do
              buildAbs noHint bTy \x -> buildBlock do
                ans <- cont i (toAtom $ sink x)
                dropSubst $ toDataAtom ans
            caseExpr <- mkCase scrut resultTyData alts'
            emit caseExpr >>= liftSimpAtom resultTy
          Nothing -> do
            split <- splitDataComponents resultTy
            (alts', closureTys, recons) <- unzip3 <$> forM (enumerate altBinderTys) \(i, bTy) -> do
               simplifyAlt split bTy $ cont i
            let closureSumTy = TyCon $ SumType closureTys
            let newNonDataTy = nonDataTy split
            alts'' <- forM (enumerate alts') \(i, alt) -> injectAltResult closureTys i alt
            caseExpr <- mkCase scrut  (PairTy (dataTy split) closureSumTy) alts''
            caseResult <- emit $ caseExpr
            (dataVal, sumVal) <- fromPair caseResult
            reconAlts <- forM (zip closureTys recons) \(ty, recon) ->
              buildAbs noHint ty \v -> applyRecon (sink recon) (toAtom v)
            nonDataVal <- reduceACase sumVal reconAlts newNonDataTy
            Distinct <- getDistinct
            fromSplit split dataVal nonDataVal
  
  simplifyAlt
    :: SplitDataNonData n
    -> SType o
    -> (forall o'. (Emits o', DExt o o') => SAtom o' -> SimplifyM i o' (CAtom o'))
    -> SimplifyM i o (Alt SimpIR o, SType o, ReconstructAtom o)
  simplifyAlt split ty cont = do
    withFreshBinder noHint ty \b -> do
      ab <- buildScoped $ cont $ sink $ toAtom $ binderVar b
      (body, recon) <- refreshAbs ab \decls result -> do
        let locals = toScopeFrag b >>> toScopeFrag decls
        -- TODO: this might be too cautious. The type only needs to
        -- be hoistable above the decls. In principle it can still
        -- mention vars from the lambda binders.
        Distinct <- getDistinct
        (resultData, resultNonData) <- toSplit split result
        (newResult, reconAbs) <- telescopicCapture locals resultNonData
        return (Abs decls (PairVal resultData newResult), LamRecon reconAbs)
      body' <- mkBlock body
      PairTy _ nonDataType <- return $ getType body'
      let nonDataType' = ignoreHoistFailure $ hoist b nonDataType
      return (Abs b body', nonDataType', recon)
  
  simplifyApp :: Emits o => CType o -> ConcreteCAtom o -> [CAtom o] -> SimplifyM i o (CAtom o)
  simplifyApp resultTy f xs = case f of
    CCCon (WithSubst s con) -> case con of
      Lam (CoreLamExpr _ lam) -> withSubst s $ withInstantiated lam xs \body -> simplifyExpr body
      _ -> error "not a function"
    CCFun ccFun -> case ccFun of
      CCLiftSimpFun _ lam -> do
        xs' <- dropSubst $ mapM toDataAtom xs
        result <- instantiate lam xs' >>= emit
        liftSimpAtom resultTy result
      CCNoInlineFun v _ _ -> simplifyTopFunApp v xs
      CCFFIFun _ f' -> do
        xs' <- dropSubst $ mapM toDataAtom xs
        liftSimpAtom resultTy =<< naryTopApp f' xs'
    CCACase aCase -> forceACase aCase \f' -> simplifyApp (sink resultTy) f' (sink <$> xs)
    CCTabLam _     -> error "not a function"
    CCLiftSimp _ _ -> error "not a function"
  
  simplifyTopFunApp :: Emits n => CAtomVar n -> [CAtom n] -> SimplifyM i n (CAtom n)
  simplifyTopFunApp fName xs = do
    fTy@(TyCon (Pi piTy)) <- return $ getType fName
    resultTy <- typeOfApp fTy xs
    isData resultTy >>= \case
      True -> do
        (xsGeneralized, runtimeArgs) <- generalizeArgs piTy xs
        let spec = AppSpecialization fName xsGeneralized
        Just specializedFunction <- getSpecializedFunction spec >>= emitHoistedEnv
        runtimeArgs' <- dropSubst $ mapM toDataAtom runtimeArgs
        liftSimpAtom resultTy =<< naryTopApp specializedFunction runtimeArgs'
      False ->
        -- TODO: we should probably just fall back to inlining in this case,
        -- especially if we want make everything @noinline by default.
        error $ "can't specialize " ++ pprint fName ++ " " ++ pprint xs
  
  getSpecializedFunction :: EnvReader m => SpecializationSpec n -> m n (Abs TopEnvFrag TopFunName n)
  getSpecializedFunction s = do
    querySpecializationCache s >>= \case
      Just name -> return $ Abs emptyOutFrag name
      Nothing -> liftTopBuilderHoisted $ emitSpecialization (sink s)
  
  emitSpecialization :: (Mut n, TopBuilder m) => SpecializationSpec n -> m n (TopFunName n)
  emitSpecialization s = do
    let hint = getNameHint s
    fCore <- specializedFunCoreDefinition s
    fSimp <- simplifyTopFunction fCore
    name <- emitTopFunBinding hint (Specialization s) fSimp
    extendSpecializationCache s name
    return name
  
  specializedFunCoreDefinition :: (Mut n, TopBuilder m) => SpecializationSpec n -> m n (TopLam CoreIR n)
  specializedFunCoreDefinition (AppSpecialization f (Abs bs staticArgs)) = do
    (asTopLam =<<) $ liftBuilder $ buildLamExpr (EmptyAbs bs) \runtimeArgs -> do
      -- This avoids an infinite loop. Otherwise, in simplifyTopFunction,
      -- where we eta-expand and try to simplify `App f args`, we would see `f` as a
      -- "noinline" function and defer its simplification.
      NoinlineFun _ f' <- lookupAtomName (atomVarName (sink f))
      ListE staticArgs' <- applyRename (bs@@>(atomVarName <$> runtimeArgs)) staticArgs
      naryApp f' staticArgs'
  
  simplifyTabApp ::Emits o => ConcreteCAtom o -> CAtom o -> SimplifyM i o (CAtom o)
  simplifyTabApp f x = case f of
    CCLiftSimp fTy f' -> do
      f'' <- mkStuck f'
      resultTy <- typeOfTabApp fTy x
      x' <- dropSubst $ toDataAtom x
      liftSimpAtom resultTy =<< tabApp f'' x'
    CCACase aCase -> forceACase aCase \f' -> simplifyTabApp f' (sink x)
    CCTabLam (WithSubst s (PairE _ (Abs b ab))) -> do
      x' <- dropSubst $ toDataAtom x
      withSubst s $ extendSubst (b@>(SubstVal x')) $ substM ab
    _ -> error "not a table"
  
  simplifyIxDict :: Dict CoreIR i -> SimplifyM i o (SDict o)
  simplifyIxDict (StuckDict _ stuck) = forceStuck stuck >>= \case
    CCCon (WithSubst s con) -> case con of
      DictConAtom con' -> withSubst s $ simplifyIxDict (DictCon con')
      _ -> error "not a dict"
    CCLiftSimp _ _   -> error "not a dict"
    CCFun      _     -> error "not a dict"
    CCTabLam   _     -> error "not a dict"
    CCACase    _     -> error "not implemented" -- TODO: consider what to do about this
  simplifyIxDict (DictCon con) = case con of
    IxFin n -> DictCon <$> IxRawFin <$> toDataAtomAssumeNoDecls n
    IxRawFin n -> DictCon <$> IxRawFin <$> toDataAtomAssumeNoDecls n
    InstanceDict _ _ _ -> do
      d <- DictCon <$> substM con
      (dictAbs, params) <- generalizeIxDict d
      params' <- dropSubst $ mapM toDataAtomAssumeNoDecls params
      sdName <- requireIxDictCache dictAbs
      return $ DictCon $ IxSpecialized sdName params'
    DataData _ -> error "not an Ix dict"
  
  requireIxDictCache
    :: (HoistingTopBuilder TopEnvFrag m) => AbsDict n -> m n (Name SpecializedDictNameC n)
  requireIxDictCache dictAbs = do
    queryIxDictCache dictAbs >>= \case
      Just name -> return name
      Nothing -> do
        ab <- liftTopBuilderHoisted do
          dName <- emitBinding "d" $ sink $ SpecializedDictBinding $ SpecializedDict dictAbs Nothing
          updateTopEnv $ ExtendCache $ mempty { ixDictCache = eMapSingleton (sink dictAbs) dName }
          methods <- forM [minBound..maxBound] \method -> simplifyDictMethod (sink dictAbs) method
          updateTopEnv $ FinishDictSpecialization dName methods
          return dName
        maybeD <- emitHoistedEnv ab
        case maybeD of
          Just name -> return name
          Nothing -> error "Couldn't hoist specialized dictionary"
  {-# INLINE requireIxDictCache #-}
  
  simplifyDictMethod :: Mut n => AbsDict n -> IxMethod -> TopBuilderM n (TopLam SimpIR n)
  simplifyDictMethod absDict@(Abs bs dict) method = do
    ty <- liftEnvReaderM $ ixMethodType method absDict
    lamExpr <- liftBuilder $ buildTopLamFromPi ty \allArgs -> do
      let (extraArgs, methodArgs) = splitAt (nestLength bs) allArgs
      dict' <- applyRename (bs @@> (atomVarName <$> extraArgs)) dict
      emit =<< mkApplyMethod dict' (fromEnum method) (toAtom <$> methodArgs)
    simplifyTopFunction lamExpr
  
  ixMethodType :: IxMethod -> AbsDict n -> EnvReaderM n (PiType CoreIR n)
  ixMethodType method absDict = do
    refreshAbs absDict \extraArgBs dict -> do
      CorePiType _ _ methodArgs (EffTy _ resultTy) <- getMethodType dict (fromEnum method)
      let allBs = extraArgBs >>> methodArgs
      return $ PiType allBs (EffTy Pure resultTy)
  
  simplifyAtom :: CAtom i -> SimplifyM i o (CAtom o)
  simplifyAtom = substM
  
  -- Assumes first order (args/results are "data", allowing newtypes), monormophic
  simplifyLam
    :: LamExpr CoreIR i
    -> SimplifyM i o (LamExpr SimpIR o, Abs (Nest (AtomNameBinder SimpIR)) ReconstructAtom o)
  simplifyLam (LamExpr bsTop body) = case bsTop of
    Nest b bs -> withSimplifiedBinder b \b'@(b'':>_) -> do
      (LamExpr bs' body', Abs bsRecon recon) <- simplifyLam $ LamExpr bs body
      return (LamExpr (Nest b' bs') body', Abs (Nest b'' bsRecon) recon)
    Empty -> do
      SimplifiedBlock body' recon <- buildSimplifiedBlock $ simplifyExpr body
      return (LamExpr Empty body', Abs Empty recon)
  
  data SplitDataNonData n = SplitDataNonData
    { dataTy    :: Type SimpIR n
    , nonDataTy :: Type CoreIR n
    , toSplit   :: forall i l . CAtom l -> SimplifyM i l (SAtom l, CAtom l)
    , fromSplit :: forall i l . DExt n l => SAtom l -> CAtom l -> SimplifyM i l (CAtom l) }
  
  -- bijection between that type and a (data, non-data) pair type.
  splitDataComponents :: Type CoreIR n -> SimplifyM i n (SplitDataNonData n)
  splitDataComponents = \case
    TyCon (ProdType tys) -> do
      splits <- mapM splitDataComponents tys
      return $ SplitDataNonData
        { dataTy    = TyCon $ ProdType $ map dataTy    splits
        , nonDataTy = TyCon $ ProdType $ map nonDataTy splits
        , toSplit = \xProd -> do
            xs <- getUnpackedReduced xProd
            (ys, zs) <- unzip <$> forM (zip xs splits) \(x, split) -> toSplit split x
            return (Con $ ProdCon ys, Con $ ProdCon zs)
        , fromSplit = \xsProd ysProd -> do
            xs <- getUnpackedReduced xsProd
            ys <- getUnpackedReduced ysProd
            zs <- forM (zip (zip xs ys) splits) \((x, y), split) -> fromSplit split x y
            return $ Con $ ProdCon zs }
    ty -> tryGetRepType ty >>= \case
      Just repTy -> return $ SplitDataNonData
        { dataTy = repTy
        , nonDataTy = UnitTy
        , toSplit = \x -> (,UnitVal) <$> (dropSubst $ toDataAtomAssumeNoDecls x)
        , fromSplit = \x _ -> liftSimpAtom (sink ty) x }
      Nothing -> return $ SplitDataNonData
        { dataTy = UnitTy
        , nonDataTy = ty
        , toSplit = \x -> return (UnitVal, x)
        , fromSplit = \_ x -> return x }
  
  buildSimplifiedBlock
    :: (forall o'. (Emits o', DExt o o') => SimplifyM i o' (CAtom o'))
    -> SimplifyM i o (SimplifiedBlock o)
  buildSimplifiedBlock cont = do
    Abs decls eitherResult <- buildScoped do
      ans <- cont
      tryAsDataAtom ans >>= \case
        Nothing -> return $ LeftE ans
        Just (dataResult, _) -> do
          ansTy <- return $ getType ans
          return $ RightE (dataResult `PairE` ansTy)
    case eitherResult of
      LeftE ans -> do
        (blockAbs, recon) <- refreshAbs (Abs decls ans) \decls' ans' -> do
          (newResult, reconAbs) <- telescopicCapture (toScopeFrag decls') ans'
          return (Abs decls' newResult, LamRecon reconAbs)
        block' <- mkBlock blockAbs
        return $ SimplifiedBlock block' recon
      RightE (ans `PairE` ty) -> do
        let ty' = ignoreHoistFailure $ hoist (toScopeFrag decls) ty
        block <- mkBlock $ Abs decls ans
        return $ SimplifiedBlock block (CoerceRecon ty')
  
  simplifyOp :: Emits o => PrimOp CoreIR i -> SimplifyM i o (CAtom o)
  simplifyOp op = case op of
    Hof (TypedHof (EffTy _ ty) hof) -> do
      ty' <- substM ty
      simplifyHof ty' hof
    MemOp    op' -> simplifyGenericOp op'
    VectorOp op' -> simplifyGenericOp op'
    RefOp ref eff -> do
      ref' <- toDataAtom ref
      liftResult =<< simplifyRefOp eff ref'
    BinOp binop x y -> do
      x' <- toDataAtom x
      y' <- toDataAtom y
      liftResult =<< emit (BinOp binop x' y')
    UnOp unOp x -> do
      x' <- toDataAtom x
      liftResult =<< emit (UnOp unOp x')
    MiscOp op' -> case op' of
      ShowAny x -> do
        x' <- simplifyAtom x
        dropSubst $ showAny x' >>= simplifyExpr
      _ -> simplifyGenericOp op'
    where
      liftResult x = do
        ty <- substM $ getType op
        liftSimpAtom ty x
  
  simplifyGenericOp
    :: (GenericOp op, ToExpr (op SimpIR) SimpIR, HasType CoreIR (op CoreIR), Emits o,
        OpConst op CoreIR ~ OpConst op SimpIR)
    => op CoreIR i
    -> SimplifyM i o (CAtom o)
  simplifyGenericOp op = do
    ty <- substM $ getType op
    op' <- traverseOp op getRepType toDataAtom (error "shouldn't have lambda left")
    liftSimpAtom ty =<< emit op'
  {-# INLINE simplifyGenericOp #-}
  
  pattern CoerceReconAbs :: Abs (Nest b) ReconstructAtom n
  pattern CoerceReconAbs <- Abs _ (CoerceRecon _)
  
  applyDictMethod :: Emits o => CType o -> CDict o -> Int -> [CAtom o] -> SimplifyM i o (CAtom o)
  applyDictMethod resultTy d i methodArgs = case d of
    DictCon (InstanceDict _ instanceName instanceArgs) -> dropSubst do
      instanceArgs' <- mapM simplifyAtom instanceArgs
      instanceDef <- lookupInstanceDef instanceName
      withInstantiated instanceDef instanceArgs' \(PairE _ body) -> do
        let InstanceBody _ methods = body
        let method = methods !! i
        method' <- forceConstructor method
        simplifyApp resultTy method' methodArgs
    DictCon (IxFin n) -> applyIxFinMethod (toEnum i) n methodArgs
    d' -> error $ "Not a simplified dict: " ++ pprint d'
    where
      applyIxFinMethod :: EnvReader m => IxMethod -> CAtom n -> [CAtom n] -> m n (CAtom n)
      applyIxFinMethod method n args = do
        case (method, args) of
          (Size, []) -> return n  -- result : Nat
          (Ordinal, [ix]) -> reduceUnwrap ix -- result : Nat
          (UnsafeFromOrdinal, [ix]) -> return $ toAtom $ NewtypeCon (FinCon n) ix
          _ -> error "bad ix args"
  
  simplifyHof :: Emits o => CType o -> Hof CoreIR i -> SimplifyM i o (CAtom o)
  simplifyHof resultTy = \case
    For d (IxType ixTy ixDict) lam -> do
      (lam', Abs (UnaryNest bIx) recon) <- simplifyLam lam
      ixTy' <- getRepType ixTy
      ixDict' <- simplifyIxDict ixDict
      ans <- emitHof $ For d (IxType ixTy' ixDict') lam'
      case recon of
        CoerceRecon _ -> liftSimpAtom resultTy ans
        LamRecon (Abs bsClosure reconResult) -> do
          ab <- buildAbs noHint ixTy' \i -> do
            xs <- unpackTelescope bsClosure =<< reduceTabApp (sink ans) (toAtom i)
            applySubst (bIx@>Rename (atomVarName i) <.> bsClosure @@> map SubstVal xs) reconResult
          TyCon (TabPi resultTy') <- return resultTy
          mkStuck $ TabLam $ resultTy' `PairE` ab
    While body -> do
      SimplifiedBlock body' (CoerceRecon _) <- buildSimplifiedBlock $ simplifyExpr body
      result <- emitHof $ While body'
      liftSimpAtom resultTy result
    RunReader r lam -> do
      r' <- toDataAtom r
      (lam', Abs b recon) <- simplifyLam lam
      ans <- emitHof $ RunReader r' lam'
      let recon' = ignoreHoistFailure $ hoist b recon
      applyRecon recon' ans
    RunWriter Nothing (BaseMonoid e combine) lam -> do
      LamExpr (BinaryNest h (_:>RefTy _ wTy)) _ <- return lam
      wTy' <- substM $ ignoreHoistFailure $ hoist h wTy
      e' <- toDataAtom e
      (combine', CoerceReconAbs) <- simplifyLam combine
      (lam', Abs b recon) <- simplifyLam lam
      (ans, w) <- fromPair =<< emitHof (RunWriter Nothing (BaseMonoid e' combine') lam')
      let recon' = ignoreHoistFailure $ hoist b recon
      ans' <- applyRecon recon' ans
      w' <- liftSimpAtom wTy' w
      return $ PairVal ans' w'
    RunWriter _ _ _ -> error "Shouldn't see a RunWriter with a dest in Simplify"
    RunState Nothing s lam -> do
      s' <- toDataAtom s
      sTy <- substM $ getType s
      (lam', Abs b recon) <- simplifyLam lam
      resultPair <- emitHof $ RunState Nothing s' lam'
      (ans, sOut) <- fromPair resultPair
      let recon' = ignoreHoistFailure $ hoist b recon
      ans' <- applyRecon recon' ans
      sOut' <- liftSimpAtom sTy sOut
      return $ PairVal ans' sOut'
    RunState _ _ _ -> error "Shouldn't see a RunState with a dest in Simplify"
    RunIO body -> do
      SimplifiedBlock body' recon <- buildSimplifiedBlock $ simplifyExpr body
      ans <- emitHof $ RunIO body'
      applyRecon recon ans
    RunInit body -> do
      SimplifiedBlock body' recon <- buildSimplifiedBlock $ simplifyExpr body
      ans <- emitHof $ RunInit body'
      applyRecon recon ans
    Linearize lam x -> do
      x' <- toDataAtom x
      -- XXX: we're ignoring the result type here, which only makes sense if we're
      -- dealing with functions on simple types.
      (lam', recon) <- simplifyLam lam
      CoerceReconAbs <- return recon
      (result, linFun) <- liftDoubleBuilderToSimplifyM $ linearize lam' x'
      PairTy lamResultTy linFunTy <- return resultTy
      result' <- liftSimpAtom lamResultTy result
      linFun' <- liftSimpFun linFunTy linFun
      return $ PairVal result' linFun'
    Transpose lam x -> do
      (lam', CoerceReconAbs) <- simplifyLam lam
      x' <- toDataAtom x
      result <- transpose lam' x'
      liftSimpAtom resultTy result
    CatchException _ body-> do
      SimplifiedBlock body' recon <- buildSimplifiedBlock $ simplifyExpr body
      block <- liftBuilder $ runSubstReaderT idSubst $ buildBlock $
        exceptToMaybeExpr body'
      result <- emit block
      case recon of
        CoerceRecon ty -> do
          maybeTy <- makePreludeMaybeTy ty
          liftSimpAtom maybeTy result
        LamRecon reconAbs -> fmapMaybe result (applyReconAbs $ sink reconAbs)
  
  -- takes an internal SimpIR Maybe to a CoreIR "prelude Maybe"
  fmapMaybe
    :: SAtom n -> (forall l. DExt n l => SAtom l -> SimplifyM i l (CAtom l))
    -> SimplifyM i n (CAtom n)
  fmapMaybe scrut f = do
    ~(MaybeTy justTy) <- return $ getType scrut
    (justAlt, resultJustTy) <- withFreshBinder noHint justTy \b -> do
      result <- f (toAtom $ binderVar b)
      resultTy <- return $ ignoreHoistFailure $ hoist b (getType result)
      result' <- preludeJustVal result
      return (Abs b result', resultTy)
    nothingAlt <- buildAbs noHint UnitTy \_ -> preludeNothingVal $ sink resultJustTy
    resultMaybeTy <- makePreludeMaybeTy resultJustTy
    reduceACase scrut [nothingAlt, justAlt] resultMaybeTy
  
  -- This is wrong! The correct implementation is below. And yet there's some
  -- compensatory bug somewhere that means that the wrong answer works and the
  -- right answer doesn't. Need to investigate.
  preludeJustVal :: EnvReader m => CAtom n -> m n (CAtom n)
  preludeJustVal x = return x
    -- xTy <- getType x
    -- con <- preludeMaybeNewtypeCon xTy
    -- return $ NewtypeCon con (JustAtom xTy x)
  
  preludeNothingVal :: EnvReader m => CType n -> m n (CAtom n)
  preludeNothingVal ty = do
    con <- preludeMaybeNewtypeCon ty
    return $ Con $ NewtypeCon con (NothingAtom ty)
  
  preludeMaybeNewtypeCon :: EnvReader m => CType n -> m n (NewtypeCon n)
  preludeMaybeNewtypeCon ty = do
    ~(Just (UTyConVar tyConName)) <- lookupSourceMap "Maybe"
    TyConDef sn _ _ _ <- lookupTyCon tyConName
    let params = TyConParams [Explicit] [toAtom ty]
    return $ UserADTData sn tyConName params
  
  liftSimpFun :: EnvReader m => Type CoreIR n -> LamExpr SimpIR n -> m n (CAtom n)
  liftSimpFun (TyCon (Pi piTy)) f = mkStuck $ LiftSimpFun piTy f
  liftSimpFun _ _ = error "not a pi type"
  
  -- === simplifying custom linearizations ===
  
  linearizeTopFun :: (Mut n, Fallible1 m, TopBuilder m) => LinearizationSpec n -> m n (TopFunName n, TopFunName n)
  linearizeTopFun spec = do
    -- TODO: package up this caching pattern so we don't keep reimplementing it
    queryLinearizationCache spec >>= \case
      Just (primal, tangent) -> return (primal, tangent)
      Nothing -> do
        (primal, tangent) <- linearizeTopFunNoCache spec
        extendLinearizationCache spec (primal, tangent)
        return (primal, tangent)
  
  linearizeTopFunNoCache :: (Mut n, TopBuilder m) => LinearizationSpec n -> m n (TopFunName n, TopFunName n)
  linearizeTopFunNoCache spec@(LinearizationSpec f actives) = do
    TopFunBinding ~(DexTopFun _ lam _) <- lookupEnv f
    PairE fPrimal fTangent <- liftSimplifyM $ tryGetCustomRule (sink f) >>= \case
      Just (absParams, rule) -> simplifyCustomLinearization (sink absParams) actives (sink rule)
      Nothing -> liftM toPairE $ liftDoubleBuilderToSimplifyM $ linearizeTopLam (sink lam) actives
    fTangentT <- transposeTopFun fTangent
    fPrimal'   <- emitTopFunBinding "primal"   (LinearizationPrimal  spec) fPrimal
    fTangent'  <- emitTopFunBinding "tangent"  (LinearizationTangent spec) fTangent
    fTangentT' <- emitTopFunBinding "tangent"  (LinearizationTangent spec) fTangentT
    updateTransposeRelation fTangent' fTangentT'
    return (fPrimal', fTangent')
  
  tryGetCustomRule :: EnvReader m => TopFunName n -> m n (Maybe (Abstracted CoreIR (ListE CAtom) n, AtomRules n))
  tryGetCustomRule f' = do
    ~(TopFunBinding f) <- lookupEnv f'
    case f of
      DexTopFun def _ _ -> case def of
        Specialization (AppSpecialization fCore absParams) ->
          fmap (absParams,) <$> lookupCustomRules (atomVarName fCore)
        _ -> return Nothing
      _ -> return Nothing
  
  type Linearized = Abs (Nest SBinder)    -- primal args
                        (Abs (Nest SDecl) -- primal decls
                        (PairE SAtom      -- primal result
                               SLam))     -- tangent function
  
  simplifyCustomLinearization
    :: Abstracted CoreIR (ListE CAtom) n -> [Active] -> AtomRules n
    -> SimplifyM i n (PairE STopLam STopLam n)
  simplifyCustomLinearization (Abs runtimeBs staticArgs) actives rule = do
    CustomLinearize nImplicit nExplicit zeros fCustom <- return rule
    linearized <- withSimplifiedBinders runtimeBs \runtimeBs' runtimeArgs -> do
        Abs runtimeBs' <$> buildScoped do
          ListE staticArgs' <- instantiate (sink $ Abs runtimeBs staticArgs) (sink <$> runtimeArgs)
          fCustom' <- sinkM fCustom
          -- TODO: give a HasType instance to ConcreteCAtom
          resultTy <- typeOfApp (getType fCustom') staticArgs'
          fCustom'' <- dropSubst $ forceConstructor fCustom'
          pairResult <- dropSubst $ simplifyApp resultTy fCustom'' staticArgs'
          (primalResult, fLin) <- fromPairReduced pairResult
          primalResult' <- dropSubst $ toDataAtom primalResult
          let explicitPrimalArgs = drop nImplicit staticArgs'
          allTangentTys <- forM explicitPrimalArgs \primalArg -> do
            tangentType =<< dropSubst (getRepType (getType primalArg))
          let actives' = drop (length actives - nExplicit) actives
          activeTangentTys <- catMaybes <$> forM (zip allTangentTys actives')
            \(t, active) -> return case active of True  -> Just t; False -> Nothing
          fLin' <- buildUnaryLamExpr "t" (toType $ ProdType activeTangentTys) \activeTangentArg -> do
            activeTangentArgs <- getUnpacked $ toAtom activeTangentArg
            ListE allTangentTys' <- sinkM $ ListE allTangentTys
            tangentArgs <- buildTangentArgs zeros (zip allTangentTys' actives') activeTangentArgs
            -- TODO: we're throwing away core type information here. Once we
            -- support core-level tangent types we should make an effort to
            -- correctly restore the core types before applying `fLin`. Right now,
            -- a custom linearization defined for a function on ADTs will
            -- not work.
            fLin' <- sinkM fLin
            TyCon (Pi (CorePiType _ _ bs _)) <- return $ getType fLin'
            let tangentCoreTys = fromNonDepNest bs
            tangentArgs' <- zipWithM liftSimpAtom tangentCoreTys tangentArgs
            resultTyTangent <- typeOfApp (getType fLin') tangentArgs'
            fLin'' <- dropSubst $ forceConstructor fLin'
            tangentResult <- dropSubst $ simplifyApp resultTyTangent fLin'' tangentArgs'
            dropSubst $ toDataAtom tangentResult
          return $ PairE primalResult' fLin'
    PairE primalFun tangentFun <- defuncLinearized linearized
    primalFun' <- asTopLam primalFun
    tangentFun' <- asTopLam tangentFun
    return $ PairE primalFun' tangentFun'
   where
      buildTangentArgs :: Emits n => SymbolicZeros -> [(SType n, Active)] -> [SAtom n] -> SimplifyM i n [SAtom n]
      buildTangentArgs _ [] [] = return []
      buildTangentArgs zeros ((t, False):tys) activeArgs = do
        inactiveArg <- case zeros of
          SymbolicZeros -> symbolicTangentZero t
          InstantiateZeros -> zeroAt t
        rest <- buildTangentArgs zeros tys activeArgs
        return $ inactiveArg:rest
      buildTangentArgs zeros ((_, True):tys) (activeArg:activeArgs) = do
        activeArg' <- case zeros of
          SymbolicZeros -> symbolicTangentNonZero activeArg
          InstantiateZeros -> return activeArg
        rest <- buildTangentArgs zeros tys activeArgs
        return $ activeArg':rest
      buildTangentArgs _ _ _ = error "zip error"
  
      fromNonDepNest :: Nest CBinder n l -> [CType n]
      fromNonDepNest Empty = []
      fromNonDepNest (Nest b bs) =
        case ignoreHoistFailure $ hoist b (Abs bs UnitE) of
          Abs bs' UnitE -> binderType b : fromNonDepNest bs'
  
  defuncLinearized :: EnvReader m => Linearized n -> m n (PairE SLam SLam n)
  defuncLinearized ab = liftBuilder $ refreshAbs ab \bs ab' -> do
    (declsAndResult, reconAbs, residualsTangentsBs) <-
      refreshAbs ab' \decls (PairE primalResult fLin) -> do
        (residuals, reconAbs) <- telescopicCapture (toScopeFrag decls) fLin
        let rTy = getType residuals
        LamExpr tBs _ <- return fLin
        residualsTangentsBs <- withFreshBinder "residual" rTy \rB -> do
          Abs tBs' UnitE <- sinkM $ Abs tBs UnitE
          return $ Abs (Nest rB tBs') UnitE
        residualsTangentsBs' <- return $ ignoreHoistFailure $ hoist decls residualsTangentsBs
        return (Abs decls (PairVal primalResult residuals), reconAbs, residualsTangentsBs')
    primalFun <- LamExpr bs <$> mkBlock declsAndResult
    LamExpr residualAndTangentBs tangentBody <- buildLamExpr residualsTangentsBs \(residuals:tangents) -> do
      LamExpr tangentBs' body <- applyReconAbs (sink reconAbs) (toAtom residuals)
      applyRename (tangentBs' @@> (atomVarName <$> tangents)) body >>= emit
    let tangentFun = LamExpr (bs >>> residualAndTangentBs) tangentBody
    return $ PairE primalFun tangentFun
  
  -- === exception-handling pass ===
  
  type HandlerM = SubstReaderT AtomSubstVal (BuilderM SimpIR)
  
  exceptToMaybeBlock :: Emits o => SType o -> SBlock i -> HandlerM i o (SAtom o)
  exceptToMaybeBlock ty (Abs Empty result) = do
    result' <- exceptToMaybeExpr result
    return $ JustAtom ty result'
  exceptToMaybeBlock resultTy (Abs (Nest (Let b (DeclBinding _ rhs)) decls) finalResult) = do
    maybeResult <- exceptToMaybeExpr rhs
    case maybeResult of
      -- This case is just an optimization (but an important one!)
      JustAtom _ x  ->
        extendSubst (b@> SubstVal x) $ exceptToMaybeBlock resultTy (Abs decls finalResult)
      _ -> emitMaybeCase maybeResult (MaybeTy resultTy)
            (return $ NothingAtom $ sink resultTy)
            (\v -> extendSubst (b@> SubstVal v) $
                     exceptToMaybeBlock (sink resultTy) (Abs decls finalResult))
  
  exceptToMaybeExpr :: Emits o => SExpr i -> HandlerM i o (SAtom o)
  exceptToMaybeExpr expr = case expr of
    Block (EffTy _ ty) body -> do
      ty' <- substM ty
      exceptToMaybeBlock ty' body
    Case e alts (EffTy _ resultTy) -> do
      e' <- substM e
      resultTy' <- substM $ MaybeTy resultTy
      buildCase e' resultTy' \i v -> do
        Abs b body <- return $ alts !! i
        extendSubst (b @> SubstVal v) do
          exceptToMaybeExpr body
    Atom x -> do
      x' <- substM x
      let ty = getType x'
      return $ JustAtom ty x'
    PrimOp (Hof (TypedHof _ (For ann ixTy' (UnaryLamExpr b body)))) -> do
      ixTy <- substM ixTy'
      maybes <- buildFor (getNameHint b) ann ixTy \i -> do
        extendSubst (b@>Rename (atomVarName i)) $ exceptToMaybeExpr body
      catMaybesE maybes
    PrimOp (MiscOp (ThrowException _)) -> do
      ty <- substM $ getType expr
      return $ NothingAtom ty
    PrimOp (Hof (TypedHof _ (RunState Nothing s lam))) -> do
      s' <- substM s
      BinaryLamExpr h ref body <- return lam
      result  <- emitRunState noHint s' \h' ref' ->
        extendSubst (h @> Rename (atomVarName h') <.> ref @> Rename (atomVarName ref')) do
          exceptToMaybeExpr body
      (maybeAns, newState) <- fromPair result
      a <- substM $ getType expr
      emitMaybeCase maybeAns (MaybeTy a)
         (return $ NothingAtom $ sink a)
         (\ans -> return $ JustAtom (sink a) $ PairVal ans (sink newState))
    PrimOp (Hof (TypedHof (EffTy _ resultTy) (RunWriter Nothing monoid (BinaryLamExpr h ref body)))) -> do
      monoid' <- substM monoid
      PairTy _ accumTy <- substM resultTy
      result <- emitRunWriter noHint accumTy monoid' \h' ref' ->
        extendSubst (h @> Rename (atomVarName h') <.> ref @> Rename (atomVarName ref')) do
          exceptToMaybeExpr body
      (maybeAns, accumResult) <- fromPair result
      a <- substM $ getType expr
      emitMaybeCase maybeAns (MaybeTy a)
        (return $ NothingAtom $ sink a)
        (\ans -> return $ JustAtom (sink a) $ PairVal ans (sink accumResult))
    PrimOp (Hof (TypedHof _ (While body))) -> runMaybeWhile $ exceptToMaybeExpr body
    _ -> do
      expr' <- substM expr
      case hasExceptions expr' of
        True -> error $ "Unexpected exception-throwing expression: " ++ pprint expr
        False -> do
          v <- emit expr'
          let ty = getType v
          return $ JustAtom ty v
  
  hasExceptions :: SExpr n -> Bool
  hasExceptions expr = case getEffects expr of
    EffectRow effs NoTail -> ExceptionEffect `eSetMember` effs
  
  -- === instances ===
  
  instance GenericE ReconstructAtom where
    type RepE ReconstructAtom = EitherE2 (Type CoreIR) (ReconAbs SimpIR CAtom)
  
    fromE = \case
      CoerceRecon ty -> Case0 ty
      LamRecon ab    -> Case1 ab
    {-# INLINE fromE #-}
    toE = \case
      Case0 ty -> CoerceRecon ty
      Case1 ab -> LamRecon ab
      _ -> error "impossible"
    {-# INLINE toE #-}
  
  instance SinkableE  ReconstructAtom
  instance HoistableE ReconstructAtom
  instance RenameE    ReconstructAtom
  
  instance Pretty (ReconstructAtom n) where
    pretty (CoerceRecon ty) = "Coercion reconstruction: " <> pretty ty
    pretty (LamRecon ab) = "Reconstruction abs: " <> pretty ab
  
  -- === GHC performance hacks ===
  
  -- Note [Confuse GHC]
  -- I can't explain this, but for some reason using this function in strategic
  -- places makes GHC produce significantly better code. If we define
  --
  -- simplifyAtom = \case
  --   ...
  --   Con con -> traverse simplifyAtom con
  --   ...
  --
  -- then GHC is reluctant to generate a fast-path worker function for simplifyAtom
  -- that would return unboxed tuples, because (at least that's my guess) it's afraid
  -- that it will have to allocate a reader closure for the traverse, which does not
  -- get inlined. For some reason writing the `confuseGHC >>= \_ -> case atom of ...`
  -- makes GHC do the right thing, i.e. generate unboxed worker + a tiny wrapper that
  -- allocates -- a closure to be passed into traverse.
  --
  -- What's so special about this, I don't know. `return ()` is insufficient and doesn't
  -- make the optimization go through. I'll just take the win for now...
  --
  -- NB: We should revise this whenever we upgrade to a newer GHC version.
  confuseGHC :: SimplifyM i o (DistinctEvidence o)
  confuseGHC = getDistinct
  {-# INLINE confuseGHC #-}
  