{-# 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,
    RuleDoLoop,

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

    -- * Bottom-up rules
    BottomUp,
    BottomUpRule,
    BottomUpRuleGeneric,
    BottomUpRuleBasicOp,
    BottomUpRuleMatch,
    BottomUpRuleDoLoop,
    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 -> 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
<$ :: forall a b. a -> RuleM rep b -> RuleM rep a
$c<$ :: forall rep a b. a -> RuleM rep b -> RuleM rep a
fmap :: forall a b. (a -> b) -> RuleM rep a -> RuleM rep b
$cfmap :: forall rep a b. (a -> b) -> RuleM rep a -> RuleM rep b
Functor,
      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
<* :: forall a b. RuleM rep a -> RuleM rep b -> RuleM rep a
$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 b
$c*> :: forall rep a b. RuleM rep a -> RuleM rep b -> RuleM rep b
liftA2 :: forall a b c.
(a -> b -> c) -> RuleM rep a -> RuleM rep b -> RuleM rep c
$cliftA2 :: forall rep a b c.
(a -> b -> c) -> RuleM rep a -> RuleM rep b -> RuleM rep c
<*> :: forall a b. RuleM rep (a -> b) -> RuleM rep a -> RuleM rep b
$c<*> :: forall rep a b. RuleM rep (a -> b) -> RuleM rep a -> RuleM rep b
pure :: forall a. a -> RuleM rep a
$cpure :: forall rep a. a -> RuleM rep a
Applicative,
      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
return :: forall a. a -> RuleM rep a
$creturn :: forall rep a. a -> RuleM rep a
>> :: 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 b
>>= :: forall a b. RuleM rep a -> (a -> RuleM rep b) -> RuleM rep b
$c>>= :: forall rep a b. RuleM rep a -> (a -> RuleM rep b) -> RuleM rep b
Monad,
      RuleM rep VNameSource
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
putNameSource :: VNameSource -> RuleM rep ()
$cputNameSource :: forall rep. VNameSource -> RuleM rep ()
getNameSource :: RuleM rep VNameSource
$cgetNameSource :: forall rep. RuleM rep VNameSource
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 = forall rep a.
BuilderT rep (StateT VNameSource Maybe) a -> RuleM rep a
RuleM forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadBuilder m =>
Pat (LetDec (Rep m)) -> Exp (Rep m) -> m (ExpDec (Rep m))
mkExpDecM Pat (LetDec (Rep (RuleM rep)))
pat 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 = forall rep a.
BuilderT rep (StateT VNameSource Maybe) a -> RuleM rep a
RuleM forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadBuilder m =>
Stms (Rep m) -> Result -> m (Body (Rep m))
mkBodyM 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 = forall rep a.
BuilderT rep (StateT VNameSource Maybe) a -> RuleM rep a
RuleM forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m (Stm (Rep m))
mkLetNamesM [VName]
pat Exp (Rep (RuleM rep))
e

  addStms :: Stms (Rep (RuleM rep)) -> RuleM rep ()
addStms = forall rep a.
BuilderT rep (StateT VNameSource Maybe) a -> RuleM rep a
RuleM forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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) = forall rep a.
BuilderT rep (StateT VNameSource Maybe) a -> RuleM rep a
RuleM forall a b. (a -> b) -> a -> b
$ 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 = forall a. Maybe a
Nothing
simplify Scope rep
scope VNameSource
src (Simplify (RuleM BuilderT rep (StateT VNameSource Maybe) ()
m)) =
  forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (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 = forall rep a.
BuilderT rep (StateT VNameSource Maybe) a -> RuleM rep a
RuleM forall a b. (a -> b) -> a -> b
$ forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a. Maybe a
Nothing

liftMaybe :: Maybe a -> RuleM rep a
liftMaybe :: forall a rep. Maybe a -> RuleM rep a
liftMaybe Maybe a
Nothing = forall rep a. RuleM rep a
cannotSimplify
liftMaybe (Just a
x) = 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 RuleDoLoop rep a =
  a ->
  Pat (LetDec rep) ->
  StmAux (ExpDec rep) ->
  ( [(FParam rep, SubExp)],
    LoopForm rep,
    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)
  | RuleDoLoop (RuleDoLoop 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]
rulesDoLoop :: [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 =
    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 forall a. Semigroup a => a -> a -> a
<> [SimplificationRule rep a]
as2) ([SimplificationRule rep a]
bs1 forall a. Semigroup a => a -> a -> a
<> [SimplificationRule rep a]
bs2) ([SimplificationRule rep a]
cs1 forall a. Semigroup a => a -> a -> a
<> [SimplificationRule rep a]
cs2) ([SimplificationRule rep a]
ds1 forall a. Semigroup a => a -> a -> a
<> [SimplificationRule rep a]
ds2) ([SimplificationRule rep a]
es1 forall a. Semigroup a => a -> a -> a
<> [SimplificationRule rep a]
es2)

instance Monoid (Rules rep a) where
  mempty :: Rules rep a
mempty = forall rep a.
[SimplificationRule rep a]
-> [SimplificationRule rep a]
-> [SimplificationRule rep a]
-> [SimplificationRule rep a]
-> [SimplificationRule rep a]
-> Rules rep a
Rules forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty 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 TopDownRuleDoLoop rep = RuleDoLoop 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 BottomUpRuleDoLoop rep = RuleDoLoop 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 = forall rep. TopDownRules rep -> BottomUpRules rep -> RuleBook rep
RuleBook (TopDownRules rep
ts1 forall a. Semigroup a => a -> a -> a
<> TopDownRules rep
ts2) (BottomUpRules rep
bs1 forall a. Semigroup a => a -> a -> a
<> BottomUpRules rep
bs2)

instance Monoid (RuleBook rep) where
  mempty :: RuleBook rep
mempty = forall rep. TopDownRules rep -> BottomUpRules rep -> RuleBook rep
RuleBook forall a. Monoid a => a
mempty 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 =
  forall rep. TopDownRules rep -> BottomUpRules rep -> RuleBook rep
RuleBook (forall m a. [SimplificationRule m a] -> Rules m a
groupRules [TopDownRule m]
topdowns) (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 = forall a. (a -> Bool) -> [a] -> [a]
filter forall {rep} {a}. SimplificationRule rep a -> Bool
forBasicOp [SimplificationRule m a]
rs,
          rulesMatch :: [SimplificationRule m a]
rulesMatch = forall a. (a -> Bool) -> [a] -> [a]
filter forall {rep} {a}. SimplificationRule rep a -> Bool
forMatch [SimplificationRule m a]
rs,
          rulesDoLoop :: [SimplificationRule m a]
rulesDoLoop = forall a. (a -> Bool) -> [a] -> [a]
filter forall {rep} {a}. SimplificationRule rep a -> Bool
forDoLoop [SimplificationRule m a]
rs,
          rulesOp :: [SimplificationRule m a]
rulesOp = forall a. (a -> Bool) -> [a] -> [a]
filter 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

    forDoLoop :: SimplificationRule rep a -> Bool
forDoLoop RuleDoLoop {} = Bool
True
    forDoLoop RuleGeneric {} = Bool
True
    forDoLoop 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 = forall (m :: * -> *) rep a.
(MonadFreshNames m, HasScope rep m, PrettyRep rep) =>
Rules rep a -> a -> Stm rep -> m (Maybe (Stms rep))
applyRules forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 = forall (m :: * -> *) rep a.
(MonadFreshNames m, HasScope rep m, PrettyRep rep) =>
Rules rep a -> a -> Stm rep -> m (Maybe (Stms rep))
applyRules forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 forall rep. Stm rep -> Exp rep
stmExp Stm rep
stm of
  BasicOp {} -> forall rep a. Rules rep a -> [SimplificationRule rep a]
rulesBasicOp
  DoLoop {} -> forall rep a. Rules rep a -> [SimplificationRule rep a]
rulesDoLoop
  Op {} -> forall rep a. Rules rep a -> [SimplificationRule rep a]
rulesOp
  Match {} -> forall rep a. Rules rep a -> [SimplificationRule rep a]
rulesMatch
  Exp rep
_ -> 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 (RuleDoLoop RuleDoLoop rep a
f) a
a (Let Pat (LetDec rep)
pat StmAux (ExpDec rep)
aux (DoLoop [(FParam rep, SubExp)]
merge LoopForm rep
form Body rep
body)) =
  RuleDoLoop rep a
f a
a Pat (LetDec rep)
pat StmAux (ExpDec rep)
aux ([(FParam rep, SubExp)]
merge, LoopForm rep
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
_ =
  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 rep
scope <- forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope

  forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource forall a b. (a -> b) -> a -> b
$ \VNameSource
src ->
    let applyRules' :: [SimplificationRule rep a] -> Maybe (Stms rep, VNameSource)
applyRules' [] = forall a. Maybe a
Nothing
        applyRules' (SimplificationRule rep a
rule : [SimplificationRule rep a]
rules) =
          case forall rep.
Scope rep
-> VNameSource -> Rule rep -> Maybe (Stms rep, VNameSource)
simplify Scope rep
scope VNameSource
src (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 -> 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' forall a b. (a -> b) -> a -> b
$ 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') -> (forall a. a -> Maybe a
Just Stms rep
stms, VNameSource
src')
          Maybe (Stms rep, VNameSource)
Nothing -> (forall a. Maybe a
Nothing, VNameSource
src)