{-# LANGUAGE TypeFamilies #-}

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

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

import Control.Monad
import Control.Monad.Except
import Control.Monad.Reader
import Control.Monad.State
import Data.List (find)
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, RetAls)]
  )

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 -> [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 -> 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
$cfmap :: forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
fmap :: forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
$c<$ :: forall a b. a -> InternaliseM b -> InternaliseM a
<$ :: forall a b. a -> InternaliseM b -> InternaliseM a
Functor,
      Functor InternaliseM
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
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
$cpure :: forall a. a -> InternaliseM a
pure :: forall a. a -> InternaliseM a
$c<*> :: forall a b.
InternaliseM (a -> b) -> InternaliseM a -> InternaliseM b
<*> :: forall a b.
InternaliseM (a -> b) -> InternaliseM a -> InternaliseM b
$cliftA2 :: forall a b c.
(a -> b -> c) -> InternaliseM a -> InternaliseM b -> InternaliseM c
liftA2 :: forall a b c.
(a -> b -> c) -> InternaliseM a -> InternaliseM b -> InternaliseM c
$c*> :: forall a b. InternaliseM a -> InternaliseM b -> InternaliseM b
*> :: forall a b. InternaliseM a -> InternaliseM b -> InternaliseM b
$c<* :: forall a b. InternaliseM a -> InternaliseM b -> InternaliseM a
<* :: forall a b. InternaliseM a -> InternaliseM b -> InternaliseM a
Applicative,
      Applicative InternaliseM
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
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
$c>>= :: forall a b.
InternaliseM a -> (a -> InternaliseM b) -> InternaliseM b
>>= :: forall a b.
InternaliseM a -> (a -> InternaliseM b) -> InternaliseM b
$c>> :: forall a b. InternaliseM a -> InternaliseM b -> InternaliseM b
>> :: forall a b. InternaliseM a -> InternaliseM b -> InternaliseM b
$creturn :: forall a. a -> InternaliseM a
return :: forall a. a -> InternaliseM a
Monad,
      MonadReader InternaliseEnv,
      MonadState InternaliseState,
      Monad InternaliseM
InternaliseM VNameSource
Monad InternaliseM =>
InternaliseM VNameSource
-> (VNameSource -> InternaliseM ()) -> MonadFreshNames InternaliseM
VNameSource -> InternaliseM ()
forall (m :: * -> *).
Monad m =>
m VNameSource -> (VNameSource -> m ()) -> MonadFreshNames m
$cgetNameSource :: InternaliseM VNameSource
getNameSource :: InternaliseM VNameSource
$cputNameSource :: VNameSource -> InternaliseM ()
putNameSource :: VNameSource -> InternaliseM ()
MonadFreshNames,
      HasScope SOACS
    )

-- Internalisation has to deal with the risk of multiple binding of
-- the same variable (although always of the same type) in the
-- program; in particular this might imply shadowing a constant.  The
-- LocalScope instance for BuilderT does not handle this properly (and
-- doing so would make it slower).  So we remove already-known
-- variables before passing the scope on.
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 <- InternaliseM (Scope SOACS)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
    BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
-> InternaliseM a
forall a.
BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
-> InternaliseM a
InternaliseM (BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
 -> InternaliseM a)
-> BuilderT
     SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
-> InternaliseM a
forall a b. (a -> b) -> a -> b
$ Scope SOACS
-> BuilderT
     SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
-> BuilderT
     SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
forall a.
Scope SOACS
-> BuilderT
     SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
-> BuilderT
     SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (Scope SOACS
scope Scope SOACS -> Scope SOACS -> Scope SOACS
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 = (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 = 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 = BuilderT
  SOACS
  (ReaderT InternaliseEnv (State InternaliseState))
  (ExpDec (Rep InternaliseM))
-> InternaliseM (ExpDec (Rep InternaliseM))
forall a.
BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
-> InternaliseM a
InternaliseM (BuilderT
   SOACS
   (ReaderT InternaliseEnv (State InternaliseState))
   (ExpDec (Rep InternaliseM))
 -> InternaliseM (ExpDec (Rep InternaliseM)))
-> BuilderT
     SOACS
     (ReaderT InternaliseEnv (State InternaliseState))
     (ExpDec (Rep InternaliseM))
-> InternaliseM (ExpDec (Rep InternaliseM))
forall a b. (a -> b) -> a -> b
$ Pat
  (LetDec
     (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 (LetDec (Rep m)) -> Exp (Rep m) -> m (ExpDec (Rep m))
mkExpDecM Pat
  (LetDec
     (Rep
        (BuilderT
           SOACS (ReaderT InternaliseEnv (State InternaliseState)))))
Pat (LetDec (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 (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
forall a.
BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
-> InternaliseM a
InternaliseM (BuilderT
   SOACS
   (ReaderT InternaliseEnv (State InternaliseState))
   (Body (Rep InternaliseM))
 -> InternaliseM (Body (Rep InternaliseM)))
-> BuilderT
     SOACS
     (ReaderT InternaliseEnv (State InternaliseState))
     (Body (Rep InternaliseM))
-> InternaliseM (Body (Rep InternaliseM))
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 (Rep InternaliseM))
-> InternaliseM (Stm (Rep InternaliseM))
forall a.
BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
-> InternaliseM a
InternaliseM (BuilderT
   SOACS
   (ReaderT InternaliseEnv (State InternaliseState))
   (Stm (Rep InternaliseM))
 -> InternaliseM (Stm (Rep InternaliseM)))
-> BuilderT
     SOACS
     (ReaderT InternaliseEnv (State InternaliseState))
     (Stm (Rep InternaliseM))
-> InternaliseM (Stm (Rep InternaliseM))
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
  (Rep
     (BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState))))
-> BuilderT
     SOACS (ReaderT InternaliseEnv (State InternaliseState)) ()
Stms SOACS
-> BuilderT
     SOACS (ReaderT InternaliseEnv (State InternaliseState)) ()
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) = BuilderT
  SOACS
  (ReaderT InternaliseEnv (State InternaliseState))
  (a, Stms (Rep InternaliseM))
-> InternaliseM (a, Stms (Rep InternaliseM))
forall a.
BuilderT SOACS (ReaderT InternaliseEnv (State InternaliseState)) a
-> InternaliseM a
InternaliseM (BuilderT
   SOACS
   (ReaderT InternaliseEnv (State InternaliseState))
   (a, Stms (Rep InternaliseM))
 -> InternaliseM (a, Stms (Rep InternaliseM)))
-> BuilderT
     SOACS
     (ReaderT InternaliseEnv (State InternaliseState))
     (a, Stms (Rep InternaliseM))
-> InternaliseM (a, Stms (Rep InternaliseM))
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 a.
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 (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) =
  (VNameSource
 -> ((OpaqueTypes, Stms SOACS, [FunDef SOACS]), VNameSource))
-> m (OpaqueTypes, Stms SOACS, [FunDef SOACS])
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource ((VNameSource
  -> ((OpaqueTypes, Stms SOACS, [FunDef SOACS]), VNameSource))
 -> m (OpaqueTypes, Stms SOACS, [FunDef SOACS]))
-> (VNameSource
    -> ((OpaqueTypes, Stms SOACS, [FunDef SOACS]), VNameSource))
-> m (OpaqueTypes, 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 ( (InternaliseState -> OpaqueTypes
stateTypes InternaliseState
s, 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
        { 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
        { stateNameSource :: VNameSource
stateNameSource = VNameSource
src,
          stateFunTable :: FunTable
stateFunTable = Map
  VName
  ([VName], [DeclType], [Param DeclType],
   [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)])
FunTable
forall a. Monoid a => a
mempty,
          stateConstSubsts :: VarSubsts
stateConstSubsts = VarSubsts
forall a. Monoid a => a
mempty,
          stateFuns :: [FunDef SOACS]
stateFuns = [FunDef SOACS]
forall a. Monoid a => a
mempty,
          stateTypes :: OpaqueTypes
stateTypes = OpaqueTypes
forall a. Monoid a => a
mempty
        }

substitutingVars :: VarSubsts -> InternaliseM a -> InternaliseM a
substitutingVars :: forall a. VarSubsts -> InternaliseM a -> InternaliseM a
substitutingVars VarSubsts
substs = (InternaliseEnv -> InternaliseEnv)
-> InternaliseM a -> InternaliseM a
forall a.
(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 = substs <> envSubsts 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 a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (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 a. Maybe a -> Maybe a -> Maybe a
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` Maybe [SubExp]
const_substs

-- | Add opaque types.  If the types are already known, they will not
-- be added.
addOpaques :: OpaqueTypes -> InternaliseM ()
addOpaques :: OpaqueTypes -> InternaliseM ()
addOpaques ts :: OpaqueTypes
ts@(OpaqueTypes [(Name, OpaqueType)]
ts') = (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 ->
  -- TODO: handle this better (#1960)
  case ((Name, OpaqueType) -> Bool)
-> [(Name, OpaqueType)] -> Maybe (Name, OpaqueType)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (OpaqueTypes -> (Name, OpaqueType) -> Bool
knownButDifferent (InternaliseState -> OpaqueTypes
stateTypes InternaliseState
s)) [(Name, OpaqueType)]
ts' of
    Just (Name
x, OpaqueType
_) -> [Char] -> InternaliseState
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseState) -> [Char] -> InternaliseState
forall a b. (a -> b) -> a -> b
$ [Char]
"addOpaques: multiple incompatible definitions of type " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Name -> [Char]
nameToString Name
x
    Maybe (Name, OpaqueType)
Nothing -> InternaliseState
s {stateTypes = stateTypes s <> ts}
  where
    knownButDifferent :: OpaqueTypes -> (Name, OpaqueType) -> Bool
knownButDifferent (OpaqueTypes [(Name, OpaqueType)]
old_ts) (Name
v, OpaqueType
def) =
      ((Name, OpaqueType) -> Bool) -> [(Name, OpaqueType)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\(Name
v_old, OpaqueType
v_def) -> Name
v Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
v_old Bool -> Bool -> Bool
&& OpaqueType
def OpaqueType -> OpaqueType -> Bool
forall a. Eq a => a -> a -> Bool
/= OpaqueType
v_def) [(Name, OpaqueType)]
old_ts

-- | 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 = fd : stateFuns s}

lookupFunction :: VName -> InternaliseM FunInfo
lookupFunction :: VName -> InternaliseM FunInfo
lookupFunction VName
fname = InternaliseM
  ([VName], [DeclType], [Param DeclType],
   [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)])
-> (([VName], [DeclType], [Param DeclType],
     [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)])
    -> InternaliseM
         ([VName], [DeclType], [Param DeclType],
          [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)]))
-> Maybe
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)])
-> InternaliseM
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)])
forall b a. b -> (a -> b) -> Maybe a -> b
maybe InternaliseM
  ([VName], [DeclType], [Param DeclType],
   [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)])
bad ([VName], [DeclType], [Param DeclType],
 [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)])
-> InternaliseM
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)])
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe
   ([VName], [DeclType], [Param DeclType],
    [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)])
 -> InternaliseM
      ([VName], [DeclType], [Param DeclType],
       [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)]))
-> InternaliseM
     (Maybe
        ([VName], [DeclType], [Param DeclType],
         [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)]))
-> InternaliseM
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)])
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (InternaliseState
 -> Maybe
      ([VName], [DeclType], [Param DeclType],
       [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)]))
-> InternaliseM
     (Maybe
        ([VName], [DeclType], [Param DeclType],
         [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)]))
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (VName
-> Map
     VName
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)])
-> Maybe
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)])
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, RetAls)])
 -> Maybe
      ([VName], [DeclType], [Param DeclType],
       [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)]))
-> (InternaliseState
    -> Map
         VName
         ([VName], [DeclType], [Param DeclType],
          [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)]))
-> InternaliseState
-> Maybe
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InternaliseState
-> Map
     VName
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)])
InternaliseState -> FunTable
stateFunTable)
  where
    bad :: InternaliseM
  ([VName], [DeclType], [Param DeclType],
   [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)])
bad = [Char]
-> InternaliseM
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)])
forall a. HasCallStack => [Char] -> a
error ([Char]
 -> InternaliseM
      ([VName], [DeclType], [Param DeclType],
       [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)]))
-> [Char]
-> InternaliseM
     ([VName], [DeclType], [Param DeclType],
      [(SubExp, Type)] -> Maybe [(DeclExtType, RetAls)])
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]
prettyString 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 = do
  Bool
is_var <- (Scope SOACS -> Bool) -> InternaliseM Bool
forall a. (Scope SOACS -> a) -> InternaliseM a
forall rep (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) -> Maybe [SubExp] -> InternaliseM (Maybe [SubExp])
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe [SubExp] -> InternaliseM (Maybe [SubExp]))
-> Maybe [SubExp] -> InternaliseM (Maybe [SubExp])
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Maybe [SubExp]
forall a. a -> Maybe a
Just [SubExp]
ses
    (Bool
True, Maybe [SubExp]
_) -> Maybe [SubExp] -> InternaliseM (Maybe [SubExp])
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe [SubExp] -> InternaliseM (Maybe [SubExp]))
-> Maybe [SubExp] -> InternaliseM (Maybe [SubExp])
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Maybe [SubExp]
forall a. a -> Maybe a
Just [VName -> SubExp
Var VName
fname]
    (Bool, Maybe [SubExp])
_ -> Maybe [SubExp] -> InternaliseM (Maybe [SubExp])
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe [SubExp]
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
  (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 = M.insert fname info $ stateFunTable s}

bindConstant :: VName -> FunDef SOACS -> InternaliseM ()
bindConstant :: VName -> FunDef SOACS -> InternaliseM ()
bindConstant VName
cname FunDef SOACS
fd = do
  Stms (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
addStms (Stms (Rep InternaliseM) -> InternaliseM ())
-> Stms (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ Body (Rep InternaliseM) -> Stms (Rep InternaliseM)
forall rep. Body rep -> Stms rep
bodyStms (Body (Rep InternaliseM) -> Stms (Rep InternaliseM))
-> Body (Rep InternaliseM) -> Stms (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ FunDef SOACS -> Body SOACS
forall rep. FunDef rep -> Body rep
funDefBody FunDef SOACS
fd

  case (SubExpRes -> SubExp) -> Result -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp (Result -> [SubExp])
-> (FunDef SOACS -> Result) -> FunDef SOACS -> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Body SOACS -> Result
forall rep. Body rep -> Result
bodyResult (Body SOACS -> Result)
-> (FunDef SOACS -> Body SOACS) -> FunDef SOACS -> Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FunDef SOACS -> Body SOACS
forall rep. FunDef rep -> Body rep
funDefBody (FunDef SOACS -> [SubExp]) -> FunDef SOACS -> [SubExp]
forall a b. (a -> b) -> a -> b
$ FunDef SOACS
fd of
    [SubExp
se] -> do
      [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
cname] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
    [SubExp]
ses -> do
      let substs :: [SubExp]
substs =
            Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
drop (Set Int -> Int
forall a. Set a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([DeclExtType] -> Set Int
forall u. [TypeBase ExtShape u] -> Set Int
shapeContext (((DeclExtType, RetAls) -> DeclExtType)
-> [(DeclExtType, RetAls)] -> [DeclExtType]
forall a b. (a -> b) -> [a] -> [b]
map (DeclExtType, RetAls) -> DeclExtType
forall a b. (a, b) -> a
fst (FunDef SOACS -> [(RetType SOACS, RetAls)]
forall rep. FunDef rep -> [(RetType rep, RetAls)]
funDefRetType FunDef SOACS
fd)))) [SubExp]
ses
      (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 = M.insert cname substs $ stateConstSubsts s
          }

-- | 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 -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
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 a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure 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 a. a -> [a]
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