    {-# 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