{-# LANGUAGE TypeFamilies #-}
module Futhark.Internalise.Monad
( InternaliseM,
runInternaliseM,
throwError,
VarSubsts,
InternaliseEnv (..),
FunInfo,
substitutingVars,
lookupSubst,
addOpaques,
addFunDef,
lookupFunction,
lookupFunction',
lookupConst,
bindFunction,
bindConstant,
assert,
module Futhark.Tools,
)
where
import Control.Monad.Except
import Control.Monad.Reader
import Control.Monad.State
import Data.Map.Strict qualified as M
import Futhark.IR.SOACS
import Futhark.MonadFreshNames
import Futhark.Tools
type FunInfo =
( [VName],
[DeclType],
[FParam SOACS],
[(SubExp, Type)] -> Maybe [DeclExtType]
)
type FunTable = M.Map VName FunInfo
type VarSubsts = M.Map VName [SubExp]
data InternaliseEnv = InternaliseEnv
{ InternaliseEnv -> VarSubsts
envSubsts :: VarSubsts,
InternaliseEnv -> Bool
envDoBoundsChecks :: Bool,
InternaliseEnv -> Bool
envSafe :: Bool,
InternaliseEnv -> Attrs
envAttrs :: Attrs
}
data InternaliseState = InternaliseState
{ InternaliseState -> VNameSource
stateNameSource :: VNameSource,
InternaliseState -> FunTable
stateFunTable :: FunTable,
InternaliseState -> VarSubsts
stateConstSubsts :: VarSubsts,
InternaliseState -> [FunDef SOACS]
stateFuns :: [FunDef SOACS],
InternaliseState -> OpaqueTypes
stateTypes :: OpaqueTypes
}
newtype InternaliseM a
= InternaliseM
(BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a)
deriving
( forall a b. a -> InternaliseM b -> InternaliseM a
forall a b. (a -> b) -> InternaliseM a -> InternaliseM 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 -> InternaliseM b -> InternaliseM a
$c<$ :: forall a b. a -> InternaliseM b -> InternaliseM a
fmap :: forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
$cfmap :: forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
Functor,
Functor InternaliseM
forall a. a -> InternaliseM a
forall a b. InternaliseM a -> InternaliseM b -> InternaliseM a
forall a b. InternaliseM a -> InternaliseM b -> InternaliseM b
forall a b.
InternaliseM (a -> b) -> InternaliseM a -> InternaliseM b
forall a b c.
(a -> b -> c) -> InternaliseM a -> InternaliseM b -> InternaliseM 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. InternaliseM a -> InternaliseM b -> InternaliseM a
$c<* :: forall a b. InternaliseM a -> InternaliseM b -> InternaliseM a
*> :: forall a b. InternaliseM a -> InternaliseM b -> InternaliseM b
$c*> :: forall a b. InternaliseM a -> InternaliseM b -> InternaliseM b
liftA2 :: forall a b c.
(a -> b -> c) -> InternaliseM a -> InternaliseM b -> InternaliseM c
$cliftA2 :: forall a b c.
(a -> b -> c) -> InternaliseM a -> InternaliseM b -> InternaliseM c
<*> :: forall a b.
InternaliseM (a -> b) -> InternaliseM a -> InternaliseM b
$c<*> :: forall a b.
InternaliseM (a -> b) -> InternaliseM a -> InternaliseM b
pure :: forall a. a -> InternaliseM a
$cpure :: forall a. a -> InternaliseM a
Applicative,
Applicative InternaliseM
forall a. a -> InternaliseM a
forall a b. InternaliseM a -> InternaliseM b -> InternaliseM b
forall a b.
InternaliseM a -> (a -> InternaliseM b) -> InternaliseM 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 -> InternaliseM a
$creturn :: forall a. a -> InternaliseM a
>> :: forall a b. InternaliseM a -> InternaliseM b -> InternaliseM b
$c>> :: forall a b. InternaliseM a -> InternaliseM b -> InternaliseM b
>>= :: forall a b.
InternaliseM a -> (a -> InternaliseM b) -> InternaliseM b
$c>>= :: forall a b.
InternaliseM a -> (a -> InternaliseM b) -> InternaliseM b
Monad,
MonadReader InternaliseEnv,
MonadState InternaliseState,
Monad InternaliseM
InternaliseM VNameSource
VNameSource -> InternaliseM ()
forall (m :: * -> *).
Monad m
-> m VNameSource -> (VNameSource -> m ()) -> MonadFreshNames m
putNameSource :: VNameSource -> InternaliseM ()
$cputNameSource :: VNameSource -> InternaliseM ()
getNameSource :: InternaliseM VNameSource
$cgetNameSource :: InternaliseM VNameSource
MonadFreshNames,
HasScope SOACS
)
instance LocalScope SOACS InternaliseM where
localScope :: forall a. Scope SOACS -> InternaliseM a -> InternaliseM a
localScope Scope SOACS
scope (InternaliseM BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
m) = do
Scope SOACS
old_scope <- forall {k} (rep :: k) (m :: * -> *).
HasScope rep m =>
m (Scope rep)
askScope
forall a.
BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
-> InternaliseM a
InternaliseM forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k) (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (Scope SOACS
scope forall k a b. Ord k => Map k a -> Map k b -> Map k a
`M.difference` Scope SOACS
old_scope) BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
m
instance MonadFreshNames (State InternaliseState) where
getNameSource :: State InternaliseState VNameSource
getNameSource = forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets InternaliseState -> VNameSource
stateNameSource
putNameSource :: VNameSource -> State InternaliseState ()
putNameSource VNameSource
src = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \InternaliseState
s -> InternaliseState
s {stateNameSource :: VNameSource
stateNameSource = VNameSource
src}
instance MonadBuilder InternaliseM where
type Rep InternaliseM = SOACS
mkExpDecM :: Pat (LetDec (Rep InternaliseM))
-> Exp (Rep InternaliseM)
-> InternaliseM (ExpDec (Rep InternaliseM))
mkExpDecM Pat (LetDec (Rep InternaliseM))
pat Exp (Rep InternaliseM)
e = forall a.
BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
-> InternaliseM a
InternaliseM 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 InternaliseM))
pat Exp (Rep InternaliseM)
e
mkBodyM :: Stms (Rep InternaliseM)
-> Result -> InternaliseM (Body (Rep InternaliseM))
mkBodyM Stms (Rep InternaliseM)
stms Result
res = forall a.
BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
-> InternaliseM a
InternaliseM forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadBuilder m =>
Stms (Rep m) -> Result -> m (Body (Rep m))
mkBodyM Stms (Rep InternaliseM)
stms Result
res
mkLetNamesM :: [VName]
-> Exp (Rep InternaliseM) -> InternaliseM (Stm (Rep InternaliseM))
mkLetNamesM [VName]
pat Exp (Rep InternaliseM)
e = forall a.
BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
-> InternaliseM a
InternaliseM forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m (Stm (Rep m))
mkLetNamesM [VName]
pat Exp (Rep InternaliseM)
e
addStms :: Stms (Rep InternaliseM) -> InternaliseM ()
addStms = forall a.
BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
-> InternaliseM a
InternaliseM forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
addStms
collectStms :: forall a.
InternaliseM a -> InternaliseM (a, Stms (Rep InternaliseM))
collectStms (InternaliseM BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
m) = forall a.
BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
-> InternaliseM a
InternaliseM forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
m
runInternaliseM ::
MonadFreshNames m =>
Bool ->
InternaliseM () ->
m (OpaqueTypes, Stms SOACS, [FunDef SOACS])
runInternaliseM :: forall (m :: * -> *).
MonadFreshNames m =>
Bool
-> InternaliseM () -> m (OpaqueTypes, Stms SOACS, [FunDef SOACS])
runInternaliseM Bool
safe (InternaliseM BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) ()
m) =
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource forall a b. (a -> b) -> a -> b
$ \VNameSource
src ->
let ((()
_, Stms SOACS
consts), InternaliseState
s) =
forall s a. State s a -> s -> (a, s)
runState (forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT (forall {k} (m :: * -> *) (rep :: k) a.
MonadFreshNames m =>
BuilderT rep m a -> Scope rep -> m (a, Stms rep)
runBuilderT BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) ()
m forall a. Monoid a => a
mempty) InternaliseEnv
newEnv) (VNameSource -> InternaliseState
newState VNameSource
src)
in ( (InternaliseState -> OpaqueTypes
stateTypes InternaliseState
s, Stms SOACS
consts, forall a. [a] -> [a]
reverse forall a b. (a -> b) -> a -> b
$ InternaliseState -> [FunDef SOACS]
stateFuns InternaliseState
s),
InternaliseState -> VNameSource
stateNameSource InternaliseState
s
)
where
newEnv :: InternaliseEnv
newEnv =
InternaliseEnv
{ envSubsts :: VarSubsts
envSubsts = forall a. Monoid a => a
mempty,
envDoBoundsChecks :: Bool
envDoBoundsChecks = Bool
True,
envSafe :: Bool
envSafe = Bool
safe,
envAttrs :: Attrs
envAttrs = forall a. Monoid a => a
mempty
}
newState :: VNameSource -> InternaliseState
newState VNameSource
src =
InternaliseState
{ stateNameSource :: VNameSource
stateNameSource = VNameSource
src,
stateFunTable :: FunTable
stateFunTable = forall a. Monoid a => a
mempty,
stateConstSubsts :: VarSubsts
stateConstSubsts = forall a. Monoid a => a
mempty,
stateFuns :: [FunDef SOACS]
stateFuns = forall a. Monoid a => a
mempty,
stateTypes :: OpaqueTypes
stateTypes = forall a. Monoid a => a
mempty
}
substitutingVars :: VarSubsts -> InternaliseM a -> InternaliseM a
substitutingVars :: forall a. VarSubsts -> InternaliseM a -> InternaliseM a
substitutingVars VarSubsts
substs = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall a b. (a -> b) -> a -> b
$ \InternaliseEnv
env -> InternaliseEnv
env {envSubsts :: VarSubsts
envSubsts = VarSubsts
substs forall a. Semigroup a => a -> a -> a
<> InternaliseEnv -> VarSubsts
envSubsts InternaliseEnv
env}
lookupSubst :: VName -> InternaliseM (Maybe [SubExp])
lookupSubst :: VName -> InternaliseM (Maybe [SubExp])
lookupSubst VName
v = do
Maybe [SubExp]
env_substs <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v forall b c a. (b -> c) -> (a -> b) -> a -> c
. InternaliseEnv -> VarSubsts
envSubsts
Maybe [SubExp]
const_substs <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v forall b c a. (b -> c) -> (a -> b) -> a -> c
. InternaliseState -> VarSubsts
stateConstSubsts
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Maybe [SubExp]
env_substs forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` Maybe [SubExp]
const_substs
addOpaques :: OpaqueTypes -> InternaliseM ()
addOpaques :: OpaqueTypes -> InternaliseM ()
addOpaques OpaqueTypes
ts = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \InternaliseState
s ->
InternaliseState
s {stateTypes :: OpaqueTypes
stateTypes = InternaliseState -> OpaqueTypes
stateTypes InternaliseState
s forall a. Semigroup a => a -> a -> a
<> OpaqueTypes
ts}
addFunDef :: FunDef SOACS -> InternaliseM ()
addFunDef :: FunDef SOACS -> InternaliseM ()
addFunDef FunDef SOACS
fd = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \InternaliseState
s -> InternaliseState
s {stateFuns :: [FunDef SOACS]
stateFuns = FunDef SOACS
fd forall a. a -> [a] -> [a]
: InternaliseState -> [FunDef SOACS]
stateFuns InternaliseState
s}
lookupFunction' :: VName -> InternaliseM (Maybe FunInfo)
lookupFunction' :: VName -> InternaliseM (Maybe FunInfo)
lookupFunction' VName
fname = forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
fname forall b c a. (b -> c) -> (a -> b) -> a -> c
. InternaliseState -> FunTable
stateFunTable
lookupFunction :: VName -> InternaliseM FunInfo
lookupFunction :: VName -> InternaliseM FunInfo
lookupFunction VName
fname = forall b a. b -> (a -> b) -> Maybe a -> b
maybe InternaliseM
([VName], [DeclType], [Param DeclType],
[(SubExp, Type)] -> Maybe [DeclExtType])
bad forall (f :: * -> *) a. Applicative f => a -> f a
pure forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> InternaliseM (Maybe FunInfo)
lookupFunction' VName
fname
where
bad :: InternaliseM
([VName], [DeclType], [Param DeclType],
[(SubExp, Type)] -> Maybe [DeclExtType])
bad = forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Internalise.lookupFunction: Function '" forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString VName
fname forall a. [a] -> [a] -> [a]
++ [Char]
"' not found."
lookupConst :: VName -> InternaliseM (Maybe [SubExp])
lookupConst :: VName -> InternaliseM (Maybe [SubExp])
lookupConst VName
fname = do
Bool
is_var <- forall {k} (rep :: k) (m :: * -> *) a.
HasScope rep m =>
(Scope rep -> a) -> m a
asksScope (VName
fname `M.member`)
Maybe [SubExp]
fname_subst <- VName -> InternaliseM (Maybe [SubExp])
lookupSubst VName
fname
case (Bool
is_var, Maybe [SubExp]
fname_subst) of
(Bool
_, Just [SubExp]
ses) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just [SubExp]
ses
(Bool
True, Maybe [SubExp]
_) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just [VName -> SubExp
Var VName
fname]
(Bool, Maybe [SubExp])
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
bindFunction :: VName -> FunDef SOACS -> FunInfo -> InternaliseM ()
bindFunction :: VName -> FunDef SOACS -> FunInfo -> InternaliseM ()
bindFunction VName
fname FunDef SOACS
fd FunInfo
info = do
FunDef SOACS -> InternaliseM ()
addFunDef FunDef SOACS
fd
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \InternaliseState
s -> InternaliseState
s {stateFunTable :: FunTable
stateFunTable = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
fname FunInfo
info forall a b. (a -> b) -> a -> b
$ InternaliseState -> FunTable
stateFunTable InternaliseState
s}
bindConstant :: VName -> FunDef SOACS -> InternaliseM ()
bindConstant :: VName -> FunDef SOACS -> InternaliseM ()
bindConstant VName
cname FunDef SOACS
fd = do
forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
addStms forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Body rep -> Stms rep
bodyStms forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). FunDef rep -> Body rep
funDefBody FunDef SOACS
fd
case forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k). Body rep -> Result
bodyResult forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k). FunDef rep -> Body rep
funDefBody forall a b. (a -> b) -> a -> b
$ FunDef SOACS
fd of
[SubExp
se] -> do
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
cname] forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
[SubExp]
ses -> do
let substs :: [SubExp]
substs =
forall a. Int -> [a] -> [a]
drop (forall (t :: * -> *) a. Foldable t => t a -> Int
length (forall u. [TypeBase ExtShape u] -> Set Int
shapeContext (forall {k} (rep :: k). FunDef rep -> [RetType rep]
funDefRetType FunDef SOACS
fd))) [SubExp]
ses
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \InternaliseState
s ->
InternaliseState
s
{ stateConstSubsts :: VarSubsts
stateConstSubsts = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
cname [SubExp]
substs forall a b. (a -> b) -> a -> b
$ InternaliseState -> VarSubsts
stateConstSubsts InternaliseState
s
}
assert ::
String ->
SubExp ->
ErrorMsg SubExp ->
SrcLoc ->
InternaliseM Certs
assert :: [Char] -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert [Char]
desc SubExp
se ErrorMsg SubExp
msg SrcLoc
loc = InternaliseM VName -> InternaliseM Certs
assertingOne forall a b. (a -> b) -> a -> b
$ do
Attrs
attrs <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a -> b) -> a -> b
$ Attrs -> Attrs
attrsForAssert forall b c a. (b -> c) -> (a -> b) -> a -> c
. InternaliseEnv -> Attrs
envAttrs
forall (m :: * -> *) a. MonadBuilder m => Attrs -> m a -> m a
attributing Attrs
attrs forall a b. (a -> b) -> a -> b
$
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
desc forall a b. (a -> b) -> a -> b
$
forall {k} (rep :: k). BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$
SubExp -> ErrorMsg SubExp -> (SrcLoc, [SrcLoc]) -> BasicOp
Assert SubExp
se ErrorMsg SubExp
msg (SrcLoc
loc, forall a. Monoid a => a
mempty)
asserting ::
InternaliseM Certs ->
InternaliseM Certs
asserting :: InternaliseM Certs -> InternaliseM Certs
asserting InternaliseM Certs
m = do
Bool
doBoundsChecks <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks InternaliseEnv -> Bool
envDoBoundsChecks
if Bool
doBoundsChecks
then InternaliseM Certs
m
else forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Monoid a => a
mempty
assertingOne ::
InternaliseM VName ->
InternaliseM Certs
assertingOne :: InternaliseM VName -> InternaliseM Certs
assertingOne InternaliseM VName
m = InternaliseM Certs -> InternaliseM Certs
asserting forall a b. (a -> b) -> a -> b
$ [VName] -> Certs
Certs forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a. Applicative f => a -> f a
pure forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> InternaliseM VName
m