-- | Defines simplification functions for 'PrimExp's.
module Futhark.Analysis.PrimExp.Simplify (simplifyPrimExp, simplifyExtPrimExp) where

import Futhark.Analysis.PrimExp
import Futhark.IR
import Futhark.Optimise.Simplify.Engine as Engine

-- | Simplify a 'PrimExp', including copy propagation.  If a 'LeafExp'
-- refers to a name that is a 'Constant', the node turns into a
-- 'ValueExp'.
simplifyPrimExp ::
  SimplifiableRep rep =>
  PrimExp VName ->
  SimpleM rep (PrimExp VName)
simplifyPrimExp :: forall rep.
SimplifiableRep rep =>
PrimExp VName -> SimpleM rep (PrimExp VName)
simplifyPrimExp = forall rep a.
SimplifiableRep rep =>
(a -> PrimType -> SimpleM rep (PrimExp a))
-> PrimExp a -> SimpleM rep (PrimExp a)
simplifyAnyPrimExp forall {rep}.
(ASTRep rep, Simplifiable (LetDec rep),
 Simplifiable (FParamInfo rep), Simplifiable (LParamInfo rep),
 Simplifiable (RetType rep), Simplifiable (BranchType rep),
 TraverseOpStms (Wise rep), CanBeWise (OpC rep),
 IndexOp (OpC rep (Wise rep)), AliasedOp (OpC rep (Wise rep)),
 BuilderOps (Wise rep)) =>
VName -> PrimType -> SimpleM rep (PrimExp VName)
onLeaf
  where
    onLeaf :: VName -> PrimType -> SimpleM rep (PrimExp VName)
onLeaf VName
v PrimType
pt = do
      SubExp
se <- forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
v
      case SubExp
se of
        Var VName
v' -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall v. v -> PrimType -> PrimExp v
LeafExp VName
v' PrimType
pt
        Constant PrimValue
pv -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall v. PrimValue -> PrimExp v
ValueExp PrimValue
pv

-- | Like 'simplifyPrimExp', but where leaves may be 'Ext's.
simplifyExtPrimExp ::
  SimplifiableRep rep =>
  PrimExp (Ext VName) ->
  SimpleM rep (PrimExp (Ext VName))
simplifyExtPrimExp :: forall rep.
SimplifiableRep rep =>
PrimExp (Ext VName) -> SimpleM rep (PrimExp (Ext VName))
simplifyExtPrimExp = forall rep a.
SimplifiableRep rep =>
(a -> PrimType -> SimpleM rep (PrimExp a))
-> PrimExp a -> SimpleM rep (PrimExp a)
simplifyAnyPrimExp forall {rep}.
(ASTRep rep, Simplifiable (LetDec rep),
 Simplifiable (FParamInfo rep), Simplifiable (LParamInfo rep),
 Simplifiable (RetType rep), Simplifiable (BranchType rep),
 TraverseOpStms (Wise rep), CanBeWise (OpC rep),
 IndexOp (OpC rep (Wise rep)), AliasedOp (OpC rep (Wise rep)),
 BuilderOps (Wise rep)) =>
Ext VName -> PrimType -> SimpleM rep (PrimExp (Ext VName))
onLeaf
  where
    onLeaf :: Ext VName -> PrimType -> SimpleM rep (PrimExp (Ext VName))
onLeaf (Free VName
v) PrimType
pt = do
      SubExp
se <- forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
simplify forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
v
      case SubExp
se of
        Var VName
v' -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall v. v -> PrimType -> PrimExp v
LeafExp (forall a. a -> Ext a
Free VName
v') PrimType
pt
        Constant PrimValue
pv -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall v. PrimValue -> PrimExp v
ValueExp PrimValue
pv
    onLeaf (Ext Int
i) PrimType
pt = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall v. v -> PrimType -> PrimExp v
LeafExp (forall a. Int -> Ext a
Ext Int
i) PrimType
pt

simplifyAnyPrimExp ::
  SimplifiableRep rep =>
  (a -> PrimType -> SimpleM rep (PrimExp a)) ->
  PrimExp a ->
  SimpleM rep (PrimExp a)
simplifyAnyPrimExp :: forall rep a.
SimplifiableRep rep =>
(a -> PrimType -> SimpleM rep (PrimExp a))
-> PrimExp a -> SimpleM rep (PrimExp a)
simplifyAnyPrimExp a -> PrimType -> SimpleM rep (PrimExp a)
f (LeafExp a
v PrimType
pt) = a -> PrimType -> SimpleM rep (PrimExp a)
f a
v PrimType
pt
simplifyAnyPrimExp a -> PrimType -> SimpleM rep (PrimExp a)
_ (ValueExp PrimValue
pv) =
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall v. PrimValue -> PrimExp v
ValueExp PrimValue
pv
simplifyAnyPrimExp a -> PrimType -> SimpleM rep (PrimExp a)
f (BinOpExp BinOp
bop PrimExp a
e1 PrimExp a
e2) =
  forall v. BinOp -> PrimExp v -> PrimExp v -> PrimExp v
BinOpExp BinOp
bop forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep a.
SimplifiableRep rep =>
(a -> PrimType -> SimpleM rep (PrimExp a))
-> PrimExp a -> SimpleM rep (PrimExp a)
simplifyAnyPrimExp a -> PrimType -> SimpleM rep (PrimExp a)
f PrimExp a
e1 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall rep a.
SimplifiableRep rep =>
(a -> PrimType -> SimpleM rep (PrimExp a))
-> PrimExp a -> SimpleM rep (PrimExp a)
simplifyAnyPrimExp a -> PrimType -> SimpleM rep (PrimExp a)
f PrimExp a
e2
simplifyAnyPrimExp a -> PrimType -> SimpleM rep (PrimExp a)
f (CmpOpExp CmpOp
cmp PrimExp a
e1 PrimExp a
e2) =
  forall v. CmpOp -> PrimExp v -> PrimExp v -> PrimExp v
CmpOpExp CmpOp
cmp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep a.
SimplifiableRep rep =>
(a -> PrimType -> SimpleM rep (PrimExp a))
-> PrimExp a -> SimpleM rep (PrimExp a)
simplifyAnyPrimExp a -> PrimType -> SimpleM rep (PrimExp a)
f PrimExp a
e1 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall rep a.
SimplifiableRep rep =>
(a -> PrimType -> SimpleM rep (PrimExp a))
-> PrimExp a -> SimpleM rep (PrimExp a)
simplifyAnyPrimExp a -> PrimType -> SimpleM rep (PrimExp a)
f PrimExp a
e2
simplifyAnyPrimExp a -> PrimType -> SimpleM rep (PrimExp a)
f (UnOpExp UnOp
op PrimExp a
e) =
  forall v. UnOp -> PrimExp v -> PrimExp v
UnOpExp UnOp
op forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep a.
SimplifiableRep rep =>
(a -> PrimType -> SimpleM rep (PrimExp a))
-> PrimExp a -> SimpleM rep (PrimExp a)
simplifyAnyPrimExp a -> PrimType -> SimpleM rep (PrimExp a)
f PrimExp a
e
simplifyAnyPrimExp a -> PrimType -> SimpleM rep (PrimExp a)
f (ConvOpExp ConvOp
conv PrimExp a
e) =
  forall v. ConvOp -> PrimExp v -> PrimExp v
ConvOpExp ConvOp
conv forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep a.
SimplifiableRep rep =>
(a -> PrimType -> SimpleM rep (PrimExp a))
-> PrimExp a -> SimpleM rep (PrimExp a)
simplifyAnyPrimExp a -> PrimType -> SimpleM rep (PrimExp a)
f PrimExp a
e
simplifyAnyPrimExp a -> PrimType -> SimpleM rep (PrimExp a)
f (FunExp String
h [PrimExp a]
args PrimType
t) =
  forall v. String -> [PrimExp v] -> PrimType -> PrimExp v
FunExp String
h forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall rep a.
SimplifiableRep rep =>
(a -> PrimType -> SimpleM rep (PrimExp a))
-> PrimExp a -> SimpleM rep (PrimExp a)
simplifyAnyPrimExp a -> PrimType -> SimpleM rep (PrimExp a)
f) [PrimExp a]
args forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure PrimType
t