-- | 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 {k} (rep :: k).
SimplifiableRep rep =>
PrimExp VName -> SimpleM rep (PrimExp VName)
simplifyPrimExp = forall {k} (rep :: k) a.
SimplifiableRep rep =>
(a -> PrimType -> SimpleM rep (PrimExp a))
-> PrimExp a -> SimpleM rep (PrimExp a)
simplifyAnyPrimExp forall {k} {rep :: k}.
(ASTRep rep, Simplifiable (LetDec rep),
 Simplifiable (FParamInfo rep), Simplifiable (LParamInfo rep),
 Simplifiable (RetType rep), Simplifiable (BranchType rep),
 TraverseOpStms (Wise rep), CanBeWise (Op rep),
 IndexOp (OpWithWisdom (Op 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 {k} (rep :: k).
(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 {k} (rep :: k).
SimplifiableRep rep =>
PrimExp (Ext VName) -> SimpleM rep (PrimExp (Ext VName))
simplifyExtPrimExp = forall {k} (rep :: k) a.
SimplifiableRep rep =>
(a -> PrimType -> SimpleM rep (PrimExp a))
-> PrimExp a -> SimpleM rep (PrimExp a)
simplifyAnyPrimExp forall {k} {rep :: k}.
(ASTRep rep, Simplifiable (LetDec rep),
 Simplifiable (FParamInfo rep), Simplifiable (LParamInfo rep),
 Simplifiable (RetType rep), Simplifiable (BranchType rep),
 TraverseOpStms (Wise rep), CanBeWise (Op rep),
 IndexOp (OpWithWisdom (Op 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 {k} (rep :: k).
(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 {k} (rep :: k) 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 {k} (rep :: k) 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 {k} (rep :: k) 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 {k} (rep :: k) 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 {k} (rep :: k) 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 {k} (rep :: k) 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 {k} (rep :: k) 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 {k} (rep :: k) 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