{-# 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 rep
scope <- m (Scope rep)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope

  (VNameSource -> (Maybe (Stms rep), VNameSource))
-> m (Maybe (Stms rep))
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource ((VNameSource -> (Maybe (Stms rep), VNameSource))
 -> m (Maybe (Stms rep)))
-> (VNameSource -> (Maybe (Stms rep), VNameSource))
-> m (Maybe (Stms rep))
forall a b. (a -> b) -> a -> b
$ \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)