    {-# LANGUAGE TypeFamilies #-}

    module Futhark.IR.Mem.Simplify
      ( simplifyProgGeneric,
        simplifyStmsGeneric,
        simpleGeneric,
        SimplifyMemory,
        memRuleBook,
      )
    where

    import Futhark.Analysis.SymbolTable qualified as ST
    import Futhark.Analysis.UsageTable qualified as UT
    import Futhark.Construct
    import Futhark.IR.Mem
    import Futhark.IR.Prop.Aliases (AliasedOp)
    import Futhark.Optimise.Simplify qualified as Simplify
    import Futhark.Optimise.Simplify.Engine qualified as Engine
    import Futhark.Optimise.Simplify.Rep
    import Futhark.Optimise.Simplify.Rule
    import Futhark.Optimise.Simplify.Rules
    import Futhark.Pass
    import Futhark.Pass.ExplicitAllocations (simplifiable)

    -- | Some constraints that must hold for the simplification rules to work.
    type SimplifyMemory rep inner =
      ( Simplify.SimplifiableRep rep,
        LetDec rep ~ LetDecMem,
        ExpDec rep ~ (),
        BodyDec rep ~ (),
        CanBeWise (OpC rep),
        BuilderOps (Wise rep),
        OpReturns inner,
        ST.IndexOp (inner (Wise rep)),
        AliasedOp inner,
        Mem rep inner,
        CanBeWise inner,
        RephraseOp inner,
        ASTConstraints (inner (Engine.Wise rep))
      )

    simpleGeneric ::
      (SimplifyMemory rep inner) =>
      (inner (Wise rep) -> UT.UsageTable) ->
      Simplify.SimplifyOp rep (inner (Wise rep)) ->
      Simplify.SimpleOps rep
    simpleGeneric :: forall rep (inner :: * -> *).
    SimplifyMemory rep inner =>
    (inner (Wise rep) -> UsageTable)
    -> SimplifyOp rep (inner (Wise rep)) -> SimpleOps rep
    simpleGeneric = (inner (Wise rep) -> UsageTable)
    -> (inner (Wise rep)
        -> SimpleM rep (inner (Wise rep), Stms (Wise rep)))
    -> SimpleOps rep
    forall rep (inner :: * -> *).
    (SimplifiableRep rep, LetDec rep ~ LetDecMem, ExpDec rep ~ (),
     BodyDec rep ~ (), Mem (Wise rep) inner, CanBeWise inner,
     RephraseOp inner, IsOp inner, OpReturns inner, AliasedOp inner,
     IndexOp (inner (Wise rep))) =>
    (inner (Wise rep) -> UsageTable)
    -> (inner (Wise rep)
        -> SimpleM rep (inner (Wise rep), Stms (Wise rep)))
    -> SimpleOps rep
    simplifiable

    simplifyProgGeneric ::
      (SimplifyMemory rep inner) =>
      RuleBook (Wise rep) ->
      Simplify.SimpleOps rep ->
      Prog rep ->
      PassM (Prog rep)
    simplifyProgGeneric :: forall rep (inner :: * -> *).
    SimplifyMemory rep inner =>
    RuleBook (Wise rep)
    -> SimpleOps rep -> Prog rep -> PassM (Prog rep)
    simplifyProgGeneric RuleBook (Wise rep)
    rules SimpleOps rep
    ops =
      SimpleOps rep
    -> RuleBook (Wise rep)
    -> HoistBlockers rep
    -> Prog rep
    -> PassM (Prog rep)
    forall rep.
    SimplifiableRep rep =>
    SimpleOps rep
    -> RuleBook (Wise rep)
    -> HoistBlockers rep
    -> Prog rep
    -> PassM (Prog rep)
    Simplify.simplifyProg
        SimpleOps rep
    ops
        RuleBook (Wise rep)
    rules
        HoistBlockers rep
    forall rep (inner :: * -> *).
    (OpC rep ~ MemOp inner) =>
    HoistBlockers rep
    blockers {Engine.blockHoistBranch = blockAllocs}
      where
        blockAllocs :: SymbolTable rep -> p -> Stm rep -> Bool
    blockAllocs SymbolTable rep
    vtable p
    _ (Let Pat (LetDec rep)
    _ StmAux (ExpDec rep)
    _ (Op Alloc {})) =
          Bool -> Bool
    not (Bool -> Bool) -> Bool -> Bool
    forall a b. (a -> b) -> a -> b
    $ SymbolTable rep -> Bool
    forall rep. SymbolTable rep -> Bool
    ST.simplifyMemory SymbolTable rep
    vtable
        -- Do not hoist statements that produce arrays.  This is
        -- because in the KernelsMem representation, multiple
        -- arrays can be located in the same memory block, and moving
        -- their creation out of a branch can thus cause memory
        -- corruption.  At this point in the compiler we have probably
        -- already moved all the array creations that matter.
        blockAllocs SymbolTable rep
    _ p
    _ (Let Pat (LetDec rep)
    pat StmAux (ExpDec rep)
    _ Exp rep
    _) =
          Bool -> Bool
    not (Bool -> Bool) -> Bool -> Bool
    forall a b. (a -> b) -> a -> b
    $ (Type -> Bool) -> [Type] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all Type -> Bool
    forall shape u. TypeBase shape u -> Bool
    primType ([Type] -> Bool) -> [Type] -> Bool
    forall a b. (a -> b) -> a -> b
    $ Pat (LetDec rep) -> [Type]
    forall dec. Typed dec => Pat dec -> [Type]
    patTypes Pat (LetDec rep)
    pat

    simplifyStmsGeneric ::
      ( HasScope rep m,
        MonadFreshNames m,
        SimplifyMemory rep inner
      ) =>
      RuleBook (Wise rep) ->
      Simplify.SimpleOps rep ->
      Stms rep ->
      m (Stms rep)
    simplifyStmsGeneric :: forall rep (m :: * -> *) (inner :: * -> *).
    (HasScope rep m, MonadFreshNames m, SimplifyMemory rep inner) =>
    RuleBook (Wise rep) -> SimpleOps rep -> Stms rep -> m (Stms rep)
    simplifyStmsGeneric RuleBook (Wise rep)
    rules SimpleOps rep
    ops Stms rep
    stms = do
      Scope rep
    scope <- m (Scope rep)
    forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
    askScope
      SimpleOps rep
    -> RuleBook (Wise rep)
    -> HoistBlockers rep
    -> Scope rep
    -> Stms rep
    -> m (Stms rep)
    forall (m :: * -> *) rep.
    (MonadFreshNames m, SimplifiableRep rep) =>
    SimpleOps rep
    -> RuleBook (Wise rep)
    -> HoistBlockers rep
    -> Scope rep
    -> Stms rep
    -> m (Stms rep)
    Simplify.simplifyStms SimpleOps rep
    ops RuleBook (Wise rep)
    rules HoistBlockers rep
    forall rep (inner :: * -> *).
    (OpC rep ~ MemOp inner) =>
    HoistBlockers rep
    blockers Scope rep
    scope Stms rep
    stms

    isResultAlloc :: (OpC rep ~ MemOp op) => Engine.BlockPred rep
    isResultAlloc :: forall rep (op :: * -> *). (OpC rep ~ MemOp op) => BlockPred rep
    isResultAlloc SymbolTable rep
    _ UsageTable
    usage (Let (Pat [PatElem (LetDec rep)
    pe]) StmAux (ExpDec rep)
    _ (Op Alloc {})) =
      VName -> UsageTable -> Bool
    UT.isInResult (PatElem (LetDec rep) -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem (LetDec rep)
    pe) UsageTable
    usage
    isResultAlloc SymbolTable rep
    _ UsageTable
    _ Stm rep
    _ = Bool
    False

    isAlloc :: (OpC rep ~ MemOp op) => Engine.BlockPred rep
    isAlloc :: forall rep (op :: * -> *). (OpC rep ~ MemOp op) => BlockPred rep
    isAlloc SymbolTable rep
    _ UsageTable
    _ (Let Pat (LetDec rep)
    _ StmAux (ExpDec rep)
    _ (Op Alloc {})) = Bool
    True
    isAlloc SymbolTable rep
    _ UsageTable
    _ Stm rep
    _ = Bool
    False

    blockers ::
      (OpC rep ~ MemOp inner) =>
      Simplify.HoistBlockers rep
    blockers :: forall rep (inner :: * -> *).
    (OpC rep ~ MemOp inner) =>
    HoistBlockers rep
    blockers =
      HoistBlockers rep
    forall {k} (rep :: k). HoistBlockers rep
    Engine.noExtraHoistBlockers
        { Engine.blockHoistPar = isAlloc,
          Engine.blockHoistSeq = isResultAlloc,
          Engine.isAllocation = isAlloc mempty mempty
        }

    -- | Standard collection of simplification rules for representations
    -- with memory.
    memRuleBook :: (SimplifyMemory rep inner) => RuleBook (Wise rep)
    memRuleBook :: forall rep (inner :: * -> *).
    SimplifyMemory rep inner =>
    RuleBook (Wise rep)
    memRuleBook =
      RuleBook (Wise rep)
    forall rep. (BuilderOps rep, TraverseOpStms rep) => RuleBook rep
    standardRules
        RuleBook (Wise rep) -> RuleBook (Wise rep) -> RuleBook (Wise rep)
    forall a. Semigroup a => a -> a -> a
    <> [TopDownRule (Wise rep)]
    -> [BottomUpRule (Wise rep)] -> RuleBook (Wise rep)
    forall m. [TopDownRule m] -> [BottomUpRule m] -> RuleBook m
    ruleBook
          [ RuleOp (Wise rep) (TopDown (Wise rep)) -> TopDownRule (Wise rep)
    forall rep a. RuleOp rep a -> SimplificationRule rep a
    RuleOp RuleOp (Wise rep) (TopDown (Wise rep))
    forall rep (inner :: * -> *).
    SimplifyMemory rep inner =>
    TopDownRuleOp (Wise rep)
    decertifySafeAlloc
          ]
          []

    -- If an allocation is statically known to be safe, then we can remove
    -- the certificates on it.  This can help hoist things that would
    -- otherwise be stuck inside loops or branches.
    decertifySafeAlloc :: (SimplifyMemory rep inner) => TopDownRuleOp (Wise rep)
    decertifySafeAlloc :: forall rep (inner :: * -> *).
    SimplifyMemory rep inner =>
    TopDownRuleOp (Wise rep)
    decertifySafeAlloc TopDown (Wise rep)
    _ Pat (LetDec (Wise rep))
    pat (StmAux Certs
    cs Attrs
    attrs ExpDec (Wise rep)
    _) Op (Wise rep)
    op
      | Certs
    cs Certs -> Certs -> Bool
    forall a. Eq a => a -> a -> Bool
    /= Certs
    forall a. Monoid a => a
    mempty,
        [Mem Space
    _] <- Pat (VarWisdom, LetDecMem) -> [Type]
    forall dec. Typed dec => Pat dec -> [Type]
    patTypes Pat (VarWisdom, LetDecMem)
    Pat (LetDec (Wise rep))
    pat,
        MemOp inner (Wise rep) -> Bool
    forall rep. ASTRep rep => MemOp inner rep -> Bool
    forall (op :: * -> *) rep. (IsOp op, ASTRep rep) => op rep -> Bool
    safeOp Op (Wise rep)
    MemOp inner (Wise rep)
    op =
          RuleM (Wise rep) () -> Rule (Wise rep)
    forall rep. RuleM rep () -> Rule rep
    Simplify (RuleM (Wise rep) () -> Rule (Wise rep))
    -> RuleM (Wise rep) () -> Rule (Wise rep)
    forall a b. (a -> b) -> a -> b
    $ Attrs -> RuleM (Wise rep) () -> RuleM (Wise rep) ()
    forall (m :: * -> *) a. MonadBuilder m => Attrs -> m a -> m a
    attributing Attrs
    attrs (RuleM (Wise rep) () -> RuleM (Wise rep) ())
    -> RuleM (Wise rep) () -> RuleM (Wise rep) ()
    forall a b. (a -> b) -> a -> b
    $ Pat (LetDec (Rep (RuleM (Wise rep))))
    -> Exp (Rep (RuleM (Wise rep))) -> RuleM (Wise rep) ()
    forall (m :: * -> *).
    MonadBuilder m =>
    Pat (LetDec (Rep m)) -> Exp (Rep m) -> m ()
    letBind Pat (LetDec (Rep (RuleM (Wise rep))))
    Pat (LetDec (Wise rep))
    pat (Exp (Rep (RuleM (Wise rep))) -> RuleM (Wise rep) ())
    -> Exp (Rep (RuleM (Wise rep))) -> RuleM (Wise rep) ()
    forall a b. (a -> b) -> a -> b
    $ Op (Wise rep) -> Exp (Wise rep)
    forall rep. Op rep -> Exp rep
    Op Op (Wise rep)
    op
    decertifySafeAlloc TopDown (Wise rep)
    _ Pat (LetDec (Wise rep))
    _ StmAux (ExpDec (Wise rep))
    _ Op (Wise rep)
    _ = Rule (Wise rep)
    forall rep. Rule rep
    Skip