    {-# LANGUAGE TypeFamilies #-}
    {-# LANGUAGE UndecidableInstances #-}
    {-# OPTIONS_GHC -fno-warn-redundant-constraints #-}

    -- | This module defines the concept of a simplification rule for
    -- bindings.  The intent is that you pass some context (such as symbol
    -- table) and a binding, and is given back a sequence of bindings that
    -- compute the same result, but are "better" in some sense.
    --
    -- These rewrite rules are "local", in that they do not maintain any
    -- state or look at the program as a whole.  Compare this to the
    -- fusion algorithm in @Futhark.Optimise.Fusion.Fusion@, which must be implemented
    -- as its own pass.
    module Futhark.Optimise.Simplify.Rule
      ( -- * The rule monad
        RuleM,
        cannotSimplify,
        liftMaybe,

        -- * Rule definition
        Rule (..),
        SimplificationRule (..),
        RuleGeneric,
        RuleBasicOp,
        RuleMatch,
        RuleLoop,

        -- * Top-down rules
        TopDown,
        TopDownRule,
        TopDownRuleGeneric,
        TopDownRuleBasicOp,
        TopDownRuleMatch,
        TopDownRuleLoop,
        TopDownRuleOp,

        -- * Bottom-up rules
        BottomUp,
        BottomUpRule,
        BottomUpRuleGeneric,
        BottomUpRuleBasicOp,
        BottomUpRuleMatch,
        BottomUpRuleLoop,
        BottomUpRuleOp,

        -- * Assembling rules
        RuleBook,
        ruleBook,

        -- * Applying rules
        topDownSimplifyStm,
        bottomUpSimplifyStm,
      )
    where

    import Control.Monad.State
    import Futhark.Analysis.SymbolTable qualified as ST
    import Futhark.Analysis.UsageTable qualified as UT
    import Futhark.Builder
    import Futhark.IR

    -- | The monad in which simplification rules are evaluated.
    newtype RuleM rep a = RuleM (BuilderT rep (StateT VNameSource Maybe) a)
      deriving
        ( (forall a b. (a -> b) -> RuleM rep a -> RuleM rep b)
    -> (forall a b. a -> RuleM rep b -> RuleM rep a)
    -> Functor (RuleM rep)
    forall a b. a -> RuleM rep b -> RuleM rep a
    forall a b. (a -> b) -> RuleM rep a -> RuleM rep b
    forall rep a b. a -> RuleM rep b -> RuleM rep a
    forall rep a b. (a -> b) -> RuleM rep a -> RuleM rep b
    forall (f :: * -> *).
    (forall a b. (a -> b) -> f a -> f b)
    -> (forall a b. a -> f b -> f a) -> Functor f
    $cfmap :: forall rep a b. (a -> b) -> RuleM rep a -> RuleM rep b
    fmap :: forall a b. (a -> b) -> RuleM rep a -> RuleM rep b
    $c<$ :: forall rep a b. a -> RuleM rep b -> RuleM rep a
    <$ :: forall a b. a -> RuleM rep b -> RuleM rep a
    Functor,
          Functor (RuleM rep)
    Functor (RuleM rep) =>
    (forall a. a -> RuleM rep a)
    -> (forall a b. RuleM rep (a -> b) -> RuleM rep a -> RuleM rep b)
    -> (forall a b c.
        (a -> b -> c) -> RuleM rep a -> RuleM rep b -> RuleM rep c)
    -> (forall a b. RuleM rep a -> RuleM rep b -> RuleM rep b)
    -> (forall a b. RuleM rep a -> RuleM rep b -> RuleM rep a)
    -> Applicative (RuleM rep)
    forall rep. Functor (RuleM rep)
    forall a. a -> RuleM rep a
    forall rep a. a -> RuleM rep a
    forall a b. RuleM rep a -> RuleM rep b -> RuleM rep a
    forall a b. RuleM rep a -> RuleM rep b -> RuleM rep b
    forall a b. RuleM rep (a -> b) -> RuleM rep a -> RuleM rep b
    forall rep a b. RuleM rep a -> RuleM rep b -> RuleM rep a
    forall rep a b. RuleM rep a -> RuleM rep b -> RuleM rep b
    forall rep a b. RuleM rep (a -> b) -> RuleM rep a -> RuleM rep b
    forall a b c.
    (a -> b -> c) -> RuleM rep a -> RuleM rep b -> RuleM rep c
    forall rep a b c.
    (a -> b -> c) -> RuleM rep a -> RuleM rep b -> RuleM rep c
    forall (f :: * -> *).
    Functor f =>
    (forall a. a -> f a)
    -> (forall a b. f (a -> b) -> f a -> f b)
    -> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
    -> (forall a b. f a -> f b -> f b)
    -> (forall a b. f a -> f b -> f a)
    -> Applicative f
    $cpure :: forall rep a. a -> RuleM rep a
    pure :: forall a. a -> RuleM rep a
    $c<*> :: forall rep a b. RuleM rep (a -> b) -> RuleM rep a -> RuleM rep b
    <*> :: forall a b. RuleM rep (a -> b) -> RuleM rep a -> RuleM rep b
    $cliftA2 :: forall rep a b c.
    (a -> b -> c) -> RuleM rep a -> RuleM rep b -> RuleM rep c
    liftA2 :: forall a b c.
    (a -> b -> c) -> RuleM rep a -> RuleM rep b -> RuleM rep c
    $c*> :: forall rep a b. RuleM rep a -> RuleM rep b -> RuleM rep b
    *> :: forall a b. RuleM rep a -> RuleM rep b -> RuleM rep b
    $c<* :: forall rep a b. RuleM rep a -> RuleM rep b -> RuleM rep a
    <* :: forall a b. RuleM rep a -> RuleM rep b -> RuleM rep a
    Applicative,
          Applicative (RuleM rep)
    Applicative (RuleM rep) =>
    (forall a b. RuleM rep a -> (a -> RuleM rep b) -> RuleM rep b)
    -> (forall a b. RuleM rep a -> RuleM rep b -> RuleM rep b)
    -> (forall a. a -> RuleM rep a)
    -> Monad (RuleM rep)
    forall rep. Applicative (RuleM rep)
    forall a. a -> RuleM rep a
    forall rep a. a -> RuleM rep a
    forall a b. RuleM rep a -> RuleM rep b -> RuleM rep b
    forall a b. RuleM rep a -> (a -> RuleM rep b) -> RuleM rep b
    forall rep a b. RuleM rep a -> RuleM rep b -> RuleM rep b
    forall rep a b. RuleM rep a -> (a -> RuleM rep b) -> RuleM rep b
    forall (m :: * -> *).
    Applicative m =>
    (forall a b. m a -> (a -> m b) -> m b)
    -> (forall a b. m a -> m b -> m b)
    -> (forall a. a -> m a)
    -> Monad m
    $c>>= :: forall rep a b. RuleM rep a -> (a -> RuleM rep b) -> RuleM rep b
    >>= :: forall a b. RuleM rep a -> (a -> RuleM rep b) -> RuleM rep b
    $c>> :: forall rep a b. RuleM rep a -> RuleM rep b -> RuleM rep b
    >> :: forall a b. RuleM rep a -> RuleM rep b -> RuleM rep b
    $creturn :: forall rep a. a -> RuleM rep a
    return :: forall a. a -> RuleM rep a
    Monad,
          Monad (RuleM rep)
    RuleM rep VNameSource
    Monad (RuleM rep) =>
    RuleM rep VNameSource
    -> (VNameSource -> RuleM rep ()) -> MonadFreshNames (RuleM rep)
    VNameSource -> RuleM rep ()
    forall rep. Monad (RuleM rep)
    forall rep. RuleM rep VNameSource
    forall rep. VNameSource -> RuleM rep ()
    forall (m :: * -> *).
    Monad m =>
    m VNameSource -> (VNameSource -> m ()) -> MonadFreshNames m
    $cgetNameSource :: forall rep. RuleM rep VNameSource
    getNameSource :: RuleM rep VNameSource
    $cputNameSource :: forall rep. VNameSource -> RuleM rep ()
    putNameSource :: VNameSource -> RuleM rep ()
    MonadFreshNames,
          HasScope rep,
          LocalScope rep
        )

    instance (BuilderOps rep) => MonadBuilder (RuleM rep) where
      type Rep (RuleM rep) = rep
      mkExpDecM :: Pat (LetDec (Rep (RuleM rep)))
    -> Exp (Rep (RuleM rep)) -> RuleM rep (ExpDec (Rep (RuleM rep)))
    mkExpDecM Pat (LetDec (Rep (RuleM rep)))
    pat Exp (Rep (RuleM rep))
    e = BuilderT rep (StateT VNameSource Maybe) (ExpDec (Rep (RuleM rep)))
    -> RuleM rep (ExpDec (Rep (RuleM rep)))
    forall rep a.
    BuilderT rep (StateT VNameSource Maybe) a -> RuleM rep a
    RuleM (BuilderT rep (StateT VNameSource Maybe) (ExpDec (Rep (RuleM rep)))
     -> RuleM rep (ExpDec (Rep (RuleM rep))))
    -> BuilderT
         rep (StateT VNameSource Maybe) (ExpDec (Rep (RuleM rep)))
    -> RuleM rep (ExpDec (Rep (RuleM rep)))
    forall a b. (a -> b) -> a -> b
    $ Pat (LetDec (Rep (BuilderT rep (StateT VNameSource Maybe))))
    -> Exp (Rep (BuilderT rep (StateT VNameSource Maybe)))
    -> BuilderT
         rep
         (StateT VNameSource Maybe)
         (ExpDec (Rep (BuilderT rep (StateT VNameSource Maybe))))
    forall (m :: * -> *).
    MonadBuilder m =>
    Pat (LetDec (Rep m)) -> Exp (Rep m) -> m (ExpDec (Rep m))
    mkExpDecM Pat (LetDec (Rep (BuilderT rep (StateT VNameSource Maybe))))
    Pat (LetDec (Rep (RuleM rep)))
    pat Exp (Rep (BuilderT rep (StateT VNameSource Maybe)))
    Exp (Rep (RuleM rep))
    e
      mkBodyM :: Stms (Rep (RuleM rep))
    -> Result -> RuleM rep (Body (Rep (RuleM rep)))
    mkBodyM Stms (Rep (RuleM rep))
    stms Result
    res = BuilderT rep (StateT VNameSource Maybe) (Body (Rep (RuleM rep)))
    -> RuleM rep (Body (Rep (RuleM rep)))
    forall rep a.
    BuilderT rep (StateT VNameSource Maybe) a -> RuleM rep a
    RuleM (BuilderT rep (StateT VNameSource Maybe) (Body (Rep (RuleM rep)))
     -> RuleM rep (Body (Rep (RuleM rep))))
    -> BuilderT rep (StateT VNameSource Maybe) (Body (Rep (RuleM rep)))
    -> RuleM rep (Body (Rep (RuleM rep)))
    forall a b. (a -> b) -> a -> b
    $ Stms (Rep (BuilderT rep (StateT VNameSource Maybe)))
    -> Result
    -> BuilderT
         rep
         (StateT VNameSource Maybe)
         (Body (Rep (BuilderT rep (StateT VNameSource Maybe))))
    forall (m :: * -> *).
    MonadBuilder m =>
    Stms (Rep m) -> Result -> m (Body (Rep m))
    mkBodyM Stms (Rep (BuilderT rep (StateT VNameSource Maybe)))
    Stms (Rep (RuleM rep))
    stms Result
    res
      mkLetNamesM :: [VName]
    -> Exp (Rep (RuleM rep)) -> RuleM rep (Stm (Rep (RuleM rep)))
    mkLetNamesM [VName]
    pat Exp (Rep (RuleM rep))
    e = BuilderT rep (StateT VNameSource Maybe) (Stm (Rep (RuleM rep)))
    -> RuleM rep (Stm (Rep (RuleM rep)))
    forall rep a.
    BuilderT rep (StateT VNameSource Maybe) a -> RuleM rep a
    RuleM (BuilderT rep (StateT VNameSource Maybe) (Stm (Rep (RuleM rep)))
     -> RuleM rep (Stm (Rep (RuleM rep))))
    -> BuilderT rep (StateT VNameSource Maybe) (Stm (Rep (RuleM rep)))
    -> RuleM rep (Stm (Rep (RuleM rep)))
    forall a b. (a -> b) -> a -> b
    $ [VName]
    -> Exp (Rep (BuilderT rep (StateT VNameSource Maybe)))
    -> BuilderT
         rep
         (StateT VNameSource Maybe)
         (Stm (Rep (BuilderT rep (StateT VNameSource Maybe))))
    forall (m :: * -> *).
    MonadBuilder m =>
    [VName] -> Exp (Rep m) -> m (Stm (Rep m))
    mkLetNamesM [VName]
    pat Exp (Rep (BuilderT rep (StateT VNameSource Maybe)))
    Exp (Rep (RuleM rep))
    e

      addStms :: Stms (Rep (RuleM rep)) -> RuleM rep ()
    addStms = BuilderT rep (StateT VNameSource Maybe) () -> RuleM rep ()
    forall rep a.
    BuilderT rep (StateT VNameSource Maybe) a -> RuleM rep a
    RuleM (BuilderT rep (StateT VNameSource Maybe) () -> RuleM rep ())
    -> (Stms rep -> BuilderT rep (StateT VNameSource Maybe) ())
    -> Stms rep
    -> RuleM rep ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Stms rep -> BuilderT rep (StateT VNameSource Maybe) ()
    Stms (Rep (BuilderT rep (StateT VNameSource Maybe)))
    -> BuilderT rep (StateT VNameSource Maybe) ()
    forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
    addStms
      collectStms :: forall a. RuleM rep a -> RuleM rep (a, Stms (Rep (RuleM rep)))
    collectStms (RuleM BuilderT rep (StateT VNameSource Maybe) a
    m) = BuilderT rep (StateT VNameSource Maybe) (a, Stms (Rep (RuleM rep)))
    -> RuleM rep (a, Stms (Rep (RuleM rep)))
    forall rep a.
    BuilderT rep (StateT VNameSource Maybe) a -> RuleM rep a
    RuleM (BuilderT
       rep (StateT VNameSource Maybe) (a, Stms (Rep (RuleM rep)))
     -> RuleM rep (a, Stms (Rep (RuleM rep))))
    -> BuilderT
         rep (StateT VNameSource Maybe) (a, Stms (Rep (RuleM rep)))
    -> RuleM rep (a, Stms (Rep (RuleM rep)))
    forall a b. (a -> b) -> a -> b
    $ BuilderT rep (StateT VNameSource Maybe) a
    -> BuilderT
         rep
         (StateT VNameSource Maybe)
         (a, Stms (Rep (BuilderT rep (StateT VNameSource Maybe))))
    forall a.
    BuilderT rep (StateT VNameSource Maybe) a
    -> BuilderT
         rep
         (StateT VNameSource Maybe)
         (a, Stms (Rep (BuilderT rep (StateT VNameSource Maybe))))
    forall (m :: * -> *) a.
    MonadBuilder m =>
    m a -> m (a, Stms (Rep m))
    collectStms BuilderT rep (StateT VNameSource Maybe) a
    m

    -- | Execute a 'RuleM' action.  If succesful, returns the result and a
    -- list of new bindings.
    simplify ::
      Scope rep ->
      VNameSource ->
      Rule rep ->
      Maybe (Stms rep, VNameSource)
    simplify :: forall rep.
    Scope rep
    -> VNameSource -> Rule rep -> Maybe (Stms rep, VNameSource)
    simplify Scope rep
    _ VNameSource
    _ Rule rep
    Skip = Maybe (Stms rep, VNameSource)
    forall a. Maybe a
    Nothing
    simplify Scope rep
    scope VNameSource
    src (Simplify (RuleM BuilderT rep (StateT VNameSource Maybe) ()
    m)) =
      StateT VNameSource Maybe (Stms rep)
    -> VNameSource -> Maybe (Stms rep, VNameSource)
    forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
    runStateT (BuilderT rep (StateT VNameSource Maybe) ()
    -> Scope rep -> StateT VNameSource Maybe (Stms rep)
    forall (m :: * -> *) rep.
    MonadFreshNames m =>
    BuilderT rep m () -> Scope rep -> m (Stms rep)
    runBuilderT_ BuilderT rep (StateT VNameSource Maybe) ()
    m Scope rep
    scope) VNameSource
    src

    cannotSimplify :: RuleM rep a
    cannotSimplify :: forall rep a. RuleM rep a
    cannotSimplify = BuilderT rep (StateT VNameSource Maybe) a -> RuleM rep a
    forall rep a.
    BuilderT rep (StateT VNameSource Maybe) a -> RuleM rep a
    RuleM (BuilderT rep (StateT VNameSource Maybe) a -> RuleM rep a)
    -> BuilderT rep (StateT VNameSource Maybe) a -> RuleM rep a
    forall a b. (a -> b) -> a -> b
    $ StateT VNameSource Maybe a
    -> BuilderT rep (StateT VNameSource Maybe) a
    forall (m :: * -> *) a. Monad m => m a -> BuilderT rep m a
    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
    (MonadTrans t, Monad m) =>
    m a -> t m a
    lift (StateT VNameSource Maybe a
     -> BuilderT rep (StateT VNameSource Maybe) a)
    -> StateT VNameSource Maybe a
    -> BuilderT rep (StateT VNameSource Maybe) a
    forall a b. (a -> b) -> a -> b
    $ Maybe a -> StateT VNameSource Maybe a
    forall (m :: * -> *) a. Monad m => m a -> StateT VNameSource m a
    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
    (MonadTrans t, Monad m) =>
    m a -> t m a
    lift Maybe a
    forall a. Maybe a
    Nothing

    liftMaybe :: Maybe a -> RuleM rep a
    liftMaybe :: forall a rep. Maybe a -> RuleM rep a
    liftMaybe Maybe a
    Nothing = RuleM rep a
    forall rep a. RuleM rep a
    cannotSimplify
    liftMaybe (Just a
    x) = a -> RuleM rep a
    forall a. a -> RuleM rep a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure a
    x

    -- | An efficient way of encoding whether a simplification rule should even be attempted.
    data Rule rep
      = -- | Give it a shot.
        Simplify (RuleM rep ())
      | -- | Don't bother.
        Skip

    type RuleGeneric rep a = a -> Stm rep -> Rule rep

    type RuleBasicOp rep a =
      ( a ->
        Pat (LetDec rep) ->
        StmAux (ExpDec rep) ->
        BasicOp ->
        Rule rep
      )

    type RuleMatch rep a =
      a ->
      Pat (LetDec rep) ->
      StmAux (ExpDec rep) ->
      ( [SubExp],
        [Case (Body rep)],
        Body rep,
        MatchDec (BranchType rep)
      ) ->
      Rule rep

    type RuleLoop rep a =
      a ->
      Pat (LetDec rep) ->
      StmAux (ExpDec rep) ->
      ( [(FParam rep, SubExp)],
        LoopForm,
        Body rep
      ) ->
      Rule rep

    type RuleOp rep a =
      a ->
      Pat (LetDec rep) ->
      StmAux (ExpDec rep) ->
      Op rep ->
      Rule rep

    -- | A simplification rule takes some argument and a statement, and
    -- tries to simplify the statement.
    data SimplificationRule rep a
      = RuleGeneric (RuleGeneric rep a)
      | RuleBasicOp (RuleBasicOp rep a)
      | RuleMatch (RuleMatch rep a)
      | RuleLoop (RuleLoop rep a)
      | RuleOp (RuleOp rep a)

    -- | A collection of rules grouped by which forms of statements they
    -- may apply to.
    data Rules rep a = Rules
      { forall rep a. Rules rep a -> [SimplificationRule rep a]
    rulesAny :: [SimplificationRule rep a],
        forall rep a. Rules rep a -> [SimplificationRule rep a]
    rulesBasicOp :: [SimplificationRule rep a],
        forall rep a. Rules rep a -> [SimplificationRule rep a]
    rulesMatch :: [SimplificationRule rep a],
        forall rep a. Rules rep a -> [SimplificationRule rep a]
    rulesLoop :: [SimplificationRule rep a],
        forall rep a. Rules rep a -> [SimplificationRule rep a]
    rulesOp :: [SimplificationRule rep a]
      }

    instance Semigroup (Rules rep a) where
      Rules [SimplificationRule rep a]
    as1 [SimplificationRule rep a]
    bs1 [SimplificationRule rep a]
    cs1 [SimplificationRule rep a]
    ds1 [SimplificationRule rep a]
    es1 <> :: Rules rep a -> Rules rep a -> Rules rep a
    <> Rules [SimplificationRule rep a]
    as2 [SimplificationRule rep a]
    bs2 [SimplificationRule rep a]
    cs2 [SimplificationRule rep a]
    ds2 [SimplificationRule rep a]
    es2 =
        [SimplificationRule rep a]
    -> [SimplificationRule rep a]
    -> [SimplificationRule rep a]
    -> [SimplificationRule rep a]
    -> [SimplificationRule rep a]
    -> Rules rep a
    forall rep a.
    [SimplificationRule rep a]
    -> [SimplificationRule rep a]
    -> [SimplificationRule rep a]
    -> [SimplificationRule rep a]
    -> [SimplificationRule rep a]
    -> Rules rep a
    Rules ([SimplificationRule rep a]
    as1 [SimplificationRule rep a]
    -> [SimplificationRule rep a] -> [SimplificationRule rep a]
    forall a. Semigroup a => a -> a -> a
    <> [SimplificationRule rep a]
    as2) ([SimplificationRule rep a]
    bs1 [SimplificationRule rep a]
    -> [SimplificationRule rep a] -> [SimplificationRule rep a]
    forall a. Semigroup a => a -> a -> a
    <> [SimplificationRule rep a]
    bs2) ([SimplificationRule rep a]
    cs1 [SimplificationRule rep a]
    -> [SimplificationRule rep a] -> [SimplificationRule rep a]
    forall a. Semigroup a => a -> a -> a
    <> [SimplificationRule rep a]
    cs2) ([SimplificationRule rep a]
    ds1 [SimplificationRule rep a]
    -> [SimplificationRule rep a] -> [SimplificationRule rep a]
    forall a. Semigroup a => a -> a -> a
    <> [SimplificationRule rep a]
    ds2) ([SimplificationRule rep a]
    es1 [SimplificationRule rep a]
    -> [SimplificationRule rep a] -> [SimplificationRule rep a]
    forall a. Semigroup a => a -> a -> a
    <> [SimplificationRule rep a]
    es2)

    instance Monoid (Rules rep a) where
      mempty :: Rules rep a
    mempty = [SimplificationRule rep a]
    -> [SimplificationRule rep a]
    -> [SimplificationRule rep a]
    -> [SimplificationRule rep a]
    -> [SimplificationRule rep a]
    -> Rules rep a
    forall rep a.
    [SimplificationRule rep a]
    -> [SimplificationRule rep a]
    -> [SimplificationRule rep a]
    -> [SimplificationRule rep a]
    -> [SimplificationRule rep a]
    -> Rules rep a
    Rules [SimplificationRule rep a]
    forall a. Monoid a => a
    mempty [SimplificationRule rep a]
    forall a. Monoid a => a
    mempty [SimplificationRule rep a]
    forall a. Monoid a => a
    mempty [SimplificationRule rep a]
    forall a. Monoid a => a
    mempty [SimplificationRule rep a]
    forall a. Monoid a => a
    mempty

    -- | Context for a rule applied during top-down traversal of the
    -- program.  Takes a symbol table as argument.
    type TopDown rep = ST.SymbolTable rep

    type TopDownRuleGeneric rep = RuleGeneric rep (TopDown rep)

    type TopDownRuleBasicOp rep = RuleBasicOp rep (TopDown rep)

    type TopDownRuleMatch rep = RuleMatch rep (TopDown rep)

    type TopDownRuleLoop rep = RuleLoop rep (TopDown rep)

    type TopDownRuleOp rep = RuleOp rep (TopDown rep)

    type TopDownRule rep = SimplificationRule rep (TopDown rep)

    -- | Context for a rule applied during bottom-up traversal of the
    -- program.  Takes a symbol table and usage table as arguments.
    type BottomUp rep = (ST.SymbolTable rep, UT.UsageTable)

    type BottomUpRuleGeneric rep = RuleGeneric rep (BottomUp rep)

    type BottomUpRuleBasicOp rep = RuleBasicOp rep (BottomUp rep)

    type BottomUpRuleMatch rep = RuleMatch rep (BottomUp rep)

    type BottomUpRuleLoop rep = RuleLoop rep (BottomUp rep)

    type BottomUpRuleOp rep = RuleOp rep (BottomUp rep)

    type BottomUpRule rep = SimplificationRule rep (BottomUp rep)

    -- | A collection of top-down rules.
    type TopDownRules rep = Rules rep (TopDown rep)

    -- | A collection of bottom-up rules.
    type BottomUpRules rep = Rules rep (BottomUp rep)

    -- | A collection of both top-down and bottom-up rules.
    data RuleBook rep = RuleBook
      { forall rep. RuleBook rep -> TopDownRules rep
    bookTopDownRules :: TopDownRules rep,
        forall rep. RuleBook rep -> BottomUpRules rep
    bookBottomUpRules :: BottomUpRules rep
      }

    instance Semigroup (RuleBook rep) where
      RuleBook TopDownRules rep
    ts1 BottomUpRules rep
    bs1 <> :: RuleBook rep -> RuleBook rep -> RuleBook rep
    <> RuleBook TopDownRules rep
    ts2 BottomUpRules rep
    bs2 = TopDownRules rep -> BottomUpRules rep -> RuleBook rep
    forall rep. TopDownRules rep -> BottomUpRules rep -> RuleBook rep
    RuleBook (TopDownRules rep
    ts1 TopDownRules rep -> TopDownRules rep -> TopDownRules rep
    forall a. Semigroup a => a -> a -> a
    <> TopDownRules rep
    ts2) (BottomUpRules rep
    bs1 BottomUpRules rep -> BottomUpRules rep -> BottomUpRules rep
    forall a. Semigroup a => a -> a -> a
    <> BottomUpRules rep
    bs2)

    instance Monoid (RuleBook rep) where
      mempty :: RuleBook rep
    mempty = TopDownRules rep -> BottomUpRules rep -> RuleBook rep
    forall rep. TopDownRules rep -> BottomUpRules rep -> RuleBook rep
    RuleBook TopDownRules rep
    forall a. Monoid a => a
    mempty BottomUpRules rep
    forall a. Monoid a => a
    mempty

    -- | Construct a rule book from a collection of rules.
    ruleBook ::
      [TopDownRule m] ->
      [BottomUpRule m] ->
      RuleBook m
    ruleBook :: forall m. [TopDownRule m] -> [BottomUpRule m] -> RuleBook m
    ruleBook [TopDownRule m]
    topdowns [BottomUpRule m]
    bottomups =
      TopDownRules m -> BottomUpRules m -> RuleBook m
    forall rep. TopDownRules rep -> BottomUpRules rep -> RuleBook rep
    RuleBook ([TopDownRule m] -> TopDownRules m
    forall m a. [SimplificationRule m a] -> Rules m a
    groupRules [TopDownRule m]
    topdowns) ([BottomUpRule m] -> BottomUpRules m
    forall m a. [SimplificationRule m a] -> Rules m a
    groupRules [BottomUpRule m]
    bottomups)
      where
        groupRules :: [SimplificationRule m a] -> Rules m a
        groupRules :: forall m a. [SimplificationRule m a] -> Rules m a
    groupRules [SimplificationRule m a]
    rs =
          Rules
            { rulesAny :: [SimplificationRule m a]
    rulesAny = [SimplificationRule m a]
    rs,
              rulesBasicOp :: [SimplificationRule m a]
    rulesBasicOp = (SimplificationRule m a -> Bool)
    -> [SimplificationRule m a] -> [SimplificationRule m a]
    forall a. (a -> Bool) -> [a] -> [a]
    filter SimplificationRule m a -> Bool
    forall {rep} {a}. SimplificationRule rep a -> Bool
    forBasicOp [SimplificationRule m a]
    rs,
              rulesMatch :: [SimplificationRule m a]
    rulesMatch = (SimplificationRule m a -> Bool)
    -> [SimplificationRule m a] -> [SimplificationRule m a]
    forall a. (a -> Bool) -> [a] -> [a]
    filter SimplificationRule m a -> Bool
    forall {rep} {a}. SimplificationRule rep a -> Bool
    forMatch [SimplificationRule m a]
    rs,
              rulesLoop :: [SimplificationRule m a]
    rulesLoop = (SimplificationRule m a -> Bool)
    -> [SimplificationRule m a] -> [SimplificationRule m a]
    forall a. (a -> Bool) -> [a] -> [a]
    filter SimplificationRule m a -> Bool
    forall {rep} {a}. SimplificationRule rep a -> Bool
    forLoop [SimplificationRule m a]
    rs,
              rulesOp :: [SimplificationRule m a]
    rulesOp = (SimplificationRule m a -> Bool)
    -> [SimplificationRule m a] -> [SimplificationRule m a]
    forall a. (a -> Bool) -> [a] -> [a]
    filter SimplificationRule m a -> Bool
    forall {rep} {a}. SimplificationRule rep a -> Bool
    forOp [SimplificationRule m a]
    rs
            }

        forBasicOp :: SimplificationRule rep a -> Bool
    forBasicOp RuleBasicOp {} = Bool
    True
        forBasicOp RuleGeneric {} = Bool
    True
        forBasicOp SimplificationRule rep a
    _ = Bool
    False

        forMatch :: SimplificationRule rep a -> Bool
    forMatch RuleMatch {} = Bool
    True
        forMatch RuleGeneric {} = Bool
    True
        forMatch SimplificationRule rep a
    _ = Bool
    False

        forLoop :: SimplificationRule rep a -> Bool
    forLoop RuleLoop {} = Bool
    True
        forLoop RuleGeneric {} = Bool
    True
        forLoop SimplificationRule rep a
    _ = Bool
    False

        forOp :: SimplificationRule rep a -> Bool
    forOp RuleOp {} = Bool
    True
        forOp RuleGeneric {} = Bool
    True
        forOp SimplificationRule rep a
    _ = Bool
    False

    -- | @simplifyStm lookup stm@ performs simplification of the
    -- binding @stm@.  If simplification is possible, a replacement list
    -- of bindings is returned, that bind at least the same names as the
    -- original binding (and possibly more, for intermediate results).
    topDownSimplifyStm ::
      (MonadFreshNames m, HasScope rep m, PrettyRep rep) =>
      RuleBook rep ->
      ST.SymbolTable rep ->
      Stm rep ->
      m (Maybe (Stms rep))
    topDownSimplifyStm :: forall (m :: * -> *) rep.
    (MonadFreshNames m, HasScope rep m, PrettyRep rep) =>
    RuleBook rep -> SymbolTable rep -> Stm rep -> m (Maybe (Stms rep))
    topDownSimplifyStm = Rules rep (SymbolTable rep)
    -> SymbolTable rep -> Stm rep -> m (Maybe (Stms rep))
    forall (m :: * -> *) rep a.
    (MonadFreshNames m, HasScope rep m, PrettyRep rep) =>
    Rules rep a -> a -> Stm rep -> m (Maybe (Stms rep))
    applyRules (Rules rep (SymbolTable rep)
     -> SymbolTable rep -> Stm rep -> m (Maybe (Stms rep)))
    -> (RuleBook rep -> Rules rep (SymbolTable rep))
    -> RuleBook rep
    -> SymbolTable rep
    -> Stm rep
    -> m (Maybe (Stms rep))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . RuleBook rep -> Rules rep (SymbolTable rep)
    forall rep. RuleBook rep -> TopDownRules rep
    bookTopDownRules

    -- | @simplifyStm uses stm@ performs simplification of the binding
    -- @stm@.  If simplification is possible, a replacement list of
    -- bindings is returned, that bind at least the same names as the
    -- original binding (and possibly more, for intermediate results).
    -- The first argument is the set of names used after this binding.
    bottomUpSimplifyStm ::
      (MonadFreshNames m, HasScope rep m, PrettyRep rep) =>
      RuleBook rep ->
      (ST.SymbolTable rep, UT.UsageTable) ->
      Stm rep ->
      m (Maybe (Stms rep))
    bottomUpSimplifyStm :: forall (m :: * -> *) rep.
    (MonadFreshNames m, HasScope rep m, PrettyRep rep) =>
    RuleBook rep
    -> (SymbolTable rep, UsageTable) -> Stm rep -> m (Maybe (Stms rep))
    bottomUpSimplifyStm = Rules rep (SymbolTable rep, UsageTable)
    -> (SymbolTable rep, UsageTable) -> Stm rep -> m (Maybe (Stms rep))
    forall (m :: * -> *) rep a.
    (MonadFreshNames m, HasScope rep m, PrettyRep rep) =>
    Rules rep a -> a -> Stm rep -> m (Maybe (Stms rep))
    applyRules (Rules rep (SymbolTable rep, UsageTable)
     -> (SymbolTable rep, UsageTable)
     -> Stm rep
     -> m (Maybe (Stms rep)))
    -> (RuleBook rep -> Rules rep (SymbolTable rep, UsageTable))
    -> RuleBook rep
    -> (SymbolTable rep, UsageTable)
    -> Stm rep
    -> m (Maybe (Stms rep))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . RuleBook rep -> Rules rep (SymbolTable rep, UsageTable)
    forall rep. RuleBook rep -> BottomUpRules rep
    bookBottomUpRules

    rulesForStm :: Stm rep -> Rules rep a -> [SimplificationRule rep a]
    rulesForStm :: forall rep a. Stm rep -> Rules rep a -> [SimplificationRule rep a]
    rulesForStm Stm rep
    stm = case Stm rep -> Exp rep
    forall rep. Stm rep -> Exp rep
    stmExp Stm rep
    stm of
      BasicOp {} -> Rules rep a -> [SimplificationRule rep a]
    forall rep a. Rules rep a -> [SimplificationRule rep a]
    rulesBasicOp
      Loop {} -> Rules rep a -> [SimplificationRule rep a]
    forall rep a. Rules rep a -> [SimplificationRule rep a]
    rulesLoop
      Op {} -> Rules rep a -> [SimplificationRule rep a]
    forall rep a. Rules rep a -> [SimplificationRule rep a]
    rulesOp
      Match {} -> Rules rep a -> [SimplificationRule rep a]
    forall rep a. Rules rep a -> [SimplificationRule rep a]
    rulesMatch
      Exp rep
    _ -> Rules rep a -> [SimplificationRule rep a]
    forall rep a. Rules rep a -> [SimplificationRule rep a]
    rulesAny

    applyRule :: SimplificationRule rep a -> a -> Stm rep -> Rule rep
    applyRule :: forall rep a. SimplificationRule rep a -> a -> Stm rep -> Rule rep
    applyRule (RuleGeneric RuleGeneric rep a
    f) a
    a Stm rep
    stm = RuleGeneric rep a
    f a
    a Stm rep
    stm
    applyRule (RuleBasicOp RuleBasicOp rep a
    f) a
    a (Let Pat (LetDec rep)
    pat StmAux (ExpDec rep)
    aux (BasicOp BasicOp
    e)) = RuleBasicOp rep a
    f a
    a Pat (LetDec rep)
    pat StmAux (ExpDec rep)
    aux BasicOp
    e
    applyRule (RuleLoop RuleLoop rep a
    f) a
    a (Let Pat (LetDec rep)
    pat StmAux (ExpDec rep)
    aux (Loop [(FParam rep, SubExp)]
    merge LoopForm
    form Body rep
    body)) =
      RuleLoop rep a
    f a
    a Pat (LetDec rep)
    pat StmAux (ExpDec rep)
    aux ([(FParam rep, SubExp)]
    merge, LoopForm
    form, Body rep
    body)
    applyRule (RuleMatch RuleMatch rep a
    f) a
    a (Let Pat (LetDec rep)
    pat StmAux (ExpDec rep)
    aux (Match [SubExp]
    cond [Case (Body rep)]
    cases Body rep
    defbody MatchDec (BranchType rep)
    ifsort)) =
      RuleMatch rep a
    f a
    a Pat (LetDec rep)
    pat StmAux (ExpDec rep)
    aux ([SubExp]
    cond, [Case (Body rep)]
    cases, Body rep
    defbody, MatchDec (BranchType rep)
    ifsort)
    applyRule (RuleOp RuleOp rep a
    f) a
    a (Let Pat (LetDec rep)
    pat StmAux (ExpDec rep)
    aux (Op Op rep
    op)) =
      RuleOp rep a
    f a
    a Pat (LetDec rep)
    pat StmAux (ExpDec rep)
    aux Op rep
    op
    applyRule SimplificationRule rep a
    _ a
    _ Stm rep
    _ =
      Rule rep
    forall rep. Rule rep
    Skip

    applyRules ::
      (MonadFreshNames m, HasScope rep m, PrettyRep rep) =>
      Rules rep a ->
      a ->
      Stm rep ->
      m (Maybe (Stms rep))
    applyRules :: forall (m :: * -> *) rep a.
    (MonadFreshNames m, HasScope rep m, PrettyRep rep) =>
    Rules rep a -> a -> Stm rep -> m (Maybe (Stms rep))
    applyRules Rules rep a
    all_rules a
    context Stm rep
    stm = do
      scope <- m (Scope rep)
    forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
    askScope

      modifyNameSource $ \VNameSource
    src ->
        let applyRules' :: [SimplificationRule rep a] -> Maybe (Stms rep, VNameSource)
    applyRules' [] = Maybe (Stms rep, VNameSource)
    forall a. Maybe a
    Nothing
            applyRules' (SimplificationRule rep a
    rule : [SimplificationRule rep a]
    rules) =
              case Scope rep
    -> VNameSource -> Rule rep -> Maybe (Stms rep, VNameSource)
    forall rep.
    Scope rep
    -> VNameSource -> Rule rep -> Maybe (Stms rep, VNameSource)
    simplify Scope rep
    scope VNameSource
    src (SimplificationRule rep a -> a -> Stm rep -> Rule rep
    forall rep a. SimplificationRule rep a -> a -> Stm rep -> Rule rep
    applyRule SimplificationRule rep a
    rule a
    context Stm rep
    stm) of
                Just (Stms rep, VNameSource)
    x -> (Stms rep, VNameSource) -> Maybe (Stms rep, VNameSource)
    forall a. a -> Maybe a
    Just (Stms rep, VNameSource)
    x
                Maybe (Stms rep, VNameSource)
    Nothing -> [SimplificationRule rep a] -> Maybe (Stms rep, VNameSource)
    applyRules' [SimplificationRule rep a]
    rules
         in case [SimplificationRule rep a] -> Maybe (Stms rep, VNameSource)
    applyRules' ([SimplificationRule rep a] -> Maybe (Stms rep, VNameSource))
    -> [SimplificationRule rep a] -> Maybe (Stms rep, VNameSource)
    forall a b. (a -> b) -> a -> b
    $ Stm rep -> Rules rep a -> [SimplificationRule rep a]
    forall rep a. Stm rep -> Rules rep a -> [SimplificationRule rep a]
    rulesForStm Stm rep
    stm Rules rep a
    all_rules of
              Just (Stms rep
    stms, VNameSource
    src') -> (Stms rep -> Maybe (Stms rep)
    forall a. a -> Maybe a
    Just Stms rep
    stms, VNameSource
    src')
              Maybe (Stms rep, VNameSource)
    Nothing -> (Maybe (Stms rep)
    forall a. Maybe a
    Nothing, VNameSource
    src)