    -- | Compute a mapping from variables to their corresponding (fully
    -- expanded) PrimExps.
    module Futhark.Analysis.PrimExp.Table
      ( primExpTable,
        PrimExpTable,

        -- * Extensibility
        PrimExpAnalysis (..),

        -- * Testing
        stmToPrimExps,
      )
    where

    import Control.Monad.State.Strict
    import Data.Foldable
    import Data.Map.Strict qualified as M
    import Futhark.Analysis.PrimExp
    import Futhark.Analysis.PrimExp.Convert
    import Futhark.IR.Aliases
    import Futhark.IR.GPU
    import Futhark.IR.GPUMem
    import Futhark.IR.MC
    import Futhark.IR.MCMem

    -- | Maps variables to maybe PrimExps. Will map to nothing if it
    -- cannot be resolved to a PrimExp. For all uses of this analysis atm.
    -- a variable can be considered inscrutable if it cannot be resolved
    -- to a primexp.
    type PrimExpTable = M.Map VName (Maybe (PrimExp VName))

    -- | A class for extracting PrimExps from what is inside an op.
    class PrimExpAnalysis rep where
      opPrimExp :: Scope rep -> Op rep -> State PrimExpTable ()

    primExpTable :: (PrimExpAnalysis rep, RepTypes rep) => Prog rep -> PrimExpTable
    primExpTable :: forall rep.
    (PrimExpAnalysis rep, RepTypes rep) =>
    Prog rep -> PrimExpTable
    primExpTable Prog rep
    prog = PrimExpTable
    initialState PrimExpTable -> PrimExpTable -> PrimExpTable
    forall a. Semigroup a => a -> a -> a
    <> ((Scope rep, FunDef rep) -> PrimExpTable)
    -> [(Scope rep, FunDef rep)] -> PrimExpTable
    forall m a. Monoid m => (a -> m) -> [a] -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap' ((Scope rep -> FunDef rep -> PrimExpTable)
    -> (Scope rep, FunDef rep) -> PrimExpTable
    forall a b c. (a -> b -> c) -> (a, b) -> c
    uncurry Scope rep -> FunDef rep -> PrimExpTable
    forall rep.
    (PrimExpAnalysis rep, RepTypes rep) =>
    Scope rep -> FunDef rep -> PrimExpTable
    funToPrimExp) [(Scope rep, FunDef rep)]
    scopesAndFuns
      where
        scopesAndFuns :: [(Scope rep, FunDef rep)]
    scopesAndFuns = do
          let fun_defs :: [FunDef rep]
    fun_defs = Prog rep -> [FunDef rep]
    forall rep. Prog rep -> [FunDef rep]
    progFuns Prog rep
    prog
          let scopes :: [Scope rep]
    scopes = (FunDef rep -> Scope rep) -> [FunDef rep] -> [Scope rep]
    forall a b. (a -> b) -> [a] -> [b]
    map FunDef rep -> Scope rep
    forall {rep}.
    (FParamInfo rep ~ FParamInfo rep) =>
    FunDef rep -> Scope rep
    getScope [FunDef rep]
    fun_defs
          [Scope rep] -> [FunDef rep] -> [(Scope rep, FunDef rep)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Scope rep]
    scopes [FunDef rep]
    fun_defs

        getScope :: FunDef rep -> Scope rep
    getScope FunDef rep
    funDef = Stms rep -> Scope rep
    forall rep a. Scoped rep a => a -> Scope rep
    scopeOf (Prog rep -> Stms rep
    forall rep. Prog rep -> Stms rep
    progConsts Prog rep
    prog) Scope rep -> Scope rep -> Scope rep
    forall a. Semigroup a => a -> a -> a
    <> [Param (FParamInfo rep)] -> Scope rep
    forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
    scopeOfFParams (FunDef rep -> [FParam rep]
    forall rep. FunDef rep -> [FParam rep]
    funDefParams FunDef rep
    funDef)

        -- We need to have the dummy "slice" in the analysis for our "slice hack".
        initialState :: PrimExpTable
    initialState =
          VName -> Maybe (PrimExp VName) -> PrimExpTable
    forall k a. k -> a -> Map k a
    M.singleton (Name -> Int -> VName
    VName Name
    "slice" Int
    0) (Maybe (PrimExp VName) -> PrimExpTable)
    -> Maybe (PrimExp VName) -> PrimExpTable
    forall a b. (a -> b) -> a -> b
    $ PrimExp VName -> Maybe (PrimExp VName)
    forall a. a -> Maybe a
    Just (PrimExp VName -> Maybe (PrimExp VName))
    -> PrimExp VName -> Maybe (PrimExp VName)
    forall a b. (a -> b) -> a -> b
    $ VName -> PrimType -> PrimExp VName
    forall v. v -> PrimType -> PrimExp v
    LeafExp (Name -> Int -> VName
    VName Name
    "slice" Int
    0) (PrimType -> PrimExp VName) -> PrimType -> PrimExp VName
    forall a b. (a -> b) -> a -> b
    $ IntType -> PrimType
    IntType IntType
    Int64

    funToPrimExp ::
      (PrimExpAnalysis rep, RepTypes rep) =>
      Scope rep ->
      FunDef rep ->
      PrimExpTable
    funToPrimExp :: forall rep.
    (PrimExpAnalysis rep, RepTypes rep) =>
    Scope rep -> FunDef rep -> PrimExpTable
    funToPrimExp Scope rep
    scope FunDef rep
    fundef = State PrimExpTable () -> PrimExpTable -> PrimExpTable
    forall s a. State s a -> s -> s
    execState (Scope rep -> Body rep -> State PrimExpTable ()
    forall rep.
    (PrimExpAnalysis rep, RepTypes rep) =>
    Scope rep -> Body rep -> State PrimExpTable ()
    bodyToPrimExps Scope rep
    scope (FunDef rep -> Body rep
    forall rep. FunDef rep -> Body rep
    funDefBody FunDef rep
    fundef)) PrimExpTable
    forall a. Monoid a => a
    mempty

    -- | Adds the statements of a body to the PrimExpTable
    bodyToPrimExps ::
      (PrimExpAnalysis rep, RepTypes rep) =>
      Scope rep ->
      Body rep ->
      State PrimExpTable ()
    bodyToPrimExps :: forall rep.
    (PrimExpAnalysis rep, RepTypes rep) =>
    Scope rep -> Body rep -> State PrimExpTable ()
    bodyToPrimExps Scope rep
    scope Body rep
    body = (Stm rep -> State PrimExpTable ())
    -> Seq (Stm rep) -> State PrimExpTable ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ (Scope rep -> Stm rep -> State PrimExpTable ()
    forall rep.
    (PrimExpAnalysis rep, RepTypes rep) =>
    Scope rep -> Stm rep -> State PrimExpTable ()
    stmToPrimExps Scope rep
    scope') (Body rep -> Seq (Stm rep)
    forall rep. Body rep -> Stms rep
    bodyStms Body rep
    body)
      where
        scope' :: Scope rep
    scope' = Scope rep
    scope Scope rep -> Scope rep -> Scope rep
    forall a. Semigroup a => a -> a -> a
    <> Seq (Stm rep) -> Scope rep
    forall rep a. Scoped rep a => a -> Scope rep
    scopeOf (Body rep -> Seq (Stm rep)
    forall rep. Body rep -> Stms rep
    bodyStms Body rep
    body)

    -- | Adds the statements of a kernel body to the PrimExpTable
    kernelToBodyPrimExps ::
      (PrimExpAnalysis rep, RepTypes rep) =>
      Scope rep ->
      KernelBody rep ->
      State PrimExpTable ()
    kernelToBodyPrimExps :: forall rep.
    (PrimExpAnalysis rep, RepTypes rep) =>
    Scope rep -> KernelBody rep -> State PrimExpTable ()
    kernelToBodyPrimExps Scope rep
    scope KernelBody rep
    kbody = (Stm rep -> State PrimExpTable ())
    -> Seq (Stm rep) -> State PrimExpTable ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ (Scope rep -> Stm rep -> State PrimExpTable ()
    forall rep.
    (PrimExpAnalysis rep, RepTypes rep) =>
    Scope rep -> Stm rep -> State PrimExpTable ()
    stmToPrimExps Scope rep
    scope') (KernelBody rep -> Seq (Stm rep)
    forall rep. KernelBody rep -> Stms rep
    kernelBodyStms KernelBody rep
    kbody)
      where
        scope' :: Scope rep
    scope' = Scope rep
    scope Scope rep -> Scope rep -> Scope rep
    forall a. Semigroup a => a -> a -> a
    <> Seq (Stm rep) -> Scope rep
    forall rep a. Scoped rep a => a -> Scope rep
    scopeOf (KernelBody rep -> Seq (Stm rep)
    forall rep. KernelBody rep -> Stms rep
    kernelBodyStms KernelBody rep
    kbody)

    -- | Adds a statement to the PrimExpTable. If it can't be resolved as a `PrimExp`,
    -- it will be added as `Nothing`.
    stmToPrimExps ::
      forall rep.
      (PrimExpAnalysis rep, RepTypes rep) =>
      Scope rep ->
      Stm rep ->
      State PrimExpTable ()
    stmToPrimExps :: forall rep.
    (PrimExpAnalysis rep, RepTypes rep) =>
    Scope rep -> Stm rep -> State PrimExpTable ()
    stmToPrimExps Scope rep
    scope Stm rep
    stm = do
      PrimExpTable
    table <- StateT PrimExpTable Identity PrimExpTable
    forall s (m :: * -> *). MonadState s m => m s
    get
      case Stm rep
    stm of
        (Let (Pat [PatElem (LetDec rep)]
    pat_elems) StmAux (ExpDec rep)
    _ Exp rep
    e)
          | Just PrimExp VName
    primExp <- (VName -> Maybe (PrimExp VName))
    -> Exp rep -> Maybe (PrimExp VName)
    forall (m :: * -> *) rep v.
    (MonadFail m, RepTypes rep) =>
    (VName -> m (PrimExp v)) -> Exp rep -> m (PrimExp v)
    primExpFromExp (Scope rep -> PrimExpTable -> VName -> Maybe (PrimExp VName)
    forall rep.
    RepTypes rep =>
    Scope rep -> PrimExpTable -> VName -> Maybe (PrimExp VName)
    toPrimExp Scope rep
    scope PrimExpTable
    table) Exp rep
    e ->
              -- The statement can be resolved as a `PrimExp`.
              -- For each pattern element, insert the PrimExp in the table
              [PatElem (LetDec rep)]
    -> (PatElem (LetDec rep) -> State PrimExpTable ())
    -> State PrimExpTable ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ [PatElem (LetDec rep)]
    pat_elems ((PatElem (LetDec rep) -> State PrimExpTable ())
     -> State PrimExpTable ())
    -> (PatElem (LetDec rep) -> State PrimExpTable ())
    -> State PrimExpTable ()
    forall a b. (a -> b) -> a -> b
    $ \PatElem (LetDec rep)
    pe ->
                (PrimExpTable -> PrimExpTable) -> State PrimExpTable ()
    forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
    modify ((PrimExpTable -> PrimExpTable) -> State PrimExpTable ())
    -> (PrimExpTable -> PrimExpTable) -> State PrimExpTable ()
    forall a b. (a -> b) -> a -> b
    $ VName -> Maybe (PrimExp VName) -> PrimExpTable -> PrimExpTable
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    M.insert (PatElem (LetDec rep) -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem (LetDec rep)
    pe) (PrimExp VName -> Maybe (PrimExp VName)
    forall a. a -> Maybe a
    Just PrimExp VName
    primExp)
          | Bool
    otherwise -> do
              -- The statement can't be resolved as a `PrimExp`.
              Exp rep -> State PrimExpTable ()
    walk (Exp rep -> State PrimExpTable ())
    -> Exp rep -> State PrimExpTable ()
    forall a b. (a -> b) -> a -> b
    $ Stm rep -> Exp rep
    forall rep. Stm rep -> Exp rep
    stmExp Stm rep
    stm -- Traverse the rest of the AST Get the
              -- updated PrimExpTable after traversing the AST
              PrimExpTable
    table' <- StateT PrimExpTable Identity PrimExpTable
    forall s (m :: * -> *). MonadState s m => m s
    get

              -- Add pattern elements that can't be resolved as `PrimExp`
              -- to the `PrimExpTable` as `Nothing`
              [PatElem (LetDec rep)]
    -> (PatElem (LetDec rep) -> State PrimExpTable ())
    -> State PrimExpTable ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ [PatElem (LetDec rep)]
    pat_elems ((PatElem (LetDec rep) -> State PrimExpTable ())
     -> State PrimExpTable ())
    -> (PatElem (LetDec rep) -> State PrimExpTable ())
    -> State PrimExpTable ()
    forall a b. (a -> b) -> a -> b
    $ \PatElem (LetDec rep)
    pe ->
                case VName -> PrimExpTable -> Maybe (Maybe (PrimExp VName))
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup (PatElem (LetDec rep) -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem (LetDec rep)
    pe) PrimExpTable
    table' of
                  Maybe (Maybe (PrimExp VName))
    Nothing -> (PrimExpTable -> PrimExpTable) -> State PrimExpTable ()
    forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
    modify ((PrimExpTable -> PrimExpTable) -> State PrimExpTable ())
    -> (PrimExpTable -> PrimExpTable) -> State PrimExpTable ()
    forall a b. (a -> b) -> a -> b
    $ VName -> Maybe (PrimExp VName) -> PrimExpTable -> PrimExpTable
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    M.insert (PatElem (LetDec rep) -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem (LetDec rep)
    pe) Maybe (PrimExp VName)
    forall a. Maybe a
    Nothing
                  Just Maybe (PrimExp VName)
    _ -> () -> State PrimExpTable ()
    forall a. a -> StateT PrimExpTable Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
      where
        walk :: Exp rep -> State PrimExpTable ()
    walk Exp rep
    e = do
          -- Handle most cases using the walker
          Walker rep (StateT PrimExpTable Identity)
    -> Exp rep -> State PrimExpTable ()
    forall (m :: * -> *) rep.
    Monad m =>
    Walker rep m -> Exp rep -> m ()
    walkExpM Walker rep (StateT PrimExpTable Identity)
    walker Exp rep
    e
          -- Additionally, handle loop parameters
          case Exp rep
    e of
            Loop [(FParam rep, SubExp)]
    _ (ForLoop VName
    i IntType
    t SubExp
    _) Body rep
    _ ->
              (PrimExpTable -> PrimExpTable) -> State PrimExpTable ()
    forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
    modify ((PrimExpTable -> PrimExpTable) -> State PrimExpTable ())
    -> (PrimExpTable -> PrimExpTable) -> State PrimExpTable ()
    forall a b. (a -> b) -> a -> b
    $ VName -> Maybe (PrimExp VName) -> PrimExpTable -> PrimExpTable
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    M.insert VName
    i (Maybe (PrimExp VName) -> PrimExpTable -> PrimExpTable)
    -> Maybe (PrimExp VName) -> PrimExpTable -> PrimExpTable
    forall a b. (a -> b) -> a -> b
    $ PrimExp VName -> Maybe (PrimExp VName)
    forall a. a -> Maybe a
    Just (PrimExp VName -> Maybe (PrimExp VName))
    -> PrimExp VName -> Maybe (PrimExp VName)
    forall a b. (a -> b) -> a -> b
    $ VName -> PrimType -> PrimExp VName
    forall v. v -> PrimType -> PrimExp v
    LeafExp VName
    i (PrimType -> PrimExp VName) -> PrimType -> PrimExp VName
    forall a b. (a -> b) -> a -> b
    $ IntType -> PrimType
    IntType IntType
    t
            Exp rep
    _ -> () -> State PrimExpTable ()
    forall a. a -> StateT PrimExpTable Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()

        walker :: Walker rep (StateT PrimExpTable Identity)
    walker =
          (forall rep (m :: * -> *). Monad m => Walker rep m
    identityWalker @rep)
            { walkOnBody = \Scope rep
    body_scope -> Scope rep -> Body rep -> State PrimExpTable ()
    forall rep.
    (PrimExpAnalysis rep, RepTypes rep) =>
    Scope rep -> Body rep -> State PrimExpTable ()
    bodyToPrimExps (Scope rep
    scope Scope rep -> Scope rep -> Scope rep
    forall a. Semigroup a => a -> a -> a
    <> Scope rep
    body_scope),
              walkOnOp = opPrimExp scope,
              walkOnFParam = paramToPrimExp -- Loop parameters
            }

        -- Adds a loop parameter to the PrimExpTable
        paramToPrimExp :: FParam rep -> State PrimExpTable ()
        paramToPrimExp :: FParam rep -> State PrimExpTable ()
    paramToPrimExp FParam rep
    param = do
          let name :: VName
    name = FParam rep -> VName
    forall dec. Param dec -> VName
    paramName FParam rep
    param
          -- Construct a `PrimExp` from the type of the parameter
          -- and add it to the `PrimExpTable`
          case FParamInfo rep -> Type
    forall t. Typed t => t -> Type
    typeOf (FParamInfo rep -> Type) -> FParamInfo rep -> Type
    forall a b. (a -> b) -> a -> b
    $ FParam rep -> FParamInfo rep
    forall dec. Param dec -> dec
    paramDec FParam rep
    param of
            -- TODO: Handle other types?
            Prim PrimType
    pt ->
              (PrimExpTable -> PrimExpTable) -> State PrimExpTable ()
    forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
    modify ((PrimExpTable -> PrimExpTable) -> State PrimExpTable ())
    -> (PrimExpTable -> PrimExpTable) -> State PrimExpTable ()
    forall a b. (a -> b) -> a -> b
    $ VName -> Maybe (PrimExp VName) -> PrimExpTable -> PrimExpTable
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    M.insert VName
    name (PrimExp VName -> Maybe (PrimExp VName)
    forall a. a -> Maybe a
    Just (PrimExp VName -> Maybe (PrimExp VName))
    -> PrimExp VName -> Maybe (PrimExp VName)
    forall a b. (a -> b) -> a -> b
    $ VName -> PrimType -> PrimExp VName
    forall v. v -> PrimType -> PrimExp v
    LeafExp VName
    name PrimType
    pt)
            Type
    _ -> () -> State PrimExpTable ()
    forall a. a -> StateT PrimExpTable Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()

    -- | Checks if a name is in the PrimExpTable and construct a `PrimExp`
    -- if it is not
    toPrimExp :: (RepTypes rep) => Scope rep -> PrimExpTable -> VName -> Maybe (PrimExp VName)
    toPrimExp :: forall rep.
    RepTypes rep =>
    Scope rep -> PrimExpTable -> VName -> Maybe (PrimExp VName)
    toPrimExp Scope rep
    scope PrimExpTable
    table VName
    name = case VName -> PrimExpTable -> Maybe (Maybe (PrimExp VName))
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    name PrimExpTable
    table of
      Just Maybe (PrimExp VName)
    maybePrimExp
        | Just PrimExp VName
    primExp <- Maybe (PrimExp VName)
    maybePrimExp -> PrimExp VName -> Maybe (PrimExp VName)
    forall a. a -> Maybe a
    Just PrimExp VName
    primExp -- Already in the table
      Maybe (Maybe (PrimExp VName))
    _ -> case (NameInfo rep -> Type) -> Maybe (NameInfo rep) -> Maybe Type
    forall a b. (a -> b) -> Maybe a -> Maybe b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap NameInfo rep -> Type
    forall t. Typed t => t -> Type
    typeOf (Maybe (NameInfo rep) -> Maybe Type)
    -> (Scope rep -> Maybe (NameInfo rep)) -> Scope rep -> Maybe Type
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VName -> Scope rep -> Maybe (NameInfo rep)
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    name (Scope rep -> Maybe Type) -> Scope rep -> Maybe Type
    forall a b. (a -> b) -> a -> b
    $ Scope rep
    scope of
        (Just (Prim PrimType
    pt)) -> PrimExp VName -> Maybe (PrimExp VName)
    forall a. a -> Maybe a
    Just (PrimExp VName -> Maybe (PrimExp VName))
    -> PrimExp VName -> Maybe (PrimExp VName)
    forall a b. (a -> b) -> a -> b
    $ VName -> PrimType -> PrimExp VName
    forall v. v -> PrimType -> PrimExp v
    LeafExp VName
    name PrimType
    pt
        Maybe Type
    _ -> Maybe (PrimExp VName)
    forall a. Maybe a
    Nothing

    -- | Adds the parameters of a SegOp as well as the statements in its
    -- body to the PrimExpTable
    segOpToPrimExps :: (PrimExpAnalysis rep, RepTypes rep) => Scope rep -> SegOp lvl rep -> State PrimExpTable ()
    segOpToPrimExps :: forall rep lvl.
    (PrimExpAnalysis rep, RepTypes rep) =>
    Scope rep -> SegOp lvl rep -> State PrimExpTable ()
    segOpToPrimExps Scope rep
    scope SegOp lvl rep
    op = do
      [VName]
    -> (VName -> State PrimExpTable ()) -> State PrimExpTable ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ (((VName, SubExp) -> VName) -> [(VName, SubExp)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName, SubExp) -> VName
    forall a b. (a, b) -> a
    fst ([(VName, SubExp)] -> [VName]) -> [(VName, SubExp)] -> [VName]
    forall a b. (a -> b) -> a -> b
    $ SegSpace -> [(VName, SubExp)]
    unSegSpace (SegSpace -> [(VName, SubExp)]) -> SegSpace -> [(VName, SubExp)]
    forall a b. (a -> b) -> a -> b
    $ SegOp lvl rep -> SegSpace
    forall lvl rep. SegOp lvl rep -> SegSpace
    segSpace SegOp lvl rep
    op) ((VName -> State PrimExpTable ()) -> State PrimExpTable ())
    -> (VName -> State PrimExpTable ()) -> State PrimExpTable ()
    forall a b. (a -> b) -> a -> b
    $ \VName
    name ->
        (PrimExpTable -> PrimExpTable) -> State PrimExpTable ()
    forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
    modify ((PrimExpTable -> PrimExpTable) -> State PrimExpTable ())
    -> (PrimExpTable -> PrimExpTable) -> State PrimExpTable ()
    forall a b. (a -> b) -> a -> b
    $ VName -> Maybe (PrimExp VName) -> PrimExpTable -> PrimExpTable
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    M.insert VName
    name (Maybe (PrimExp VName) -> PrimExpTable -> PrimExpTable)
    -> Maybe (PrimExp VName) -> PrimExpTable -> PrimExpTable
    forall a b. (a -> b) -> a -> b
    $ PrimExp VName -> Maybe (PrimExp VName)
    forall a. a -> Maybe a
    Just (PrimExp VName -> Maybe (PrimExp VName))
    -> PrimExp VName -> Maybe (PrimExp VName)
    forall a b. (a -> b) -> a -> b
    $ VName -> PrimType -> PrimExp VName
    forall v. v -> PrimType -> PrimExp v
    LeafExp VName
    name PrimType
    int64
      Scope rep -> KernelBody rep -> State PrimExpTable ()
    forall rep.
    (PrimExpAnalysis rep, RepTypes rep) =>
    Scope rep -> KernelBody rep -> State PrimExpTable ()
    kernelToBodyPrimExps Scope rep
    scope (SegOp lvl rep -> KernelBody rep
    forall lvl rep. SegOp lvl rep -> KernelBody rep
    segBody SegOp lvl rep
    op)

    instance PrimExpAnalysis GPU where
      opPrimExp :: Scope GPU -> Op GPU -> State PrimExpTable ()
    opPrimExp Scope GPU
    scope Op GPU
    gpu_op
        | (SegOp SegOp SegLevel GPU
    op) <- Op GPU
    gpu_op = Scope GPU -> SegOp SegLevel GPU -> State PrimExpTable ()
    forall rep lvl.
    (PrimExpAnalysis rep, RepTypes rep) =>
    Scope rep -> SegOp lvl rep -> State PrimExpTable ()
    segOpToPrimExps Scope GPU
    scope SegOp SegLevel GPU
    op
        | (SizeOp SizeOp
    _) <- Op GPU
    gpu_op = () -> State PrimExpTable ()
    forall a. a -> StateT PrimExpTable Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
        | (GPUBody [Type]
    _ Body GPU
    body) <- Op GPU
    gpu_op = Scope GPU -> Body GPU -> State PrimExpTable ()
    forall rep.
    (PrimExpAnalysis rep, RepTypes rep) =>
    Scope rep -> Body rep -> State PrimExpTable ()
    bodyToPrimExps Scope GPU
    scope Body GPU
    body
        | (Futhark.IR.GPUMem.OtherOp SOAC GPU
    _) <- Op GPU
    gpu_op = () -> State PrimExpTable ()
    forall a. a -> StateT PrimExpTable Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()

    instance PrimExpAnalysis MC where
      opPrimExp :: Scope MC -> Op MC -> State PrimExpTable ()
    opPrimExp Scope MC
    scope Op MC
    mc_op
        | (ParOp Maybe (SegOp () MC)
    maybe_par_segop SegOp () MC
    seq_segop) <- Op MC
    mc_op = do
            -- Add the statements in the parallel part of the ParOp to the PrimExpTable
            case Maybe (SegOp () MC)
    maybe_par_segop of
              Maybe (SegOp () MC)
    Nothing -> () -> State PrimExpTable ()
    forall a. a -> StateT PrimExpTable Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
              Just SegOp () MC
    _ -> Maybe (SegOp () MC)
    -> (SegOp () MC -> State PrimExpTable ()) -> State PrimExpTable ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ Maybe (SegOp () MC)
    maybe_par_segop ((SegOp () MC -> State PrimExpTable ()) -> State PrimExpTable ())
    -> (SegOp () MC -> State PrimExpTable ()) -> State PrimExpTable ()
    forall a b. (a -> b) -> a -> b
    $ Scope MC -> SegOp () MC -> State PrimExpTable ()
    forall rep lvl.
    (PrimExpAnalysis rep, RepTypes rep) =>
    Scope rep -> SegOp lvl rep -> State PrimExpTable ()
    segOpToPrimExps Scope MC
    scope
            -- Add the statements in the sequential part of the ParOp to the PrimExpTable
            Scope MC -> SegOp () MC -> State PrimExpTable ()
    forall rep lvl.
    (PrimExpAnalysis rep, RepTypes rep) =>
    Scope rep -> SegOp lvl rep -> State PrimExpTable ()
    segOpToPrimExps Scope MC
    scope SegOp () MC
    seq_segop
        | (Futhark.IR.MCMem.OtherOp SOAC MC
    _) <- Op MC
    mc_op = () -> State PrimExpTable ()
    forall a. a -> StateT PrimExpTable Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()