content: -- Copyright 2022 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 Optimize
    ( optimize, hoistLoopInvariant, dceTop) where
  
  import Data.Functor
  import Control.Monad
  import Control.Monad.State.Strict
  
  import Types.Core
  import Types.Primitives
  import Types.Top
  import MTL1
  import Name
  import Subst
  import IRVariants
  import Core
  import CheapReduction
  import Builder
  import QueryType
  import Util (iota)
  import Err
  
  optimize :: EnvReader m => STopLam n -> m n (STopLam n)
  optimize = dceTop     -- Clean up user code
         >=> unrollLoops
         >=> dceTop     -- Clean up peephole-optimized code after unrolling
         >=> hoistLoopInvariant
  
  -- === Loop unrolling ===
  
  unrollLoops :: EnvReader m => STopLam n -> m n (STopLam n)
  unrollLoops lam = liftLamExpr lam unrollLoopsExpr
  {-# SCC unrollLoops #-}
  
  unrollLoopsExpr :: EnvReader m => SExpr n -> m n (SExpr n)
  unrollLoopsExpr b = liftM fst $
    liftBuilder $ runStateT1 (runSubstReaderT idSubst (runULM $ buildBlock $ ulExpr b)) (ULS 0)
  
  newtype ULS n = ULS Int deriving Show
  newtype ULM i o a = ULM { runULM :: SubstReaderT AtomSubstVal (StateT1 ULS (BuilderM SimpIR)) i o a}
    deriving (MonadFail, Fallible, Functor, Applicative, Monad, ScopeReader,
              EnvReader, EnvExtender, Builder SimpIR, SubstReader AtomSubstVal,
              ScopableBuilder SimpIR, MonadState (ULS o))
  
  instance SinkableE ULS where
    sinkingProofE _ (ULS c) = ULS c
  instance HoistableState ULS where
    hoistState _ _ (ULS c) = (ULS c)
    {-# INLINE hoistState #-}
  
  instance NonAtomRenamer (ULM i o) i o where renameN = substM
  instance Visitor (ULM i o) SimpIR i o where
    visitAtom = visitAtomDefault
    visitType = visitTypeDefault
    visitPi   = visitPiDefault
    visitLam  = visitLamEmits
  instance ExprVisitorEmits (ULM i o) SimpIR i o where
    visitExprEmits = ulExpr
  
  -- TODO: Refine the cost accounting so that operations that will become
  -- constant-foldable after inlining don't count towards it.
  ulExpr :: Emits o => SExpr i -> ULM i o (SAtom o)
  ulExpr expr = case expr of
    PrimOp (Hof (TypedHof _ (For Fwd ixTy body))) ->
      case ixTypeDict ixTy of
        DictCon (IxRawFin (IdxRepVal n)) -> do
          (body', bodyCost) <- withLocalAccounting $ visitLamEmits body
          -- We add n (in the form of (... + 1) * n) for the cost of the TabCon reconstructing the result.
          case (bodyCost + 1) * (fromIntegral n) <= unrollBlowupThreshold of
            True -> case body' of
              UnaryLamExpr b' block' -> do
                vals <- dropSubst $ forM (iota n) \i -> do
                  extendSubst (b' @> SubstVal (IdxRepVal i)) $ ulExpr block'
                inc $ fromIntegral n  -- To account for the TabCon we emit below
                getLamExprType body' >>= \case
                  PiType (UnaryNest (tb:>_)) (EffTy _ valTy) -> do
                    let tabTy = toType $ TabPiType (DictCon $ IxRawFin (IdxRepVal n)) (tb:>IdxRepTy) valTy
                    emit $ TabCon Nothing tabTy vals
                  _ -> error "Expected `for` body to have a Pi type"
              _ -> error "Expected `for` body to be a lambda expression"
            False -> do
              inc bodyCost
              ixTy' <- visitGeneric ixTy
              emitHof $ For Fwd ixTy' body'
        _ -> nothingSpecial
    -- Avoid unrolling loops with large table literals
    TabCon _ _ els -> inc (length els) >> nothingSpecial
    Block _ (Abs decls body) -> visitDeclsEmits decls $ ulExpr body
    _ -> nothingSpecial
    where
      inc i = modify \(ULS n) -> ULS (n + i)
      nothingSpecial = inc 1 >> visitGeneric expr >>= emit
      unrollBlowupThreshold = 12
      withLocalAccounting m = do
        oldCost <- get
        ans <- put (ULS 0) *> m
        ULS newCost <- get
        put oldCost $> (ans, newCost)
      {-# INLINE withLocalAccounting #-}
  
  -- === Loop invariant code motion ===
  
  data LICMTag
  type LICMM = AtomSubstBuilder LICMTag SimpIR
  
  liftLICMM :: EnvReader m => LICMM n n a -> m n a
  liftLICMM cont = liftAtomSubstBuilder cont
  
  instance NonAtomRenamer (LICMM i o) i o where renameN = substM
  instance Visitor (LICMM i o) SimpIR i o where
    visitAtom = visitAtomDefault
    visitType = visitTypeDefault
    visitPi   = visitPiDefault
    visitLam  = visitLamEmits
  
  instance ExprVisitorEmits (LICMM i o) SimpIR i o where
    visitExprEmits = licmExpr
  
  hoistLoopInvariantExpr :: EnvReader m => SExpr n -> m n (SExpr n)
  hoistLoopInvariantExpr body = liftLICMM $ buildBlock $ visitExprEmits body
  {-# SCC hoistLoopInvariantExpr #-}
  
  hoistLoopInvariant :: EnvReader m => STopLam n -> m n (STopLam n)
  hoistLoopInvariant lam = liftLamExpr lam hoistLoopInvariantExpr
  {-# INLINE hoistLoopInvariant #-}
  
  licmExpr :: Emits o => SExpr i -> LICMM i o (SAtom o)
  licmExpr = \case
    PrimOp (DAMOp (Seq _ dir ix (Con (ProdCon dests)) (LamExpr (UnaryNest b) body))) -> do
      ix' <- substM ix
      dests' <- mapM visitAtom dests
      let numCarriesOriginal = length dests'
      Abs hdecls destsAndBody <- visitBinders (UnaryNest b) \(UnaryNest b') -> do
        -- First, traverse the block, to allow any Hofs inside it to hoist their own decls.
        Abs decls ans <- buildScoped $ visitExprEmits body
        -- Now, we process the decls and decide which ones to hoist.
        liftEnvReaderM $ runSubstReaderT idSubst $
            seqLICM REmpty mempty (asNameBinder b') REmpty decls ans
      PairE (ListE extraDests) ab <- emitDecls $ Abs hdecls destsAndBody
      extraDests' <- mapM toAtomVar extraDests
      -- Append the destinations of hoisted Allocs as loop carried values.
      let dests'' = Con $ ProdCon $ dests' ++ (toAtom <$> extraDests')
      let carryTy = getType dests''
      let lbTy = case ix' of IxType ixTy _ -> PairTy ixTy carryTy
      extraDestsTyped <- forM extraDests' \(AtomVar d t) -> return (d, t)
      Abs extraDestBs (Abs lb bodyAbs) <- return $ abstractFreeVars extraDestsTyped ab
      body' <- withFreshBinder noHint lbTy \lb' -> do
        (oldIx, allCarries) <- fromPairReduced $ toAtom $ binderVar lb'
        (oldCarries, newCarries) <- splitAt numCarriesOriginal <$> getUnpackedReduced allCarries
        let oldLoopBinderVal = Con $ ProdCon [oldIx, Con $ ProdCon oldCarries]
        let s = extraDestBs @@> map SubstVal newCarries <.> lb @> SubstVal oldLoopBinderVal
        block <- mkBlock =<< applySubst s bodyAbs
        return $ UnaryLamExpr lb' block
      emitSeq dir ix' dests'' body'
    PrimOp (Hof (TypedHof _ (For dir ix (LamExpr (UnaryNest b) body)))) -> do
      ix' <- substM ix
      Abs hdecls destsAndBody <- visitBinders (UnaryNest b) \(UnaryNest b') -> do
        Abs decls ans <- buildScoped $ visitExprEmits body
        liftEnvReaderM $ runSubstReaderT idSubst $
            seqLICM REmpty mempty (asNameBinder b') REmpty decls ans
      PairE (ListE []) (Abs lnb bodyAbs) <- emitDecls $ Abs hdecls destsAndBody
      ixTy <- substM $ binderType b
      body' <- withFreshBinder noHint ixTy \i -> do
        block <- mkBlock =<< applyRename (lnb@>binderName i) bodyAbs
        return $ UnaryLamExpr i block
      emitHof $ For dir ix' body'
    Block _ (Abs decls result) -> visitDeclsEmits decls $ licmExpr result
    expr -> visitGeneric expr >>= emit
  
  seqLICM :: RNest SDecl n1 n2      -- hoisted decls
          -> [SAtomName n2]          -- hoisted dests
          -> AtomNameBinder SimpIR n2 n3   -- loop binder
          -> RNest SDecl n3 n4      -- loop-dependent decls
          -> Nest SDecl m1 m2       -- decls remaining to process
          -> SAtom m2               -- loop result
          -> SubstReaderT AtomSubstVal EnvReaderM m1 n4
               (Abs (Nest SDecl)            -- hoisted decls
                  (PairE (ListE SAtomName)  -- hoisted allocs (these should go in the loop carry)
                         (Abs (AtomNameBinder SimpIR) -- loop binder
                              (Abs (Nest SDecl)       -- non-hoisted decls
                               SAtom))) n1)           -- final result
  seqLICM !top !topDestNames !lb !reg decls ans = case decls of
    Empty -> do
      ans' <- substM ans
      return $ Abs (unRNest top) $ PairE (ListE $ reverse topDestNames) $ Abs lb $ Abs (unRNest reg) ans'
    Nest (Let bb binding) bs -> do
      binding' <- substM binding
      withFreshBinder (getNameHint bb) binding' \(bb':>_) -> do
        let b = Let bb' binding'
        let moveOn = extendRenamer (bb@>binderName bb') $
                       seqLICM top topDestNames lb (RNest reg b) bs ans
        case getEffects binding' of
          -- OPTIMIZE: We keep querying the ScopeFrag of lb and reg here, leading to quadratic runtime
          Pure -> case exchangeBs $ PairB (PairB lb reg) b of
            HoistSuccess (PairB b' lbreg@(PairB lb' reg')) -> do
              withSubscopeDistinct lbreg $ withExtEvidence b' $
                extendRenamer (bb@>sink (binderName b')) do
                  extraTopDestNames <- return case b' of
                    Let bn (DeclBinding _ (PrimOp (DAMOp (AllocDest _)))) -> [binderName bn]
                    _ -> []
                  seqLICM (RNest top b') (extraTopDestNames ++ sinkList topDestNames) lb' reg' bs ans
                where
            HoistFailure _ -> moveOn
          _ -> moveOn
  
  -- === Dead code elimination ===
  
  newtype FV n = FV (NameSet n) deriving (Semigroup, Monoid)
  instance SinkableE FV where
    sinkingProofE _ _ = todoSinkableProof
  instance HoistableState FV where
    hoistState _ b (FV ns) = FV $ hoistFilterNameSet b ns
    {-# INLINE hoistState #-}
  
  newtype DCEM n a = DCEM { runDCEM :: StateT1 FV EnvReaderM n a }
    deriving ( Functor, Applicative, Monad, EnvReader, ScopeReader
             , MonadState (FV n), EnvExtender)
  
  dceTop :: EnvReader m => STopLam n -> m n (STopLam n)
  dceTop lam = liftLamExpr lam dceExpr
  {-# INLINE dceTop #-}
  
  dceExpr :: EnvReader m => SExpr n -> m n (SExpr n)
  dceExpr b = liftEnvReaderM $ evalStateT1 (runDCEM $ dce b) mempty
  {-# SCC dceExpr #-}
  
  class HasDCE (e::E) where
    dce :: e n -> DCEM n (e n)
    default dce :: VisitGeneric e SimpIR => e n -> DCEM n (e n)
    dce = visitGeneric
  
  instance NonAtomRenamer (DCEM o) o o where renameN = dce
  instance Visitor (DCEM o) SimpIR o o where
    visitType = dce
    visitAtom = dce
    visitLam  = dce
    visitPi   = dce
  
  instance Color c => HasDCE (Name c) where
    dce n = modify (<> FV (freeVarsE n)) $> n
  
  instance HasDCE SAtom where
    dce atom = case atom of
      Stuck _ _ -> modify (<> FV (freeVarsE atom)) $> atom
      Con con   -> Con <$> visitGeneric con
  
  instance HasDCE SType where
    dce (TyCon e) = TyCon <$> visitGeneric e
  
  instance HasDCE (PiType SimpIR) where
    dce (PiType bs effTy) = do
      dceBinders bs effTy \bs' effTy' -> PiType bs' <$> dce effTy'
  
  instance HasDCE (LamExpr SimpIR) where
    dce (LamExpr bs e) = dceBinders bs e \bs' e' -> LamExpr bs' <$> dce e'
  
  instance HasDCE (Expr SimpIR) where
    dce = \case
      Block effTy block -> do
        -- The free vars accumulated in the state of DCEM should correspond to
        -- the free vars of the Abs of the block answer, by the decls traversed
        -- so far. dceNest takes care to uphold this invariant, but we temporarily
        -- reset the state to an empty map, just so that names from the surrounding
        -- block don't end up influencing elimination decisions here. Note that we
        -- restore the state (and accumulate free vars of the DCE'd block into it)
        -- right after dceNest.
        effTy' <- dce effTy
        old <- get
        put mempty
        block' <- dceBlock block
        modify (<> old)
        return $ Block effTy' block'
      e -> visitGeneric e
  
  dceBinders
    :: (HoistableB b, BindsEnv b, RenameB b, RenameE e)
    => b n l -> e l
    -> (forall l'. b n l' -> e l' -> DCEM l' a)
    -> DCEM n a
  dceBinders b e cont = do
    ans <- refreshAbs (Abs b e) \b' e' -> cont b' e'
    modify (<>FV (freeVarsB b))
    return ans
  {-# INLINE dceBinders #-}
  
  wrapWithCachedFVs :: HoistableE e => e n -> DCEM n (CachedFVs e n)
  wrapWithCachedFVs e = do
    FV fvs <- get
  #ifdef DEX_DEBUG
    let fvsAreCorrect = nameSetRawNames fvs == nameSetRawNames (freeVarsE e)
  #else
    -- Verification of this invariant defeats the performance benefits of
    -- avoiding the extra traversal (e.g. actually having linear complexity),
    -- so we only do that in debug builds.
    let fvsAreCorrect = True
  #endif
    case fvsAreCorrect of
      True -> return $ UnsafeCachedFVs fvs e
      False -> error $ "Free variables were computed incorrectly."
  
  hoistUsingCachedFVs :: (BindsNames b, HoistableE e) =>
    b n l -> e l -> DCEM l (HoistExcept (e n))
  hoistUsingCachedFVs b e = hoistViaCachedFVs b <$> wrapWithCachedFVs e
  
  data ElimResult n where
    ElimSuccess :: SBlock n -> ElimResult n
    ElimFailure :: SDecl n l -> SBlock l -> ElimResult n
  
  dceBlock :: SBlock n -> DCEM n (SBlock n)
  dceBlock (Abs decls ans) = case decls of
    Empty -> Abs Empty <$> dce ans
    Nest b@(Let _ decl) bs -> do
      -- Note that we only ever dce the abs below under this refreshAbs,
      -- which will remove any references to b upon exit (it happens
      -- because refreshAbs of StateT1 triggers hoistState, which we
      -- implement by deleting the entries that can't hoist).
      dceAttempt <- refreshAbs (Abs b (Abs bs ans)) \b' (Abs bs' ans') -> do
        below <- dceBlock $ Abs bs' ans'
        case isPure decl of
          False -> return $ ElimFailure b' below
          True  -> do
            hoistUsingCachedFVs b' below <&> \case
              HoistSuccess below' -> ElimSuccess below'
              HoistFailure _ -> ElimFailure b' below
      case dceAttempt of
        ElimSuccess below' -> return below'
        ElimFailure (Let b' (DeclBinding ann expr)) (Abs bs'' ans'') -> do
          expr' <- dce expr
          modify (<>FV (freeVarsB b'))
          return $ Abs (Nest (Let b' (DeclBinding ann expr')) bs'') ans''
  
  instance HasDCE (EffectRow   SimpIR)
  instance HasDCE (EffTy       SimpIR)
  