{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE TypeFamilies #-}

module Futhark.Internalise.Monad
  ( InternaliseM,
    runInternaliseM,
    throwError,
    VarSubsts,
    InternaliseEnv (..),
    FunInfo,
    substitutingVars,
    lookupSubst,
    addFunDef,
    lookupFunction,
    lookupFunction',
    lookupConst,
    bindFunction,
    bindConstant,
    localConstsScope,
    assert,

    -- * Convenient reexports
    module Futhark.Tools,
  )
where

import Control.Monad.Except
import Control.Monad.Reader
import Control.Monad.State
import qualified Data.Map.Strict as M
import Futhark.IR.SOACS
import Futhark.MonadFreshNames
import Futhark.Tools

type FunInfo =
  ( [VName],
    [DeclType],
    [FParam],
    [(SubExp, Type)] -> Maybe [DeclExtType]
  )

type FunTable = M.Map VName FunInfo

-- | A mapping from external variable names to the corresponding
-- internalised subexpressions.
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 -> Scope SOACS
stateConstScope :: Scope SOACS,
    InternaliseState -> [FunDef SOACS]
stateFuns :: [FunDef SOACS]
  }

newtype InternaliseM a
  = InternaliseM
      (BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a)
  deriving
    ( a -> InternaliseM b -> InternaliseM a
(a -> b) -> InternaliseM a -> InternaliseM b
(forall a b. (a -> b) -> InternaliseM a -> InternaliseM b)
-> (forall a b. a -> InternaliseM b -> InternaliseM a)
-> Functor InternaliseM
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
<$ :: a -> InternaliseM b -> InternaliseM a
$c<$ :: forall a b. a -> InternaliseM b -> InternaliseM a
fmap :: (a -> b) -> InternaliseM a -> InternaliseM b
$cfmap :: forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
Functor,
      Functor InternaliseM
a -> InternaliseM a
Functor InternaliseM
-> (forall a. a -> InternaliseM a)
-> (forall a b.
    InternaliseM (a -> b) -> InternaliseM a -> InternaliseM b)
-> (forall a b c.
    (a -> b -> c)
    -> InternaliseM a -> InternaliseM b -> InternaliseM c)
-> (forall a b. InternaliseM a -> InternaliseM b -> InternaliseM b)
-> (forall a b. InternaliseM a -> InternaliseM b -> InternaliseM a)
-> Applicative InternaliseM
InternaliseM a -> InternaliseM b -> InternaliseM b
InternaliseM a -> InternaliseM b -> InternaliseM a
InternaliseM (a -> b) -> InternaliseM a -> InternaliseM b
(a -> b -> c) -> InternaliseM a -> InternaliseM b -> InternaliseM c
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
<* :: InternaliseM a -> InternaliseM b -> InternaliseM a
$c<* :: forall a b. InternaliseM a -> InternaliseM b -> InternaliseM a
*> :: InternaliseM a -> InternaliseM b -> InternaliseM b
$c*> :: forall a b. InternaliseM a -> InternaliseM b -> InternaliseM b
liftA2 :: (a -> b -> c) -> InternaliseM a -> InternaliseM b -> InternaliseM c
$cliftA2 :: forall a b c.
(a -> b -> c) -> InternaliseM a -> InternaliseM b -> InternaliseM c
<*> :: InternaliseM (a -> b) -> InternaliseM a -> InternaliseM b
$c<*> :: forall a b.
InternaliseM (a -> b) -> InternaliseM a -> InternaliseM b
pure :: a -> InternaliseM a
$cpure :: forall a. a -> InternaliseM a
$cp1Applicative :: Functor InternaliseM
Applicative,
      Applicative InternaliseM
a -> InternaliseM a
Applicative InternaliseM
-> (forall a b.
    InternaliseM a -> (a -> InternaliseM b) -> InternaliseM b)
-> (forall a b. InternaliseM a -> InternaliseM b -> InternaliseM b)
-> (forall a. a -> InternaliseM a)
-> Monad InternaliseM
InternaliseM a -> (a -> InternaliseM b) -> InternaliseM b
InternaliseM a -> InternaliseM b -> InternaliseM b
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 :: a -> InternaliseM a
$creturn :: forall a. a -> InternaliseM a
>> :: InternaliseM a -> InternaliseM b -> InternaliseM b
$c>> :: forall a b. InternaliseM a -> InternaliseM b -> InternaliseM b
>>= :: InternaliseM a -> (a -> InternaliseM b) -> InternaliseM b
$c>>= :: forall a b.
InternaliseM a -> (a -> InternaliseM b) -> InternaliseM b
$cp1Monad :: Applicative InternaliseM
Monad,
      MonadReader InternaliseEnv,
      MonadState InternaliseState,
      Monad InternaliseM
Applicative InternaliseM
InternaliseM VNameSource
Applicative InternaliseM
-> Monad InternaliseM
-> InternaliseM VNameSource
-> (VNameSource -> InternaliseM ())
-> MonadFreshNames InternaliseM
VNameSource -> InternaliseM ()
forall (m :: * -> *).
Applicative m
-> Monad m
-> m VNameSource
-> (VNameSource -> m ())
-> MonadFreshNames m
putNameSource :: VNameSource -> InternaliseM ()
$cputNameSource :: VNameSource -> InternaliseM ()
getNameSource :: InternaliseM VNameSource
$cgetNameSource :: InternaliseM VNameSource
$cp2MonadFreshNames :: Monad InternaliseM
$cp1MonadFreshNames :: Applicative InternaliseM
MonadFreshNames,
      HasScope SOACS,
      LocalScope SOACS
    )

instance MonadFreshNames (State InternaliseState) where
  getNameSource :: State InternaliseState VNameSource
getNameSource = (InternaliseState -> VNameSource)
-> State InternaliseState VNameSource
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets InternaliseState -> VNameSource
stateNameSource
  putNameSource :: VNameSource -> State InternaliseState ()
putNameSource VNameSource
src = (InternaliseState -> InternaliseState) -> State InternaliseState ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((InternaliseState -> InternaliseState)
 -> State InternaliseState ())
-> (InternaliseState -> InternaliseState)
-> State InternaliseState ()
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 (Rep InternaliseM)
-> Exp (Rep InternaliseM)
-> InternaliseM (ExpDec (Rep InternaliseM))
mkExpDecM Pat (Rep InternaliseM)
pat Exp (Rep InternaliseM)
e = BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) ()
-> InternaliseM ()
forall a.
BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
-> InternaliseM a
InternaliseM (BuilderT
   SOACS (ReaderT InternaliseEnv (State InternaliseState)) ()
 -> InternaliseM ())
-> BuilderT
     SOACS (ReaderT InternaliseEnv (State InternaliseState)) ()
-> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ Pat
  (Rep
     (BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState))))
-> Exp
     (Rep
        (BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState))))
-> BuilderT
     SOACS
     (ReaderT InternaliseEnv (State InternaliseState))
     (ExpDec
        (Rep
           (BuilderT
              SOACS (ReaderT InternaliseEnv (State InternaliseState)))))
forall (m :: * -> *).
MonadBuilder m =>
Pat (Rep m) -> Exp (Rep m) -> m (ExpDec (Rep m))
mkExpDecM Pat
  (Rep
     (BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState))))
Pat (Rep InternaliseM)
pat Exp
  (Rep
     (BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState))))
Exp (Rep InternaliseM)
e
  mkBodyM :: Stms (Rep InternaliseM)
-> Result -> InternaliseM (Body (Rep InternaliseM))
mkBodyM Stms (Rep InternaliseM)
stms Result
res = BuilderT
  SOACS
  (ReaderT InternaliseEnv (State InternaliseState))
  (Body SOACS)
-> InternaliseM (Body SOACS)
forall a.
BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
-> InternaliseM a
InternaliseM (BuilderT
   SOACS
   (ReaderT InternaliseEnv (State InternaliseState))
   (Body SOACS)
 -> InternaliseM (Body SOACS))
-> BuilderT
     SOACS
     (ReaderT InternaliseEnv (State InternaliseState))
     (Body SOACS)
-> InternaliseM (Body SOACS)
forall a b. (a -> b) -> a -> b
$ Stms
  (Rep
     (BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState))))
-> Result
-> BuilderT
     SOACS
     (ReaderT InternaliseEnv (State InternaliseState))
     (Body
        (Rep
           (BuilderT
              SOACS (ReaderT InternaliseEnv (State InternaliseState)))))
forall (m :: * -> *).
MonadBuilder m =>
Stms (Rep m) -> Result -> m (Body (Rep m))
mkBodyM Stms
  (Rep
     (BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState))))
Stms (Rep InternaliseM)
stms Result
res
  mkLetNamesM :: [VName]
-> Exp (Rep InternaliseM) -> InternaliseM (Stm (Rep InternaliseM))
mkLetNamesM [VName]
pat Exp (Rep InternaliseM)
e = BuilderT
  SOACS (ReaderT InternaliseEnv (State InternaliseState)) (Stm SOACS)
-> InternaliseM (Stm SOACS)
forall a.
BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
-> InternaliseM a
InternaliseM (BuilderT
   SOACS (ReaderT InternaliseEnv (State InternaliseState)) (Stm SOACS)
 -> InternaliseM (Stm SOACS))
-> BuilderT
     SOACS (ReaderT InternaliseEnv (State InternaliseState)) (Stm SOACS)
-> InternaliseM (Stm SOACS)
forall a b. (a -> b) -> a -> b
$ [VName]
-> Exp
     (Rep
        (BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState))))
-> BuilderT
     SOACS
     (ReaderT InternaliseEnv (State InternaliseState))
     (Stm
        (Rep
           (BuilderT
              SOACS (ReaderT InternaliseEnv (State InternaliseState)))))
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m (Stm (Rep m))
mkLetNamesM [VName]
pat Exp
  (Rep
     (BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState))))
Exp (Rep InternaliseM)
e

  addStms :: Stms (Rep InternaliseM) -> InternaliseM ()
addStms = BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) ()
-> InternaliseM ()
forall a.
BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
-> InternaliseM a
InternaliseM (BuilderT
   SOACS (ReaderT InternaliseEnv (State InternaliseState)) ()
 -> InternaliseM ())
-> (Stms SOACS
    -> BuilderT
         SOACS (ReaderT InternaliseEnv (State InternaliseState)) ())
-> Stms SOACS
-> InternaliseM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stms SOACS
-> BuilderT
     SOACS (ReaderT InternaliseEnv (State InternaliseState)) ()
forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
addStms
  collectStms :: InternaliseM a -> InternaliseM (a, Stms (Rep InternaliseM))
collectStms (InternaliseM BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
m) = BuilderT
  SOACS
  (ReaderT InternaliseEnv (State InternaliseState))
  (a, Stms SOACS)
-> InternaliseM (a, Stms SOACS)
forall a.
BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
-> InternaliseM a
InternaliseM (BuilderT
   SOACS
   (ReaderT InternaliseEnv (State InternaliseState))
   (a, Stms SOACS)
 -> InternaliseM (a, Stms SOACS))
-> BuilderT
     SOACS
     (ReaderT InternaliseEnv (State InternaliseState))
     (a, Stms SOACS)
-> InternaliseM (a, Stms SOACS)
forall a b. (a -> b) -> a -> b
$ BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
-> BuilderT
     SOACS
     (ReaderT InternaliseEnv (State InternaliseState))
     (a,
      Stms
        (Rep
           (BuilderT
              SOACS (ReaderT InternaliseEnv (State InternaliseState)))))
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 (Stms SOACS, [FunDef SOACS])
runInternaliseM :: Bool -> InternaliseM () -> m (Stms SOACS, [FunDef SOACS])
runInternaliseM Bool
safe (InternaliseM BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) ()
m) =
  (VNameSource -> ((Stms SOACS, [FunDef SOACS]), VNameSource))
-> m (Stms SOACS, [FunDef SOACS])
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource ((VNameSource -> ((Stms SOACS, [FunDef SOACS]), VNameSource))
 -> m (Stms SOACS, [FunDef SOACS]))
-> (VNameSource -> ((Stms SOACS, [FunDef SOACS]), VNameSource))
-> m (Stms SOACS, [FunDef SOACS])
forall a b. (a -> b) -> a -> b
$ \VNameSource
src ->
    let ((()
_, Stms SOACS
consts), InternaliseState
s) =
          State InternaliseState ((), Stms SOACS)
-> InternaliseState -> (((), Stms SOACS), InternaliseState)
forall s a. State s a -> s -> (a, s)
runState (ReaderT InternaliseEnv (State InternaliseState) ((), Stms SOACS)
-> InternaliseEnv -> State InternaliseState ((), Stms SOACS)
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT (BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) ()
-> Scope SOACS
-> ReaderT InternaliseEnv (State InternaliseState) ((), Stms SOACS)
forall (m :: * -> *) rep a.
MonadFreshNames m =>
BuilderT rep m a -> Scope rep -> m (a, Stms rep)
runBuilderT BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) ()
m Scope SOACS
forall a. Monoid a => a
mempty) InternaliseEnv
newEnv) (VNameSource -> InternaliseState
newState VNameSource
src)
     in ((Stms SOACS
consts, [FunDef SOACS] -> [FunDef SOACS]
forall a. [a] -> [a]
reverse ([FunDef SOACS] -> [FunDef SOACS])
-> [FunDef SOACS] -> [FunDef SOACS]
forall a b. (a -> b) -> a -> b
$ InternaliseState -> [FunDef SOACS]
stateFuns InternaliseState
s), InternaliseState -> VNameSource
stateNameSource InternaliseState
s)
  where
    newEnv :: InternaliseEnv
newEnv =
      InternaliseEnv :: VarSubsts -> Bool -> Bool -> Attrs -> InternaliseEnv
InternaliseEnv
        { envSubsts :: VarSubsts
envSubsts = VarSubsts
forall a. Monoid a => a
mempty,
          envDoBoundsChecks :: Bool
envDoBoundsChecks = Bool
True,
          envSafe :: Bool
envSafe = Bool
safe,
          envAttrs :: Attrs
envAttrs = Attrs
forall a. Monoid a => a
mempty
        }
    newState :: VNameSource -> InternaliseState
newState VNameSource
src =
      InternaliseState :: VNameSource
-> FunTable
-> VarSubsts
-> Scope SOACS
-> [FunDef SOACS]
-> InternaliseState
InternaliseState
        { stateNameSource :: VNameSource
stateNameSource = VNameSource
src,
          stateFunTable :: FunTable
stateFunTable = FunTable
forall a. Monoid a => a
mempty,
          stateConstSubsts :: VarSubsts
stateConstSubsts = VarSubsts
forall a. Monoid a => a
mempty,
          stateConstScope :: Scope SOACS
stateConstScope = Scope SOACS
forall a. Monoid a => a
mempty,
          stateFuns :: [FunDef SOACS]
stateFuns = [FunDef SOACS]
forall a. Monoid a => a
mempty
        }

substitutingVars :: VarSubsts -> InternaliseM a -> InternaliseM a
substitutingVars :: VarSubsts -> InternaliseM a -> InternaliseM a
substitutingVars VarSubsts
substs = (InternaliseEnv -> InternaliseEnv)
-> InternaliseM a -> InternaliseM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((InternaliseEnv -> InternaliseEnv)
 -> InternaliseM a -> InternaliseM a)
-> (InternaliseEnv -> InternaliseEnv)
-> InternaliseM a
-> InternaliseM a
forall a b. (a -> b) -> a -> b
$ \InternaliseEnv
env -> InternaliseEnv
env {envSubsts :: VarSubsts
envSubsts = VarSubsts
substs VarSubsts -> VarSubsts -> VarSubsts
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 <- (InternaliseEnv -> Maybe [SubExp]) -> InternaliseM (Maybe [SubExp])
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((InternaliseEnv -> Maybe [SubExp])
 -> InternaliseM (Maybe [SubExp]))
-> (InternaliseEnv -> Maybe [SubExp])
-> InternaliseM (Maybe [SubExp])
forall a b. (a -> b) -> a -> b
$ VName -> VarSubsts -> Maybe [SubExp]
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v (VarSubsts -> Maybe [SubExp])
-> (InternaliseEnv -> VarSubsts)
-> InternaliseEnv
-> Maybe [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InternaliseEnv -> VarSubsts
envSubsts
  Maybe [SubExp]
const_substs <- (InternaliseState -> Maybe [SubExp])
-> InternaliseM (Maybe [SubExp])
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((InternaliseState -> Maybe [SubExp])
 -> InternaliseM (Maybe [SubExp]))
-> (InternaliseState -> Maybe [SubExp])
-> InternaliseM (Maybe [SubExp])
forall a b. (a -> b) -> a -> b
$ VName -> VarSubsts -> Maybe [SubExp]
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v (VarSubsts -> Maybe [SubExp])
-> (InternaliseState -> VarSubsts)
-> InternaliseState
-> Maybe [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InternaliseState -> VarSubsts
stateConstSubsts
  Maybe [SubExp] -> InternaliseM (Maybe [SubExp])
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe [SubExp] -> InternaliseM (Maybe [SubExp]))
-> Maybe [SubExp] -> InternaliseM (Maybe [SubExp])
forall a b. (a -> b) -> a -> b
$ Maybe [SubExp]
env_substs Maybe [SubExp] -> Maybe [SubExp] -> Maybe [SubExp]
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` Maybe [SubExp]
const_substs

-- | Add a function definition to the program being constructed.
addFunDef :: FunDef SOACS -> InternaliseM ()
addFunDef :: FunDef SOACS -> InternaliseM ()
addFunDef FunDef SOACS
fd = (InternaliseState -> InternaliseState) -> InternaliseM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((InternaliseState -> InternaliseState) -> InternaliseM ())
-> (InternaliseState -> InternaliseState) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \InternaliseState
s -> InternaliseState
s {stateFuns :: [FunDef SOACS]
stateFuns = FunDef SOACS
fd FunDef SOACS -> [FunDef SOACS] -> [FunDef SOACS]
forall a. a -> [a] -> [a]
: InternaliseState -> [FunDef SOACS]
stateFuns InternaliseState
s}

lookupFunction' :: VName -> InternaliseM (Maybe FunInfo)
lookupFunction' :: VName -> InternaliseM (Maybe FunInfo)
lookupFunction' VName
fname = (InternaliseState
 -> Maybe
      ([VName], [DeclType], [Param DeclType],
       [(SubExp, Type)] -> Maybe [DeclExtType]))
-> InternaliseM
     (Maybe
        ([VName], [DeclType], [Param DeclType],
         [(SubExp, Type)] -> Maybe [DeclExtType]))
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((InternaliseState
  -> Maybe
       ([VName], [DeclType], [Param DeclType],
        [(SubExp, Type)] -> Maybe [DeclExtType]))
 -> InternaliseM
      (Maybe
         ([VName], [DeclType], [Param DeclType],
          [(SubExp, Type)] -> Maybe [DeclExtType])))
-> (InternaliseState
    -> Maybe
         ([VName], [DeclType], [Param DeclType],
          [(SubExp, Type)] -> Maybe [DeclExtType]))
-> InternaliseM
     (Maybe
        ([VName], [DeclType], [Param DeclType],
         [(SubExp, Type)] -> Maybe [DeclExtType]))
forall a b. (a -> b) -> a -> b
$ VName
-> Map
     VName
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [DeclExtType])
-> Maybe
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [DeclExtType])
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
fname (Map
   VName
   ([VName], [DeclType], [Param DeclType],
    [(SubExp, Type)] -> Maybe [DeclExtType])
 -> Maybe
      ([VName], [DeclType], [Param DeclType],
       [(SubExp, Type)] -> Maybe [DeclExtType]))
-> (InternaliseState
    -> Map
         VName
         ([VName], [DeclType], [Param DeclType],
          [(SubExp, Type)] -> Maybe [DeclExtType]))
-> InternaliseState
-> Maybe
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [DeclExtType])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InternaliseState
-> Map
     VName
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [DeclExtType])
InternaliseState -> FunTable
stateFunTable

lookupFunction :: VName -> InternaliseM FunInfo
lookupFunction :: VName -> InternaliseM FunInfo
lookupFunction VName
fname = InternaliseM
  ([VName], [DeclType], [Param DeclType],
   [(SubExp, Type)] -> Maybe [DeclExtType])
-> (([VName], [DeclType], [Param DeclType],
     [(SubExp, Type)] -> Maybe [DeclExtType])
    -> InternaliseM
         ([VName], [DeclType], [Param DeclType],
          [(SubExp, Type)] -> Maybe [DeclExtType]))
-> Maybe
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [DeclExtType])
-> InternaliseM
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [DeclExtType])
forall b a. b -> (a -> b) -> Maybe a -> b
maybe InternaliseM
  ([VName], [DeclType], [Param DeclType],
   [(SubExp, Type)] -> Maybe [DeclExtType])
bad ([VName], [DeclType], [Param DeclType],
 [(SubExp, Type)] -> Maybe [DeclExtType])
-> InternaliseM
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [DeclExtType])
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe
   ([VName], [DeclType], [Param DeclType],
    [(SubExp, Type)] -> Maybe [DeclExtType])
 -> InternaliseM
      ([VName], [DeclType], [Param DeclType],
       [(SubExp, Type)] -> Maybe [DeclExtType]))
-> InternaliseM
     (Maybe
        ([VName], [DeclType], [Param DeclType],
         [(SubExp, Type)] -> Maybe [DeclExtType]))
-> InternaliseM
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [DeclExtType])
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 = [Char]
-> InternaliseM
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [DeclExtType])
forall a. HasCallStack => [Char] -> a
error ([Char]
 -> InternaliseM
      ([VName], [DeclType], [Param DeclType],
       [(SubExp, Type)] -> Maybe [DeclExtType]))
-> [Char]
-> InternaliseM
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [DeclExtType])
forall a b. (a -> b) -> a -> b
$ [Char]
"Internalise.lookupFunction: Function '" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ VName -> [Char]
forall a. Pretty a => a -> [Char]
pretty VName
fname [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"' not found."

lookupConst :: VName -> InternaliseM (Maybe [SubExp])
lookupConst :: VName -> InternaliseM (Maybe [SubExp])
lookupConst VName
fname = (InternaliseState -> Maybe [SubExp])
-> InternaliseM (Maybe [SubExp])
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((InternaliseState -> Maybe [SubExp])
 -> InternaliseM (Maybe [SubExp]))
-> (InternaliseState -> Maybe [SubExp])
-> InternaliseM (Maybe [SubExp])
forall a b. (a -> b) -> a -> b
$ VName -> VarSubsts -> Maybe [SubExp]
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
fname (VarSubsts -> Maybe [SubExp])
-> (InternaliseState -> VarSubsts)
-> InternaliseState
-> Maybe [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InternaliseState -> VarSubsts
stateConstSubsts

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
  (InternaliseState -> InternaliseState) -> InternaliseM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((InternaliseState -> InternaliseState) -> InternaliseM ())
-> (InternaliseState -> InternaliseState) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \InternaliseState
s -> InternaliseState
s {stateFunTable :: FunTable
stateFunTable = VName
-> ([VName], [DeclType], [Param DeclType],
    [(SubExp, Type)] -> Maybe [DeclExtType])
-> Map
     VName
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [DeclExtType])
-> Map
     VName
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [DeclExtType])
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
fname ([VName], [DeclType], [Param DeclType],
 [(SubExp, Type)] -> Maybe [DeclExtType])
FunInfo
info (Map
   VName
   ([VName], [DeclType], [Param DeclType],
    [(SubExp, Type)] -> Maybe [DeclExtType])
 -> Map
      VName
      ([VName], [DeclType], [Param DeclType],
       [(SubExp, Type)] -> Maybe [DeclExtType]))
-> Map
     VName
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [DeclExtType])
-> Map
     VName
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [DeclExtType])
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
  let stms :: Stms SOACS
stms = Body SOACS -> Stms SOACS
forall rep. BodyT rep -> Stms rep
bodyStms (Body SOACS -> Stms SOACS) -> Body SOACS -> Stms SOACS
forall a b. (a -> b) -> a -> b
$ FunDef SOACS -> Body SOACS
forall rep. FunDef rep -> BodyT rep
funDefBody FunDef SOACS
fd
      substs :: [SubExp]
substs =
        Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
drop (Set Int -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([DeclExtType] -> Set Int
forall u. [TypeBase ExtShape u] -> Set Int
shapeContext (FunDef SOACS -> [RetType SOACS]
forall rep. FunDef rep -> [RetType rep]
funDefRetType FunDef SOACS
fd))) ([SubExp] -> [SubExp]) -> [SubExp] -> [SubExp]
forall a b. (a -> b) -> a -> b
$
          (SubExpRes -> SubExp) -> Result -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp (Result -> [SubExp]) -> Result -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Body SOACS -> Result
forall rep. BodyT rep -> Result
bodyResult (Body SOACS -> Result) -> Body SOACS -> Result
forall a b. (a -> b) -> a -> b
$ FunDef SOACS -> Body SOACS
forall rep. FunDef rep -> BodyT rep
funDefBody FunDef SOACS
fd
  Stms (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
addStms Stms (Rep InternaliseM)
Stms SOACS
stms
  (InternaliseState -> InternaliseState) -> InternaliseM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((InternaliseState -> InternaliseState) -> InternaliseM ())
-> (InternaliseState -> InternaliseState) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \InternaliseState
s ->
    InternaliseState
s
      { stateConstSubsts :: VarSubsts
stateConstSubsts = VName -> [SubExp] -> VarSubsts -> VarSubsts
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
cname [SubExp]
substs (VarSubsts -> VarSubsts) -> VarSubsts -> VarSubsts
forall a b. (a -> b) -> a -> b
$ InternaliseState -> VarSubsts
stateConstSubsts InternaliseState
s,
        stateConstScope :: Scope SOACS
stateConstScope = Stms SOACS -> Scope SOACS
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stms SOACS
stms Scope SOACS -> Scope SOACS -> Scope SOACS
forall a. Semigroup a => a -> a -> a
<> InternaliseState -> Scope SOACS
stateConstScope InternaliseState
s
      }

localConstsScope :: InternaliseM a -> InternaliseM a
localConstsScope :: InternaliseM a -> InternaliseM a
localConstsScope InternaliseM a
m = do
  Scope SOACS
scope <- (InternaliseState -> Scope SOACS) -> InternaliseM (Scope SOACS)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets InternaliseState -> Scope SOACS
stateConstScope
  Scope SOACS -> InternaliseM a -> InternaliseM a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope Scope SOACS
scope InternaliseM a
m

-- | Construct an 'Assert' statement, but taking attributes into
-- account.  Always use this function, and never construct 'Assert'
-- directly in the internaliser!
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 (InternaliseM VName -> InternaliseM Certs)
-> InternaliseM VName -> InternaliseM Certs
forall a b. (a -> b) -> a -> b
$ do
  Attrs
attrs <- (InternaliseEnv -> Attrs) -> InternaliseM Attrs
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((InternaliseEnv -> Attrs) -> InternaliseM Attrs)
-> (InternaliseEnv -> Attrs) -> InternaliseM Attrs
forall a b. (a -> b) -> a -> b
$ Attrs -> Attrs
attrsForAssert (Attrs -> Attrs)
-> (InternaliseEnv -> Attrs) -> InternaliseEnv -> Attrs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InternaliseEnv -> Attrs
envAttrs
  Attrs -> InternaliseM VName -> InternaliseM VName
forall (m :: * -> *) a. MonadBuilder m => Attrs -> m a -> m a
attributing Attrs
attrs (InternaliseM VName -> InternaliseM VName)
-> InternaliseM VName -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$
    [Char] -> Exp (Rep InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
desc (Exp (Rep InternaliseM) -> InternaliseM VName)
-> Exp (Rep InternaliseM) -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$
      BasicOp -> ExpT SOACS
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> ErrorMsg SubExp -> (SrcLoc, [SrcLoc]) -> BasicOp
Assert SubExp
se ErrorMsg SubExp
msg (SrcLoc
loc, [SrcLoc]
forall a. Monoid a => a
mempty)

-- | Execute the given action if 'envDoBoundsChecks' is true, otherwise
-- just return an empty list.
asserting ::
  InternaliseM Certs ->
  InternaliseM Certs
asserting :: InternaliseM Certs -> InternaliseM Certs
asserting InternaliseM Certs
m = do
  Bool
doBoundsChecks <- (InternaliseEnv -> Bool) -> InternaliseM Bool
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks InternaliseEnv -> Bool
envDoBoundsChecks
  if Bool
doBoundsChecks
    then InternaliseM Certs
m
    else Certs -> InternaliseM Certs
forall (m :: * -> *) a. Monad m => a -> m a
return Certs
forall a. Monoid a => a
mempty

-- | Execute the given action if 'envDoBoundsChecks' is true, otherwise
-- just return an empty list.
assertingOne ::
  InternaliseM VName ->
  InternaliseM Certs
assertingOne :: InternaliseM VName -> InternaliseM Certs
assertingOne InternaliseM VName
m = InternaliseM Certs -> InternaliseM Certs
asserting (InternaliseM Certs -> InternaliseM Certs)
-> InternaliseM Certs -> InternaliseM Certs
forall a b. (a -> b) -> a -> b
$ [VName] -> Certs
Certs ([VName] -> Certs) -> (VName -> [VName]) -> VName -> Certs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> [VName]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Certs) -> InternaliseM VName -> InternaliseM Certs
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> InternaliseM VName
m