{-# OPTIONS_GHC -fno-warn-orphans #-}
-- | Converting back and forth between 'PrimExp's.
module Futhark.Analysis.PrimExp.Convert
  (
    primExpToExp
  , primExpFromExp
  , primExpToSubExp
  , primExpFromSubExp
  , primExpFromSubExpM
  , replaceInPrimExp
  , replaceInPrimExpM
  , substituteInPrimExp

    -- * Module reexport
    , module Futhark.Analysis.PrimExp
  ) where

import qualified Control.Monad.Fail as Fail
import           Control.Monad.Identity
import           Data.Loc
import qualified Data.Map.Strict as M
import           Data.Maybe

import           Futhark.Analysis.PrimExp
import           Futhark.Construct
import           Futhark.Representation.AST

-- | Convert a 'PrimExp' to a Futhark expression.  The provided
-- function converts the leaves.
primExpToExp :: MonadBinder m =>
                (v -> m (Exp (Lore m))) -> PrimExp v -> m (Exp (Lore m))
primExpToExp :: (v -> m (Exp (Lore m))) -> PrimExp v -> m (Exp (Lore m))
primExpToExp v -> m (Exp (Lore m))
f (BinOpExp BinOp
op PrimExp v
x PrimExp v
y) =
  BasicOp (Lore m) -> Exp (Lore m)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Lore m) -> Exp (Lore m))
-> m (BasicOp (Lore m)) -> m (Exp (Lore m))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (BinOp -> SubExp -> SubExp -> BasicOp (Lore m)
forall lore. BinOp -> SubExp -> SubExp -> BasicOp lore
BinOp BinOp
op
               (SubExp -> SubExp -> BasicOp (Lore m))
-> m SubExp -> m (SubExp -> BasicOp (Lore m))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> (v -> m (Exp (Lore m))) -> PrimExp v -> m SubExp
forall (m :: * -> *) v.
MonadBinder m =>
String -> (v -> m (Exp (Lore m))) -> PrimExp v -> m SubExp
primExpToSubExp String
"binop_x" v -> m (Exp (Lore m))
f PrimExp v
x
               m (SubExp -> BasicOp (Lore m)) -> m SubExp -> m (BasicOp (Lore m))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> String -> (v -> m (Exp (Lore m))) -> PrimExp v -> m SubExp
forall (m :: * -> *) v.
MonadBinder m =>
String -> (v -> m (Exp (Lore m))) -> PrimExp v -> m SubExp
primExpToSubExp String
"binop_y" v -> m (Exp (Lore m))
f PrimExp v
y)
primExpToExp v -> m (Exp (Lore m))
f (CmpOpExp CmpOp
op PrimExp v
x PrimExp v
y) =
  BasicOp (Lore m) -> Exp (Lore m)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Lore m) -> Exp (Lore m))
-> m (BasicOp (Lore m)) -> m (Exp (Lore m))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (CmpOp -> SubExp -> SubExp -> BasicOp (Lore m)
forall lore. CmpOp -> SubExp -> SubExp -> BasicOp lore
CmpOp CmpOp
op
               (SubExp -> SubExp -> BasicOp (Lore m))
-> m SubExp -> m (SubExp -> BasicOp (Lore m))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> (v -> m (Exp (Lore m))) -> PrimExp v -> m SubExp
forall (m :: * -> *) v.
MonadBinder m =>
String -> (v -> m (Exp (Lore m))) -> PrimExp v -> m SubExp
primExpToSubExp String
"cmpop_x" v -> m (Exp (Lore m))
f PrimExp v
x
               m (SubExp -> BasicOp (Lore m)) -> m SubExp -> m (BasicOp (Lore m))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> String -> (v -> m (Exp (Lore m))) -> PrimExp v -> m SubExp
forall (m :: * -> *) v.
MonadBinder m =>
String -> (v -> m (Exp (Lore m))) -> PrimExp v -> m SubExp
primExpToSubExp String
"cmpop_y" v -> m (Exp (Lore m))
f PrimExp v
y)
primExpToExp v -> m (Exp (Lore m))
f (UnOpExp UnOp
op PrimExp v
x) =
  BasicOp (Lore m) -> Exp (Lore m)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Lore m) -> Exp (Lore m))
-> m (BasicOp (Lore m)) -> m (Exp (Lore m))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (UnOp -> SubExp -> BasicOp (Lore m)
forall lore. UnOp -> SubExp -> BasicOp lore
UnOp UnOp
op (SubExp -> BasicOp (Lore m)) -> m SubExp -> m (BasicOp (Lore m))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> (v -> m (Exp (Lore m))) -> PrimExp v -> m SubExp
forall (m :: * -> *) v.
MonadBinder m =>
String -> (v -> m (Exp (Lore m))) -> PrimExp v -> m SubExp
primExpToSubExp String
"unop_x" v -> m (Exp (Lore m))
f PrimExp v
x)
primExpToExp v -> m (Exp (Lore m))
f (ConvOpExp ConvOp
op PrimExp v
x) =
  BasicOp (Lore m) -> Exp (Lore m)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Lore m) -> Exp (Lore m))
-> m (BasicOp (Lore m)) -> m (Exp (Lore m))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ConvOp -> SubExp -> BasicOp (Lore m)
forall lore. ConvOp -> SubExp -> BasicOp lore
ConvOp ConvOp
op (SubExp -> BasicOp (Lore m)) -> m SubExp -> m (BasicOp (Lore m))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> (v -> m (Exp (Lore m))) -> PrimExp v -> m SubExp
forall (m :: * -> *) v.
MonadBinder m =>
String -> (v -> m (Exp (Lore m))) -> PrimExp v -> m SubExp
primExpToSubExp String
"convop_x" v -> m (Exp (Lore m))
f PrimExp v
x)
primExpToExp v -> m (Exp (Lore m))
_ (ValueExp PrimValue
v) =
  Exp (Lore m) -> m (Exp (Lore m))
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp (Lore m) -> m (Exp (Lore m)))
-> Exp (Lore m) -> m (Exp (Lore m))
forall a b. (a -> b) -> a -> b
$ BasicOp (Lore m) -> Exp (Lore m)
forall lore. BasicOp lore -> ExpT lore
BasicOp (BasicOp (Lore m) -> Exp (Lore m))
-> BasicOp (Lore m) -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp (Lore m)
forall lore. SubExp -> BasicOp lore
SubExp (SubExp -> BasicOp (Lore m)) -> SubExp -> BasicOp (Lore m)
forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
Constant PrimValue
v
primExpToExp v -> m (Exp (Lore m))
f (FunExp String
h [PrimExp v]
args PrimType
t) =
  Name
-> [(SubExp, Diet)]
-> [RetType (Lore m)]
-> (Safety, SrcLoc, [SrcLoc])
-> Exp (Lore m)
forall lore.
Name
-> [(SubExp, Diet)]
-> [RetType lore]
-> (Safety, SrcLoc, [SrcLoc])
-> ExpT lore
Apply (String -> Name
nameFromString String
h) ([(SubExp, Diet)]
 -> [RetType (Lore m)]
 -> (Safety, SrcLoc, [SrcLoc])
 -> Exp (Lore m))
-> m [(SubExp, Diet)]
-> m ([RetType (Lore m)]
      -> (Safety, SrcLoc, [SrcLoc]) -> Exp (Lore m))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m [(SubExp, Diet)]
args' m ([RetType (Lore m)]
   -> (Safety, SrcLoc, [SrcLoc]) -> Exp (Lore m))
-> m [RetType (Lore m)]
-> m ((Safety, SrcLoc, [SrcLoc]) -> Exp (Lore m))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [RetType (Lore m)] -> m [RetType (Lore m)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimType -> RetType (Lore m)
forall rt. IsRetType rt => PrimType -> rt
primRetType PrimType
t] m ((Safety, SrcLoc, [SrcLoc]) -> Exp (Lore m))
-> m (Safety, SrcLoc, [SrcLoc]) -> m (Exp (Lore m))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
  (Safety, SrcLoc, [SrcLoc]) -> m (Safety, SrcLoc, [SrcLoc])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Safety
Safe, SrcLoc
forall a. IsLocation a => a
noLoc, [])
  where args' :: m [(SubExp, Diet)]
args' = [SubExp] -> [Diet] -> [(SubExp, Diet)]
forall a b. [a] -> [b] -> [(a, b)]
zip ([SubExp] -> [Diet] -> [(SubExp, Diet)])
-> m [SubExp] -> m ([Diet] -> [(SubExp, Diet)])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PrimExp v -> m SubExp) -> [PrimExp v] -> m [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> (v -> m (Exp (Lore m))) -> PrimExp v -> m SubExp
forall (m :: * -> *) v.
MonadBinder m =>
String -> (v -> m (Exp (Lore m))) -> PrimExp v -> m SubExp
primExpToSubExp String
"apply_arg" v -> m (Exp (Lore m))
f) [PrimExp v]
args m ([Diet] -> [(SubExp, Diet)]) -> m [Diet] -> m [(SubExp, Diet)]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Diet] -> m [Diet]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Diet -> [Diet]
forall a. a -> [a]
repeat Diet
Observe)
primExpToExp v -> m (Exp (Lore m))
f (LeafExp v
v PrimType
_) =
  v -> m (Exp (Lore m))
f v
v

instance ToExp v => ToExp (PrimExp v) where
  toExp :: PrimExp v -> m (Exp (Lore m))
toExp = (v -> m (Exp (Lore m))) -> PrimExp v -> m (Exp (Lore m))
forall (m :: * -> *) v.
MonadBinder m =>
(v -> m (Exp (Lore m))) -> PrimExp v -> m (Exp (Lore m))
primExpToExp v -> m (Exp (Lore m))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Lore m))
toExp

primExpToSubExp :: MonadBinder m =>
                   String -> (v -> m (Exp (Lore m))) -> PrimExp v -> m SubExp
primExpToSubExp :: String -> (v -> m (Exp (Lore m))) -> PrimExp v -> m SubExp
primExpToSubExp String
s v -> m (Exp (Lore m))
f PrimExp v
e = String -> Exp (Lore m) -> m SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
s (Exp (Lore m) -> m SubExp) -> m (Exp (Lore m)) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (v -> m (Exp (Lore m))) -> PrimExp v -> m (Exp (Lore m))
forall (m :: * -> *) v.
MonadBinder m =>
(v -> m (Exp (Lore m))) -> PrimExp v -> m (Exp (Lore m))
primExpToExp v -> m (Exp (Lore m))
f PrimExp v
e

-- | Convert an expression to a 'PrimExp'.  The provided function is
-- used to convert expressions that are not trivially 'PrimExp's.
-- This includes constants and variable names, which are passed as
-- 'SubExp's.
primExpFromExp :: (Fail.MonadFail m, Annotations lore) =>
                  (VName -> m (PrimExp v)) -> Exp lore -> m (PrimExp v)
primExpFromExp :: (VName -> m (PrimExp v)) -> Exp lore -> m (PrimExp v)
primExpFromExp VName -> m (PrimExp v)
f (BasicOp (BinOp BinOp
op SubExp
x SubExp
y)) =
  BinOp -> PrimExp v -> PrimExp v -> PrimExp v
forall v. BinOp -> PrimExp v -> PrimExp v -> PrimExp v
BinOpExp BinOp
op (PrimExp v -> PrimExp v -> PrimExp v)
-> m (PrimExp v) -> m (PrimExp v -> PrimExp v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> m (PrimExp v)) -> SubExp -> m (PrimExp v)
forall (m :: * -> *) v.
Applicative m =>
(VName -> m (PrimExp v)) -> SubExp -> m (PrimExp v)
primExpFromSubExpM VName -> m (PrimExp v)
f SubExp
x m (PrimExp v -> PrimExp v) -> m (PrimExp v) -> m (PrimExp v)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (VName -> m (PrimExp v)) -> SubExp -> m (PrimExp v)
forall (m :: * -> *) v.
Applicative m =>
(VName -> m (PrimExp v)) -> SubExp -> m (PrimExp v)
primExpFromSubExpM VName -> m (PrimExp v)
f SubExp
y
primExpFromExp VName -> m (PrimExp v)
f (BasicOp (CmpOp CmpOp
op SubExp
x SubExp
y)) =
  CmpOp -> PrimExp v -> PrimExp v -> PrimExp v
forall v. CmpOp -> PrimExp v -> PrimExp v -> PrimExp v
CmpOpExp CmpOp
op (PrimExp v -> PrimExp v -> PrimExp v)
-> m (PrimExp v) -> m (PrimExp v -> PrimExp v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> m (PrimExp v)) -> SubExp -> m (PrimExp v)
forall (m :: * -> *) v.
Applicative m =>
(VName -> m (PrimExp v)) -> SubExp -> m (PrimExp v)
primExpFromSubExpM VName -> m (PrimExp v)
f SubExp
x m (PrimExp v -> PrimExp v) -> m (PrimExp v) -> m (PrimExp v)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (VName -> m (PrimExp v)) -> SubExp -> m (PrimExp v)
forall (m :: * -> *) v.
Applicative m =>
(VName -> m (PrimExp v)) -> SubExp -> m (PrimExp v)
primExpFromSubExpM VName -> m (PrimExp v)
f SubExp
y
primExpFromExp VName -> m (PrimExp v)
f (BasicOp (UnOp UnOp
op SubExp
x)) =
  UnOp -> PrimExp v -> PrimExp v
forall v. UnOp -> PrimExp v -> PrimExp v
UnOpExp UnOp
op (PrimExp v -> PrimExp v) -> m (PrimExp v) -> m (PrimExp v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> m (PrimExp v)) -> SubExp -> m (PrimExp v)
forall (m :: * -> *) v.
Applicative m =>
(VName -> m (PrimExp v)) -> SubExp -> m (PrimExp v)
primExpFromSubExpM VName -> m (PrimExp v)
f SubExp
x
primExpFromExp VName -> m (PrimExp v)
f (BasicOp (ConvOp ConvOp
op SubExp
x)) =
  ConvOp -> PrimExp v -> PrimExp v
forall v. ConvOp -> PrimExp v -> PrimExp v
ConvOpExp ConvOp
op (PrimExp v -> PrimExp v) -> m (PrimExp v) -> m (PrimExp v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> m (PrimExp v)) -> SubExp -> m (PrimExp v)
forall (m :: * -> *) v.
Applicative m =>
(VName -> m (PrimExp v)) -> SubExp -> m (PrimExp v)
primExpFromSubExpM VName -> m (PrimExp v)
f SubExp
x
primExpFromExp VName -> m (PrimExp v)
_ (BasicOp (SubExp (Constant PrimValue
v))) =
  PrimExp v -> m (PrimExp v)
forall (m :: * -> *) a. Monad m => a -> m a
return (PrimExp v -> m (PrimExp v)) -> PrimExp v -> m (PrimExp v)
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimExp v
forall v. PrimValue -> PrimExp v
ValueExp PrimValue
v
primExpFromExp VName -> m (PrimExp v)
f (Apply Name
fname [(SubExp, Diet)]
args [RetType lore]
ts (Safety, SrcLoc, [SrcLoc])
_)
  | Name -> Bool
isBuiltInFunction Name
fname, [Prim PrimType
t] <- [RetType lore] -> [DeclExtType]
forall rt. IsRetType rt => [rt] -> [DeclExtType]
retTypeValues [RetType lore]
ts =
      String -> [PrimExp v] -> PrimType -> PrimExp v
forall v. String -> [PrimExp v] -> PrimType -> PrimExp v
FunExp (Name -> String
nameToString Name
fname) ([PrimExp v] -> PrimType -> PrimExp v)
-> m [PrimExp v] -> m (PrimType -> PrimExp v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((SubExp, Diet) -> m (PrimExp v))
-> [(SubExp, Diet)] -> m [PrimExp v]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((VName -> m (PrimExp v)) -> SubExp -> m (PrimExp v)
forall (m :: * -> *) v.
Applicative m =>
(VName -> m (PrimExp v)) -> SubExp -> m (PrimExp v)
primExpFromSubExpM VName -> m (PrimExp v)
f (SubExp -> m (PrimExp v))
-> ((SubExp, Diet) -> SubExp) -> (SubExp, Diet) -> m (PrimExp v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SubExp, Diet) -> SubExp
forall a b. (a, b) -> a
fst) [(SubExp, Diet)]
args m (PrimType -> PrimExp v) -> m PrimType -> m (PrimExp v)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> PrimType -> m PrimType
forall (f :: * -> *) a. Applicative f => a -> f a
pure PrimType
t
primExpFromExp VName -> m (PrimExp v)
_ Exp lore
_ = String -> m (PrimExp v)
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Not a PrimExp"

primExpFromSubExpM :: Applicative m => (VName -> m (PrimExp v)) -> SubExp -> m (PrimExp v)
primExpFromSubExpM :: (VName -> m (PrimExp v)) -> SubExp -> m (PrimExp v)
primExpFromSubExpM VName -> m (PrimExp v)
f (Var VName
v) = VName -> m (PrimExp v)
f VName
v
primExpFromSubExpM VName -> m (PrimExp v)
_ (Constant PrimValue
v) = PrimExp v -> m (PrimExp v)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PrimExp v -> m (PrimExp v)) -> PrimExp v -> m (PrimExp v)
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimExp v
forall v. PrimValue -> PrimExp v
ValueExp PrimValue
v

-- | Convert 'SubExp's of a given type.
primExpFromSubExp :: PrimType -> SubExp -> PrimExp VName
primExpFromSubExp :: PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
t (Var VName
v)      = VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
v PrimType
t
primExpFromSubExp PrimType
_ (Constant PrimValue
v) = PrimValue -> PrimExp VName
forall v. PrimValue -> PrimExp v
ValueExp PrimValue
v

-- | Applying a monadic transformation to the leaves in a 'PrimExp'.
replaceInPrimExpM :: Monad m =>
                     (a -> PrimType -> m (PrimExp b)) ->
                     PrimExp a -> m (PrimExp b)
replaceInPrimExpM :: (a -> PrimType -> m (PrimExp b)) -> PrimExp a -> m (PrimExp b)
replaceInPrimExpM a -> PrimType -> m (PrimExp b)
f (LeafExp a
v PrimType
pt) =
  a -> PrimType -> m (PrimExp b)
f a
v PrimType
pt
replaceInPrimExpM a -> PrimType -> m (PrimExp b)
_ (ValueExp PrimValue
v) =
  PrimExp b -> m (PrimExp b)
forall (m :: * -> *) a. Monad m => a -> m a
return (PrimExp b -> m (PrimExp b)) -> PrimExp b -> m (PrimExp b)
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimExp b
forall v. PrimValue -> PrimExp v
ValueExp PrimValue
v
replaceInPrimExpM a -> PrimType -> m (PrimExp b)
f (BinOpExp BinOp
bop PrimExp a
pe1 PrimExp a
pe2) =
  PrimExp b -> PrimExp b
forall v. PrimExp v -> PrimExp v
constFoldPrimExp (PrimExp b -> PrimExp b) -> m (PrimExp b) -> m (PrimExp b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
  (BinOp -> PrimExp b -> PrimExp b -> PrimExp b
forall v. BinOp -> PrimExp v -> PrimExp v -> PrimExp v
BinOpExp BinOp
bop (PrimExp b -> PrimExp b -> PrimExp b)
-> m (PrimExp b) -> m (PrimExp b -> PrimExp b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> PrimType -> m (PrimExp b)) -> PrimExp a -> m (PrimExp b)
forall (m :: * -> *) a b.
Monad m =>
(a -> PrimType -> m (PrimExp b)) -> PrimExp a -> m (PrimExp b)
replaceInPrimExpM a -> PrimType -> m (PrimExp b)
f PrimExp a
pe1 m (PrimExp b -> PrimExp b) -> m (PrimExp b) -> m (PrimExp b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (a -> PrimType -> m (PrimExp b)) -> PrimExp a -> m (PrimExp b)
forall (m :: * -> *) a b.
Monad m =>
(a -> PrimType -> m (PrimExp b)) -> PrimExp a -> m (PrimExp b)
replaceInPrimExpM a -> PrimType -> m (PrimExp b)
f PrimExp a
pe2)
replaceInPrimExpM a -> PrimType -> m (PrimExp b)
f (CmpOpExp CmpOp
cop PrimExp a
pe1 PrimExp a
pe2) =
  CmpOp -> PrimExp b -> PrimExp b -> PrimExp b
forall v. CmpOp -> PrimExp v -> PrimExp v -> PrimExp v
CmpOpExp CmpOp
cop (PrimExp b -> PrimExp b -> PrimExp b)
-> m (PrimExp b) -> m (PrimExp b -> PrimExp b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> PrimType -> m (PrimExp b)) -> PrimExp a -> m (PrimExp b)
forall (m :: * -> *) a b.
Monad m =>
(a -> PrimType -> m (PrimExp b)) -> PrimExp a -> m (PrimExp b)
replaceInPrimExpM a -> PrimType -> m (PrimExp b)
f PrimExp a
pe1 m (PrimExp b -> PrimExp b) -> m (PrimExp b) -> m (PrimExp b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (a -> PrimType -> m (PrimExp b)) -> PrimExp a -> m (PrimExp b)
forall (m :: * -> *) a b.
Monad m =>
(a -> PrimType -> m (PrimExp b)) -> PrimExp a -> m (PrimExp b)
replaceInPrimExpM a -> PrimType -> m (PrimExp b)
f PrimExp a
pe2
replaceInPrimExpM a -> PrimType -> m (PrimExp b)
f (UnOpExp UnOp
uop PrimExp a
pe) =
  UnOp -> PrimExp b -> PrimExp b
forall v. UnOp -> PrimExp v -> PrimExp v
UnOpExp UnOp
uop (PrimExp b -> PrimExp b) -> m (PrimExp b) -> m (PrimExp b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> PrimType -> m (PrimExp b)) -> PrimExp a -> m (PrimExp b)
forall (m :: * -> *) a b.
Monad m =>
(a -> PrimType -> m (PrimExp b)) -> PrimExp a -> m (PrimExp b)
replaceInPrimExpM a -> PrimType -> m (PrimExp b)
f PrimExp a
pe
replaceInPrimExpM a -> PrimType -> m (PrimExp b)
f (ConvOpExp ConvOp
cop PrimExp a
pe) =
  ConvOp -> PrimExp b -> PrimExp b
forall v. ConvOp -> PrimExp v -> PrimExp v
ConvOpExp ConvOp
cop (PrimExp b -> PrimExp b) -> m (PrimExp b) -> m (PrimExp b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> PrimType -> m (PrimExp b)) -> PrimExp a -> m (PrimExp b)
forall (m :: * -> *) a b.
Monad m =>
(a -> PrimType -> m (PrimExp b)) -> PrimExp a -> m (PrimExp b)
replaceInPrimExpM a -> PrimType -> m (PrimExp b)
f PrimExp a
pe
replaceInPrimExpM a -> PrimType -> m (PrimExp b)
f (FunExp String
h [PrimExp a]
args PrimType
t) =
  String -> [PrimExp b] -> PrimType -> PrimExp b
forall v. String -> [PrimExp v] -> PrimType -> PrimExp v
FunExp String
h ([PrimExp b] -> PrimType -> PrimExp b)
-> m [PrimExp b] -> m (PrimType -> PrimExp b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PrimExp a -> m (PrimExp b)) -> [PrimExp a] -> m [PrimExp b]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((a -> PrimType -> m (PrimExp b)) -> PrimExp a -> m (PrimExp b)
forall (m :: * -> *) a b.
Monad m =>
(a -> PrimType -> m (PrimExp b)) -> PrimExp a -> m (PrimExp b)
replaceInPrimExpM a -> PrimType -> m (PrimExp b)
f) [PrimExp a]
args m (PrimType -> PrimExp b) -> m PrimType -> m (PrimExp b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> PrimType -> m PrimType
forall (f :: * -> *) a. Applicative f => a -> f a
pure PrimType
t

replaceInPrimExp :: (a -> PrimType -> PrimExp b) ->
                    PrimExp a -> PrimExp b
replaceInPrimExp :: (a -> PrimType -> PrimExp b) -> PrimExp a -> PrimExp b
replaceInPrimExp a -> PrimType -> PrimExp b
f PrimExp a
e = Identity (PrimExp b) -> PrimExp b
forall a. Identity a -> a
runIdentity (Identity (PrimExp b) -> PrimExp b)
-> Identity (PrimExp b) -> PrimExp b
forall a b. (a -> b) -> a -> b
$ (a -> PrimType -> Identity (PrimExp b))
-> PrimExp a -> Identity (PrimExp b)
forall (m :: * -> *) a b.
Monad m =>
(a -> PrimType -> m (PrimExp b)) -> PrimExp a -> m (PrimExp b)
replaceInPrimExpM a -> PrimType -> Identity (PrimExp b)
forall (m :: * -> *). Monad m => a -> PrimType -> m (PrimExp b)
f' PrimExp a
e
  where f' :: a -> PrimType -> m (PrimExp b)
f' a
x PrimType
y = PrimExp b -> m (PrimExp b)
forall (m :: * -> *) a. Monad m => a -> m a
return (PrimExp b -> m (PrimExp b)) -> PrimExp b -> m (PrimExp b)
forall a b. (a -> b) -> a -> b
$ a -> PrimType -> PrimExp b
f a
x PrimType
y

-- | Substituting names in a PrimExp with other PrimExps
substituteInPrimExp :: Ord v => M.Map v (PrimExp v)
                    -> PrimExp v -> PrimExp v
substituteInPrimExp :: Map v (PrimExp v) -> PrimExp v -> PrimExp v
substituteInPrimExp Map v (PrimExp v)
tab = (v -> PrimType -> PrimExp v) -> PrimExp v -> PrimExp v
forall a b. (a -> PrimType -> PrimExp b) -> PrimExp a -> PrimExp b
replaceInPrimExp ((v -> PrimType -> PrimExp v) -> PrimExp v -> PrimExp v)
-> (v -> PrimType -> PrimExp v) -> PrimExp v -> PrimExp v
forall a b. (a -> b) -> a -> b
$ \v
v PrimType
t ->
  PrimExp v -> Maybe (PrimExp v) -> PrimExp v
forall a. a -> Maybe a -> a
fromMaybe (v -> PrimType -> PrimExp v
forall v. v -> PrimType -> PrimExp v
LeafExp v
v PrimType
t) (Maybe (PrimExp v) -> PrimExp v) -> Maybe (PrimExp v) -> PrimExp v
forall a b. (a -> b) -> a -> b
$ v -> Map v (PrimExp v) -> Maybe (PrimExp v)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup v
v Map v (PrimExp v)
tab