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
  
  module Linearize (linearize, linearizeTopLam) where
  
  import Control.Category ((>>>))
  import Control.Monad.Reader
  import Data.Functor
  import Data.List (elemIndex)
  import Data.Maybe (catMaybes, isJust)
  import qualified Data.Set as S
  import GHC.Stack
  
  import Builder
  import Core
  import CheapReduction
  import Imp
  import IRVariants
  import MTL1
  import Name
  import Subst
  import {-# SOURCE #-} Simplify (linearizeTopFun)
  import PPrint
  import QueryType
  import Types.Core
  import Types.Primitives
  import Types.Top
  import Util (enumerate)
  
  -- === linearization monad ===
  
  data ActivePrimals (n::S) = ActivePrimals
    { activeVars    :: [AtomVar SimpIR n]  -- includes refs and regions
    , activeEffs    :: EffectRow SimpIR n }
  
  emptyActivePrimals :: ActivePrimals n
  emptyActivePrimals = ActivePrimals [] Pure
  
  data TangentArgs (n::S) = TangentArgs [SAtomVar n]
  
  type PrimalM  = SubstReaderT Name (ReaderT1 ActivePrimals (DoubleBuilder SimpIR)) :: MonadKind2
  type TangentM = ReaderT1 TangentArgs (DoubleBuilder SimpIR) :: MonadKind1
  
  data WithTangent (n::S) (e1::E) (e2::E) =
    WithTangent (e1 n) (forall l. (Emits l, DExt n l) => TangentM l (e2 l))
  
  type LinM i o e1 e2 = PrimalM i o (WithTangent o e1 e2)
  
  -- TODO: maybe we shouldn't roll subst into this
  pureLin :: (RenameE e, SinkableE e) => e i -> LinM i o e e
  pureLin x = do
    x' <- renameM x
    return $ WithTangent x' (sinkM x')
  
  runPrimalMInit :: PrimalM o o a -> DoubleBuilder SimpIR o a
  runPrimalMInit cont = runPrimalM idSubst emptyActivePrimals cont
  
  runPrimalM :: Subst Name i o -> ActivePrimals o -> PrimalM i o a -> DoubleBuilder SimpIR o a
  runPrimalM subst args cont = runReaderT1 args $ runSubstReaderT subst cont
  
  activePrimalIdx :: SAtomVar o -> PrimalM i o (Maybe Int)
  activePrimalIdx v = asks \primals -> elemIndex v (activeVars primals)
  
  getActivePrimals :: PrimalM i o (ActivePrimals o)
  getActivePrimals = ask
  
  extendActiveSubst
    :: BindsAtMostOneName b (AtomNameC SimpIR)
    => b i i' -> SAtomVar o -> PrimalM i' o a -> PrimalM i o a
  extendActiveSubst b v cont = extendSubst (b@>atomVarName v) $ extendActivePrimals v cont
  
  extendActiveEffs :: Effect SimpIR o -> PrimalM i o a -> PrimalM i o a
  extendActiveEffs eff = local \primals ->
    primals { activeEffs = extendEffRow (eSetSingleton eff) (activeEffs primals)}
  
  extendActivePrimals :: SAtomVar o -> PrimalM i o a -> PrimalM i o a
  extendActivePrimals v = extendActivePrimalss [v]
  
  extendActivePrimalss :: [SAtomVar o] -> PrimalM i o a -> PrimalM i o a
  extendActivePrimalss vs =
    local \primals -> primals { activeVars = activeVars primals ++ vs }
  
  getTangentArg :: Int -> TangentM o (Atom SimpIR o)
  getTangentArg idx = asks \(TangentArgs vs) -> toAtom $ vs !! idx
  
  extendTangentArgs :: SAtomVar n -> TangentM n a -> TangentM n a
  extendTangentArgs v m = local (\(TangentArgs vs) -> TangentArgs $ vs ++ [v]) m
  
  extendTangentArgss :: [SAtomVar n] -> TangentM n a -> TangentM n a
  extendTangentArgss vs' m = local (\(TangentArgs vs) -> TangentArgs $ vs ++ vs') m
  
  getTangentArgs :: TangentM o (TangentArgs o)
  getTangentArgs = ask
  
  emitBoth
    :: (Emits o, ToExpr e' SimpIR)
    => WithTangent o e e
    -> (forall o' m. (DExt o o', Builder SimpIR m) => e o' -> m o' (e' o'))
    -> LinM i o SAtom SAtom
  emitBoth (WithTangent x tx) f = do
    Distinct <- getDistinct
    x' <- emit =<< f x
    return $ WithTangent x' do
      tx' <- tx
      emitLin =<< f tx'
  
  emitZeroT :: (Emits o, HasNamesE e', ToExpr e' SimpIR) => e' i -> LinM i o SAtom SAtom
  emitZeroT e = do
    x <- emit =<< renameM e
    return $ WithTangent x (zeroLikeT x)
  
  zeroLikeT :: (DExt o o', Emits o', HasType SimpIR e) => e o -> TangentM o' (SAtom o')
  zeroLikeT x = do
    ty <- sinkM $ getType x
    zeroAt =<< tangentType ty
  
  fmapLin
    :: Emits o
    => (forall o'. e o' -> e' o')
    -> LinM i o e  e
    -> LinM i o e' e'
  fmapLin f m = do
    WithTangent ans tx <- m
    return $ WithTangent (f ans) (f <$> tx)
  
  zipLin :: WithTangent o e1 e1 -> WithTangent o e2 e2 -> WithTangent o (PairE e1 e2) (PairE e1 e2)
  zipLin (WithTangent x1 t1) (WithTangent x2 t2) = WithTangent (PairE x1 x2) do PairE <$> t1 <*> t2
  
  seqLin
    :: Traversable f
    => f (LinM i o e e)
    -> LinM i o (ComposeE f e) (ComposeE f e)
  seqLin ms = do
    ms' <- sequence ms
    let xs = ms' <&> \(WithTangent x _) -> x
    return $ WithTangent (ComposeE xs) do
      ComposeE <$> forM ms' \(WithTangent _ t) -> t
  
  liftTangentM :: TangentArgs o -> TangentM o a -> PrimalM i o a
  liftTangentM args m = liftSubstReaderT $ lift11 $ runReaderT1 args m
  
  isTrivialForAD :: SExpr o -> PrimalM i o Bool
  isTrivialForAD expr = do
    trivialTy  <- presentAnd isSingletonType <$> (maybeTangentType $ getType expr)
    hasActiveEffs <- case getEffects expr of
                       Pure -> return False
                       -- TODO: Be more precise here, such as checking
                       -- whether the effects are themselves active.
                       _ -> return True
    hasActiveVars <- isActive expr
    return $ not hasActiveEffs && (trivialTy || not hasActiveVars)
      where presentAnd :: (a -> Bool) -> Maybe a -> Bool
            presentAnd = any
  
  isActive :: HoistableE e => e o -> PrimalM i o Bool
  isActive e = do
    vs <- (S.fromList . map atomVarName . activeVars) <$> getActivePrimals
    return $ any (`S.member` vs) (freeAtomVarsList e)
  
  -- === converision between monadic and reified version of functions ===
  
  tangentFunAsLambda
    :: Emits o
    => (forall o'. (DExt o o', Emits o') => TangentM o' (Atom SimpIR o'))
    -> PrimalM i o (SLam o)
  tangentFunAsLambda cont = do
    ActivePrimals primalVars _ <- getActivePrimals
    tangentTys <- getTangentArgTys primalVars
    buildLamExpr tangentTys \tangentVars -> do
      liftTangentM (TangentArgs $ map sink tangentVars) cont
  
  getTangentArgTys :: (Fallible1 m, EnvExtender m) => [SAtomVar n] -> m n (EmptyAbs (Nest SBinder) n)
  getTangentArgTys topVs = go mempty topVs where
    go :: (Fallible1 m, EnvExtender m)
       => EMap SAtomName SAtomVar n -> [SAtomVar n] -> m n (EmptyAbs (Nest SBinder) n)
    go _ [] = return $ EmptyAbs Empty
    go heapMap (v:vs) = case getType v of
      -- This is a hack to handle heaps/references. They normally come in pairs
      -- like this, but there's nothing to prevent users writing programs that
      -- sling around heap variables by themselves. We should try to do something
      -- better...
      TyCon HeapType -> do
        withFreshBinder (getNameHint v) (TyCon HeapType) \hb -> do
          let newHeapMap = sink heapMap <> eMapSingleton (sink (atomVarName v)) (binderVar hb)
          Abs bs UnitE <- go newHeapMap $ sinkList vs
          return $ EmptyAbs $ Nest hb bs
      RefTy (Stuck _ (Var h)) referentTy -> do
        case lookupEMap heapMap (atomVarName h) of
          Nothing -> error "shouldn't happen?"
          Just h' -> do
            tt <- tangentType referentTy
            let refTy = RefTy (toAtom h') tt
            withFreshBinder (getNameHint v) refTy \refb -> do
              Abs bs UnitE <- go (sink heapMap) $ sinkList vs
              return $ EmptyAbs $ Nest refb bs
      ty -> do
        tt <- tangentType ty
        withFreshBinder (getNameHint v) tt \b -> do
          Abs bs UnitE <- go (sink heapMap) $ sinkList vs
          return $ EmptyAbs $ Nest b bs
  
  class ReconFunctor (f :: E -> E) where
    capture
      :: (EnvReader m, HoistableE e, HoistableB b)
      => b n l -> SAtom l -> e l ->  m l (SAtom l, f e n)
    reconstruct
      :: (SubstE AtomSubstVal e, SinkableE e, Emits n)
      => SAtom n -> f e n -> PrimalM i n (SAtom n, e n)
  
  -- tangent lambda
  type LinLam = SLam
  -- tangent lambda prefixed by residual args
  type LinLamAbs = MaybeReconAbs LinLam
  
  data MaybeReconAbs (e::E) (n::S) =
     ReconWithData (ReconAbs SimpIR e n)
   | TrivialRecon (e n)
  
  data ObligateReconAbs (e::E) (n::S) =
     ObligateRecon (SType n) (ReconAbs SimpIR e n)
  
  instance ReconFunctor MaybeReconAbs where
    capture locals original toCapture = do
      (reconVal, recon) <- telescopicCapture locals toCapture
      case recon of
        Abs (ReconBinders _ Empty) toCapture' -> return (original, TrivialRecon toCapture')
        _ -> return (PairVal original reconVal, ReconWithData recon)
  
    reconstruct primalAux recon = case recon of
      TrivialRecon linLam -> return (primalAux, linLam)
      ReconWithData reconAbs -> do
        (primal, residuals) <- fromPair primalAux
        linLam' <- applyReconAbs reconAbs residuals
        return (primal, linLam')
  
  instance ReconFunctor ObligateReconAbs where
    capture locals original toCapture = do
      (reconVal, recon) <- telescopicCapture locals toCapture
      -- TODO: telescopicCapture should probably return the hoisted type
      reconValTy <- return $ ignoreHoistFailure $ hoist locals $ getType reconVal
      return (PairVal original reconVal, ObligateRecon reconValTy recon)
  
    reconstruct primalAux recon = case recon of
      ObligateRecon _ reconAbs -> do
        (primal, residuals) <- fromPair primalAux
        linLam' <- applyReconAbs reconAbs residuals
        return (primal, linLam')
  
  linearizeExprDefunc :: SExpr i -> PrimalM i o (SExpr o, LinLamAbs o)
  linearizeExprDefunc = linearizeExprDefuncGeneral emptyOutFrag
  
  linearizeExprDefuncGeneral
    :: ReconFunctor f
    => ScopeFrag o' o -> SExpr i -> PrimalM i o (SExpr o, f SLam o')
  linearizeExprDefuncGeneral locals expr = do
    Abs decls result <- buildScoped do
      WithTangent primalResult tangentFun <- linearizeExpr expr
      lam <- tangentFunAsLambda tangentFun
      return $ PairE primalResult lam
    (Abs decls' result', recon) <- refreshAbs (Abs decls result) \decls' (PairE primal lam) -> do
      (primal', recon) <- capture (locals >>> toScopeFrag decls') primal lam
      return (Abs decls' primal', recon)
    block <- mkBlock (Abs decls' result')
    return (block, recon)
  
  -- Inverse of tangentFunAsLambda. Should be used inside a returned tangent action.
  applyLinLam :: Emits o => SLam i -> SubstReaderT AtomSubstVal TangentM i o (Atom SimpIR o)
  applyLinLam (LamExpr bs body) = do
    TangentArgs args <- liftSubstReaderT $ getTangentArgs
    extendSubst (bs @@> ((Rename . atomVarName) <$> args)) do
      substM body >>= emitLin
  
  -- === actual linearization passs ===
  
  -- main API entrypoint
  linearize :: Emits n => SLam n -> SAtom n -> DoubleBuilder SimpIR n (SAtom n, SLam n)
  linearize f x = runPrimalMInit $ linearizeLambdaApp f x
  {-# SCC linearize #-}
  
  linearizeTopLam :: STopLam n -> [Active] -> DoubleBuilder SimpIR n (STopLam n, STopLam n)
  linearizeTopLam (TopLam False _ (LamExpr bs body)) actives = do
    (primalFun, tangentFun) <- runPrimalMInit $ refreshBinders bs \bs' frag -> extendSubst frag do
      let allPrimals = bindersVars bs'
      activeVs <- catMaybes <$> forM (zip actives allPrimals) \(active, v) -> case active of
        True  -> return $ Just v
        False -> return $ Nothing
      (body', linLamAbs) <- extendActivePrimalss activeVs do
        linearizeExprDefuncGeneral emptyOutFrag body
      let primalFun = LamExpr bs' body'
      ObligateRecon ty (Abs bsRecon (LamExpr bsTangent tangentBody)) <- return linLamAbs
      tangentFun <- withFreshBinder "residuals" ty \bResidual -> do
        xs <- unpackTelescope bsRecon $ toAtom $ binderVar bResidual
        Abs bsTangent' UnitE <- applySubst (bsRecon @@> map SubstVal xs) (Abs bsTangent UnitE)
        tangentTy <- TyCon <$> ProdType <$> typesFromNonDepBinderNest bsTangent'
        withFreshBinder "t" tangentTy \bTangent -> do
          tangentBody' <- buildBlock do
            ts <- getUnpacked $ toAtom $ sink $ binderVar bTangent
            let substFrag =   bsRecon   @@> map (SubstVal . sink) xs
                          <.> bsTangent @@> map (SubstVal . sink) ts
            emitLin =<< applySubst substFrag tangentBody
          return $ LamExpr (bs' >>> BinaryNest bResidual bTangent) tangentBody'
      return (primalFun, tangentFun)
    (,) <$> asTopLam primalFun <*> asTopLam tangentFun
  linearizeTopLam (TopLam True _ _) _ = error "expected a non-destination-passing function"
  
  -- reify the tangent builder as a lambda
  linearizeLambdaApp :: Emits o => SLam i -> SAtom o -> PrimalM i o (SAtom o, SLam o)
  linearizeLambdaApp (UnaryLamExpr b body) x = do
    vp <- emitToVar x
    extendActiveSubst b vp do
      WithTangent primalResult tangentAction <- linearizeExpr body
      tanFun <- tangentFunAsLambda tangentAction
      return (primalResult, tanFun)
  linearizeLambdaApp _ _ = error "not implemented"
  
  linearizeAtom :: Emits o => Atom SimpIR i -> LinM i o SAtom SAtom
  linearizeAtom (Con con) = linearizePrimCon con
  linearizeAtom (Stuck _ stuck) = linearizeStuck stuck
  
  linearizeStuck :: Emits o => Stuck SimpIR i -> LinM i o SAtom SAtom
  linearizeStuck stuck = case stuck of
    Var v -> do
      v' <- renameM v
      activePrimalIdx v' >>= \case
        Nothing -> zero
        Just idx -> return $ WithTangent (toAtom v') $ getTangentArg idx
    PtrVar _ _   -> zero
    RepValAtom _ -> zero
    -- TODO: de-dup with the Expr versions of these
    StuckProject i x -> do
      x' <- linearizeStuck x
      emitBoth x' \x'' -> mkProject i x''
    StuckTabApp x i -> do
      pt <- zipLin <$> linearizeStuck x <*> pureLin i
      emitBoth pt \(PairE x' i') -> mkTabApp x' i'
    where
      zero = do
        atom <- mkStuck =<< renameM stuck
        return $ WithTangent atom (zeroLikeT atom)
  
  linearizeDecls :: Emits o => Nest SDecl i i' -> LinM i' o e1 e2 -> LinM i  o e1 e2
  linearizeDecls Empty cont = cont
  -- TODO: as an optimization, don't bother extending the tangent args if the
  -- tangent is trivial, either because its type is a singleton or because there
  -- are no active vars.
  linearizeDecls (Nest (Let b (DeclBinding ann expr)) rest) cont = do
    expr' <- renameM expr
    isTrivialForAD expr' >>= \case
      True -> do
        v <- emitToVar expr'
        extendSubst (b@>atomVarName v) $ linearizeDecls rest cont
      False -> do
        WithTangent p tf <- linearizeExpr expr
        v <- emitDecl (getNameHint b) ann (Atom p)
        extendActiveSubst b v do
          WithTangent pRest tfRest <- linearizeDecls rest cont
          return $ WithTangent pRest do
            t <- tf
            vt <- emitDecl (getNameHint b) LinearLet (Atom t)
            extendTangentArgs vt $
              tfRest
  
  linearizeExpr :: Emits o => SExpr i -> LinM i o SAtom SAtom
  linearizeExpr expr = case expr of
    Atom x -> linearizeAtom x
    Block _ (Abs decls result) -> linearizeDecls decls $ linearizeExpr result
    TopApp _ f xs -> do
      (xs', ts) <- unzip <$> forM xs \x -> do
        x' <- renameM x
        isActive x' >>= \case
          True  -> do
            WithTangent x'' t <- dropSubst $ linearizeAtom x'
            return (x'', Just (WithTangent (unitLike x'') t))
          False -> return (x', Nothing)
      f' <- renameM f
      -- TODO(dougalm): this works, but I think that what we really want here is
      -- to hoist the argument to `linearizeTopFun`, rather than the result. We
      -- want to pop all the way up to the top level, hoisting the E-kinded
      -- `LinearizationSpec` with us, rather than working underneath all the local
      -- bindings and then only hoisting the final result.
      Just (PairE fPrimal fTan) <- liftTopBuilderAndEmit $
         liftM toPairE $ linearizeTopFun (sink $ LinearizationSpec f' (map isJust ts))
      (ans, residuals) <- fromPair =<< naryTopApp fPrimal xs'
      return $ WithTangent ans do
        ts' <- forM (catMaybes ts) \(WithTangent UnitE t) -> t
        naryTopApp (sink fTan) (sinkList xs' ++ [sink residuals, Con $ ProdCon ts'])
      where
        unitLike :: e n -> UnitE n
        unitLike _ = UnitE
    PrimOp op    -> linearizeOp op
    Case e alts (EffTy effs resultTy) -> do
      e' <- renameM e
      effs' <- renameM effs
      resultTy' <- renameM resultTy
      isActive e' >>= \case
        True -> notImplemented
        False -> do
          (alts', recons) <- unzip <$> buildCaseAlts e' \i b' -> do
            Abs b body <- return $ alts !! i
            extendSubst (b@>binderName b') do
              (block, recon) <- linearizeExprDefuncGeneral (toScopeFrag b') body
              return (Abs b' block, recon)
          let tys = recons <&> \(ObligateRecon t _) -> t
          alts'' <- forM (enumerate alts') \(i, alt) -> do
            injectAltResult tys i alt
          let fullResultTy = PairTy resultTy' $ TyCon $ SumType tys
          result <- emit $ Case e' alts'' (EffTy effs' fullResultTy)
          (primal, residualss) <- fromPair result
          resultTangentType <- tangentType resultTy'
          return $ WithTangent primal do
            emitLin =<< buildCase' (sink residualss) (sink resultTangentType) \i residuals -> do
              ObligateRecon _ (Abs bs linLam) <- return $ sinkList recons !! i
              residuals' <- unpackTelescope bs residuals
              withSubstReaderT $ extendSubst (bs @@> (SubstVal <$> residuals')) do
                applyLinLam linLam
    TabCon _ ty xs -> do
      ty' <- renameM ty
      pt <- seqLin (map linearizeAtom xs)
      emitBoth pt \(ComposeE xs') -> return $ TabCon Nothing (sink ty') xs'
    TabApp _ x i -> do
      pt <- zipLin <$> linearizeAtom x <*> pureLin i
      emitBoth pt \(PairE x' i') -> mkTabApp x' i'
    Project _ i x -> do
      x' <- linearizeAtom x
      emitBoth x' \x'' -> mkProject i x''
  
  linearizeOp :: Emits o => PrimOp SimpIR i -> LinM i o SAtom SAtom
  linearizeOp op = case op of
    Hof (TypedHof _ e) -> linearizeHof e
    DAMOp _        -> error "shouldn't occur here"
    RefOp ref m -> do
      ref' <- linearizeAtom ref
      case m of
        MAsk -> emitBoth ref' \ref'' -> return $ RefOp ref'' MAsk
        MExtend monoid x -> do
          -- TODO: check that we're dealing with a +/0 monoid
          monoid' <- renameM monoid
          x' <- linearizeAtom x
          emitBoth (zipLin ref' x') \(PairE ref'' x'') ->
            return $ RefOp ref'' $ MExtend (sink monoid') x''
        MGet -> emitBoth ref' \ref'' -> return $ RefOp ref'' MGet
        MPut x -> do
          x' <- linearizeAtom x
          emitBoth (zipLin ref' x') \(PairE ref'' x'') -> return $ RefOp ref'' $ MPut x''
        IndexRef _ i -> do
          i' <- pureLin i
          emitBoth (zipLin ref' i') \(PairE ref'' i'') -> mkIndexRef ref'' i''
        ProjRef _ i -> emitBoth ref' \ref'' -> mkProjRef ref'' i
    UnOp  uop x       -> linearizeUnOp  uop x
    BinOp bop x y     -> linearizeBinOp bop x y
    -- XXX: This assumes that pointers are always constants
    MemOp _      -> emitZeroT op
    MiscOp miscOp -> linearizeMiscOp miscOp
    VectorOp _ -> error "not implemented"
  
  linearizeMiscOp :: Emits o => MiscOp SimpIR i -> LinM i o SAtom SAtom
  linearizeMiscOp op = case op of
    SumTag _     -> zero
    ToEnum _ _   -> zero
    Select p t f -> do
      p' <- pureLin p
      t' <- linearizeAtom t
      f' <- linearizeAtom f
      emitBoth (p' `zipLin` t' `zipLin` f')
         \(p'' `PairE` t'' `PairE` f'') -> return $ Select p'' t'' f''
    CastOp t v -> do
      vt <- getType <$> renameM v
      t' <- renameM t
      vtTangentType <- tangentType vt
      tTangentType  <- tangentType t'
      ((&&) <$> (vtTangentType `alphaEq` vt)
            <*> (tTangentType  `alphaEq` t')) >>= \case
        True -> do
          v' <- linearizeAtom v
          emitBoth v' \v'' -> return $ CastOp (sink t') v''
        False -> do
          WithTangent x xt <- linearizeAtom v
          yt <- case (vtTangentType, tTangentType) of
            (_     , UnitTy) -> return $ UnitVal
            (UnitTy, tt    ) -> zeroAt tt
            _                -> error "Expected at least one side of the CastOp to have a trivial tangent type"
          y <- emit $ CastOp t' x
          return $ WithTangent y do xt >> return (sink yt)
    BitcastOp _ _    -> notImplemented
    UnsafeCoerce _ _ -> notImplemented
    GarbageVal _     -> notImplemented
    ThrowException _ -> notImplemented
    ThrowError _     -> zero
    OutputStream     -> zero
    ShowAny _ -> error "Shouldn't have ShowAny in simplified IR"
    ShowScalar _ -> error "Shouldn't have ShowScalar in simplified IR"
    where zero = emitZeroT op
  
  linearizeUnOp :: Emits o => UnOp -> Atom SimpIR i -> LinM i o SAtom SAtom
  linearizeUnOp op x'' = do
    WithTangent x' tx' <- linearizeAtom x''
    ans' <- emit $ UnOp op x'
    return $ WithTangent ans' do
      ans <- sinkM ans'
      x <- sinkM x'
      tx <- tx'
      let zero = zeroLikeT ans
      case op of
        Exp   -> emitLin $ BinOp FMul tx ans
        Exp2  -> notImplemented
        Log   -> emitLin $ BinOp FDiv tx x
        Log2  -> notImplemented
        Log10 -> notImplemented
        Log1p -> notImplemented
        Sin -> do
          c <- emit $ UnOp Cos x
          emitLin $ BinOp FMul tx c
        Cos -> do
          c <- emit =<< (UnOp FNeg <$> emit (UnOp Sin x))
          emitLin $ BinOp FMul tx c
        Tan -> notImplemented
        Sqrt -> do
          denominator <- fmul (2 `fLitLike` ans) ans
          emitLin $ BinOp FDiv tx denominator
        Floor  -> zero
        Ceil   -> zero
        Round  -> zero
        LGamma -> notImplemented
        Erf    -> notImplemented
        Erfc   -> notImplemented
        FNeg   -> emitLin $ UnOp FNeg tx
        BNot   -> zero
  
  linearizeBinOp :: Emits o => BinOp -> SAtom i -> SAtom i -> LinM i o SAtom SAtom
  linearizeBinOp op x'' y'' = do
    WithTangent x' tx' <- linearizeAtom x''
    WithTangent y' ty' <- linearizeAtom y''
    ans' <- emit $ BinOp op x' y'
    return $ WithTangent ans' do
      ans <- sinkM ans'
      x <- referToPrimal x'
      y <- referToPrimal y'
      tx <- tx'
      ty <- ty'
      let zero = zeroLikeT ans
      case op of
        IAdd   -> zero
        ISub   -> zero
        IMul   -> zero
        IDiv   -> zero
        IRem   -> zero
        ICmp _ -> zero
        FAdd -> emitLin $ BinOp FAdd tx ty
        FSub -> emitLin $ BinOp FSub tx ty
        FMul -> do
          t1 <- emitLin $ BinOp FMul ty x
          t2 <- emitLin $ BinOp FMul tx y
          emitLin $ BinOp FAdd t1 t2
        FDiv -> do
          t1  <- emitLin $ BinOp FDiv tx y
          xyy  <- fdiv x =<< fmul y y
          t2  <- emitLin $ BinOp FMul ty xyy
          emitLin $ BinOp FSub t1 t2
        FPow -> do
          ym1 <- fsub y (1.0 `fLitLike` y)
          xpowym1 <- emit $ BinOp FPow x ym1
          xlogx <- fmul x =<< emit (UnOp Log x)
          t1 <- emitLin $ BinOp FMul tx y
          t2 <- emitLin $ BinOp FMul ty xlogx
          t12 <- emitLin $ BinOp FAdd t1 t2
          emitLin $ BinOp FMul xpowym1 t12
        FCmp _ -> zero
        BAnd   -> zero
        BOr    -> zero
        BXor   -> zero
        BShL   -> zero
        BShR   -> zero
  
  -- This has the same type as `sinkM` and falls back thereto, but recomputes
  -- indexing a primal array in the tangent to avoid materializing intermediate
  -- results thereof.  We should probably have a more cogent story for
  -- rematerialization, but this suffices to remove embarrassing intermediates in
  -- matrix multiplication.
  referToPrimal :: (Builder SimpIR m, Emits l, DExt n l) => SAtom n -> m l (SAtom l)
  referToPrimal x = do
    case x of
      Stuck _ (Var v) -> lookupEnv (atomVarName $ sink v) >>= \case
        AtomNameBinding (LetBound (DeclBinding PlainLet (Atom atom))) ->
          referToPrimal atom
        AtomNameBinding (LetBound (DeclBinding PlainLet (TabApp _ tab i))) -> do
          tab' <- referToPrimal tab
          i' <- referToPrimal i
          emit =<< mkTabApp tab' i'
        _ -> sinkM x
      _ -> sinkM x
  
  linearizePrimCon :: Emits o => Con SimpIR i -> LinM i o SAtom SAtom
  linearizePrimCon con = case con of
    Lit _ -> zero
    ProdCon xs -> fmapLin (Con . ProdCon . fromComposeE) $ seqLin (fmap linearizeAtom xs)
    SumCon  _ _ _ -> notImplemented
    HeapVal -> zero
    DepPair _ _ _     -> notImplemented
    where zero = emitZeroT con
  
  linearizeHof :: Emits o => Hof SimpIR i -> LinM i o SAtom SAtom
  linearizeHof hof = case hof of
    For d ixTy' lam -> do
      UnaryLamExpr ib body <- return lam
      ixTy <- renameM ixTy'
      (lam', Abs ib' linLam) <- withFreshBinder noHint (ixTypeType ixTy) \ib' -> do
        (block', linLam) <- extendSubst (ib@>binderName ib') $ linearizeExprDefunc body
        return (UnaryLamExpr ib' block', Abs ib' linLam)
      primalsAux <- emitHof $ For d ixTy lam'
      case linLam of
        TrivialRecon linLam' ->
          return $ WithTangent primalsAux do
            Abs ib'' linLam'' <- sinkM (Abs ib' linLam')
            withSubstReaderT $ emitLin =<< mkFor noHint d (sink ixTy) \i' -> do
              extendSubst (ib''@>Rename (atomVarName i')) $ applyLinLam linLam''
        ReconWithData reconAbs -> do
          primals <- buildMap primalsAux getFst
          return $ WithTangent primals do
            Abs ib'' (Abs bs linLam') <- sinkM (Abs ib' reconAbs)
            withSubstReaderT $ emitLin =<< mkFor noHint d (sink ixTy) \i' -> do
              extendSubst (ib''@> Rename (atomVarName i')) do
                residuals' <- tabApp (sink primalsAux) (toAtom i') >>= getSnd >>= unpackTelescope bs
                extendSubst (bs @@> (SubstVal <$> residuals')) $
                  applyLinLam linLam'
    RunReader r lam -> do
      WithTangent r' rLin <- linearizeAtom r
      (lam', recon) <- linearizeEffectFun Reader lam
      primalAux <- emitHof $ RunReader r' lam'
      referentTy <- renameM $ getType r
      (primal, linLam) <- reconstruct primalAux recon
      return $ WithTangent primal do
        rLin' <- rLin
        tt <- tangentType $ sink referentTy
        tanEffLam <- buildEffLam noHint tt \h ref ->
          extendTangentArgss [h, ref] do
            withSubstReaderT $ applyLinLam $ sink linLam
        emitHofLin $ RunReader rLin' tanEffLam
    RunState Nothing sInit lam -> do
      WithTangent sInit' sLin <- linearizeAtom sInit
      (lam', recon) <- linearizeEffectFun State lam
      (primalAux, sFinal) <- fromPair =<< emitHof (RunState Nothing sInit' lam')
      referentTy <- snd <$> getTypeRWSAction lam'
      (primal, linLam) <- reconstruct primalAux recon
      return $ WithTangent (PairVal primal sFinal) do
        sLin' <- sLin
        tt <- tangentType $ sink referentTy
        tanEffLam <- buildEffLam noHint tt \h ref ->
          extendTangentArgss [h, ref] do
            withSubstReaderT $ applyLinLam $ sink linLam
        emitHofLin $ RunState Nothing sLin' tanEffLam
    RunWriter Nothing bm lam -> do
      -- TODO: check it's actually the 0/+ monoid (or should we just build that in?)
      bm' <- renameM bm
      (lam', recon) <- linearizeEffectFun Writer lam
      (primalAux, wFinal) <- fromPair =<< emitHof (RunWriter Nothing bm' lam')
      (primal, linLam) <- reconstruct primalAux recon
      referentTy <- snd <$> getTypeRWSAction lam'
      return $ WithTangent (PairVal primal wFinal) do
        bm'' <- sinkM bm'
        tt <- tangentType $ sink referentTy
        tanEffLam <- buildEffLam noHint tt \h ref ->
          extendTangentArgss [h, ref] do
            withSubstReaderT $ applyLinLam $ sink linLam
        emitHofLin $ RunWriter Nothing bm'' tanEffLam
    RunIO body -> do
      (body', recon) <- linearizeExprDefunc body
      primalAux <- emitHof $ RunIO body'
      (primal, linLam) <- reconstruct primalAux recon
      return $ WithTangent primal do
        withSubstReaderT $ applyLinLam $ sink linLam
    _ -> error $ "not implemented: " ++ pprint hof
  
  linearizeEffectFun :: RWS -> SLam i -> PrimalM i o (SLam o, LinLamAbs o)
  linearizeEffectFun rws (BinaryLamExpr hB refB body) = do
    withFreshBinder noHint (TyCon HeapType) \h -> do
      bTy <- extendSubst (hB@>binderName h) $ renameM $ binderType refB
      withFreshBinder noHint bTy \b -> do
        let ref = binderVar b
        hVar <- sinkM $ binderVar h
        (body', linLam) <- extendActiveSubst hB hVar $ extendActiveSubst refB ref $
          -- TODO: maybe we should check whether we need to extend the active effects
          extendActiveEffs (RWSEffect rws (toAtom hVar)) do
            linearizeExprDefunc body
        -- TODO: this assumes that references aren't returned. Our type system
        -- ensures that such references can never be *used* once the effect runner
        -- returns, but technically it's legal to return them.
        let linLam' = ignoreHoistFailure $ hoist (PairB h b) linLam
        return (BinaryLamExpr h b body', linLam')
  linearizeEffectFun _ _ = error "expect effect function to be a binary lambda"
  
  notImplemented :: HasCallStack => a
  notImplemented = error "Not implemented"
  
  -- === boring instances ===
  
  instance GenericE ActivePrimals where
    type RepE ActivePrimals = PairE (ListE SAtomVar) (EffectRow SimpIR)
    fromE (ActivePrimals vs effs) = ListE vs `PairE` effs
    {-# INLINE fromE #-}
    toE   (ListE vs `PairE` effs) = ActivePrimals vs effs
    {-# INLINE toE #-}
  
  instance SinkableE   ActivePrimals
  instance HoistableE  ActivePrimals
  instance AlphaEqE    ActivePrimals
  instance RenameE     ActivePrimals
  
  instance GenericE TangentArgs where
    type RepE TangentArgs = ListE SAtomVar
    fromE (TangentArgs vs) = ListE vs
    {-# INLINE fromE #-}
    toE   (ListE vs) = TangentArgs vs
    {-# INLINE toE #-}
  
  instance SinkableE   TangentArgs
  instance HoistableE  TangentArgs
  instance AlphaEqE    TangentArgs
  instance RenameE     TangentArgs
  
  instance GenericE (MaybeReconAbs e) where
    type RepE (MaybeReconAbs e) = EitherE (ReconAbs SimpIR e) e
    fromE = \case
      ReconWithData ab -> LeftE ab
      TrivialRecon e   -> RightE e
    {-# INLINE fromE #-}
  
    toE = \case
      LeftE ab -> ReconWithData ab
      RightE e -> TrivialRecon e
    {-# INLINE toE #-}
  
  instance SinkableE  e => SinkableE  (MaybeReconAbs e)
  instance HoistableE e => HoistableE (MaybeReconAbs e)
  instance RenameE    e => RenameE    (MaybeReconAbs e)
  
  instance SinkableE  e => SinkableE  (ObligateReconAbs e) where
    sinkingProofE = undefined
  