    {-# LANGUAGE LambdaCase #-}

    module Futhark.Analysis.AccessPattern
      ( analyseDimAccesses,
        analyseFunction,
        vnameFromSegOp,
        analysisPropagateByTransitivity,
        isInvariant,
        Analyse,
        IndexTable,
        ArrayName,
        DimAccess (..),
        IndexExprName,
        BodyType (..),
        SegOpName (SegmentedMap, SegmentedRed, SegmentedScan, SegmentedHist),
        Context (..),
        analyseIndex,
        VariableInfo (..),
        VarType (..),
        isCounter,
        Dependency (..),
      )
    where

    import Data.Bifunctor
    import Data.Foldable
    import Data.List qualified as L
    import Data.Map.Strict qualified as M
    import Data.Maybe
    import Futhark.IR.Aliases
    import Futhark.IR.GPU
    import Futhark.IR.GPUMem
    import Futhark.IR.MC
    import Futhark.IR.MCMem
    import Futhark.IR.SOACS
    import Futhark.IR.Seq
    import Futhark.IR.SeqMem
    import Futhark.Util.Pretty

    -- | Name of a SegOp, used to identify the SegOp that an array access is
    -- contained in.
    data SegOpName
      = SegmentedMap {SegOpName -> VName
    vnameFromSegOp :: VName}
      | SegmentedRed {vnameFromSegOp :: VName}
      | SegmentedScan {vnameFromSegOp :: VName}
      | SegmentedHist {vnameFromSegOp :: VName}
      deriving (SegOpName -> SegOpName -> Bool
    (SegOpName -> SegOpName -> Bool)
    -> (SegOpName -> SegOpName -> Bool) -> Eq SegOpName
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: SegOpName -> SegOpName -> Bool
    == :: SegOpName -> SegOpName -> Bool
    $c/= :: SegOpName -> SegOpName -> Bool
    /= :: SegOpName -> SegOpName -> Bool
    Eq, Eq SegOpName
    Eq SegOpName =>
    (SegOpName -> SegOpName -> Ordering)
    -> (SegOpName -> SegOpName -> Bool)
    -> (SegOpName -> SegOpName -> Bool)
    -> (SegOpName -> SegOpName -> Bool)
    -> (SegOpName -> SegOpName -> Bool)
    -> (SegOpName -> SegOpName -> SegOpName)
    -> (SegOpName -> SegOpName -> SegOpName)
    -> Ord SegOpName
    SegOpName -> SegOpName -> Bool
    SegOpName -> SegOpName -> Ordering
    SegOpName -> SegOpName -> SegOpName
    forall a.
    Eq a =>
    (a -> a -> Ordering)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> a)
    -> (a -> a -> a)
    -> Ord a
    $ccompare :: SegOpName -> SegOpName -> Ordering
    compare :: SegOpName -> SegOpName -> Ordering
    $c< :: SegOpName -> SegOpName -> Bool
    < :: SegOpName -> SegOpName -> Bool
    $c<= :: SegOpName -> SegOpName -> Bool
    <= :: SegOpName -> SegOpName -> Bool
    $c> :: SegOpName -> SegOpName -> Bool
    > :: SegOpName -> SegOpName -> Bool
    $c>= :: SegOpName -> SegOpName -> Bool
    >= :: SegOpName -> SegOpName -> Bool
    $cmax :: SegOpName -> SegOpName -> SegOpName
    max :: SegOpName -> SegOpName -> SegOpName
    $cmin :: SegOpName -> SegOpName -> SegOpName
    min :: SegOpName -> SegOpName -> SegOpName
    Ord, Int -> SegOpName -> ShowS
    [SegOpName] -> ShowS
    SegOpName -> [Char]
    (Int -> SegOpName -> ShowS)
    -> (SegOpName -> [Char])
    -> ([SegOpName] -> ShowS)
    -> Show SegOpName
    forall a.
    (Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> SegOpName -> ShowS
    showsPrec :: Int -> SegOpName -> ShowS
    $cshow :: SegOpName -> [Char]
    show :: SegOpName -> [Char]
    $cshowList :: [SegOpName] -> ShowS
    showList :: [SegOpName] -> ShowS
    Show)

    -- | Name of an array indexing expression. Taken from the pattern of
    -- the expression.
    type IndexExprName = VName

    data BodyType
      = SegOpName SegOpName
      | LoopBodyName VName
      | CondBodyName VName
      deriving (Int -> BodyType -> ShowS
    [BodyType] -> ShowS
    BodyType -> [Char]
    (Int -> BodyType -> ShowS)
    -> (BodyType -> [Char]) -> ([BodyType] -> ShowS) -> Show BodyType
    forall a.
    (Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> BodyType -> ShowS
    showsPrec :: Int -> BodyType -> ShowS
    $cshow :: BodyType -> [Char]
    show :: BodyType -> [Char]
    $cshowList :: [BodyType] -> ShowS
    showList :: [BodyType] -> ShowS
    Show, Eq BodyType
    Eq BodyType =>
    (BodyType -> BodyType -> Ordering)
    -> (BodyType -> BodyType -> Bool)
    -> (BodyType -> BodyType -> Bool)
    -> (BodyType -> BodyType -> Bool)
    -> (BodyType -> BodyType -> Bool)
    -> (BodyType -> BodyType -> BodyType)
    -> (BodyType -> BodyType -> BodyType)
    -> Ord BodyType
    BodyType -> BodyType -> Bool
    BodyType -> BodyType -> Ordering
    BodyType -> BodyType -> BodyType
    forall a.
    Eq a =>
    (a -> a -> Ordering)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> a)
    -> (a -> a -> a)
    -> Ord a
    $ccompare :: BodyType -> BodyType -> Ordering
    compare :: BodyType -> BodyType -> Ordering
    $c< :: BodyType -> BodyType -> Bool
    < :: BodyType -> BodyType -> Bool
    $c<= :: BodyType -> BodyType -> Bool
    <= :: BodyType -> BodyType -> Bool
    $c> :: BodyType -> BodyType -> Bool
    > :: BodyType -> BodyType -> Bool
    $c>= :: BodyType -> BodyType -> Bool
    >= :: BodyType -> BodyType -> Bool
    $cmax :: BodyType -> BodyType -> BodyType
    max :: BodyType -> BodyType -> BodyType
    $cmin :: BodyType -> BodyType -> BodyType
    min :: BodyType -> BodyType -> BodyType
    Ord, BodyType -> BodyType -> Bool
    (BodyType -> BodyType -> Bool)
    -> (BodyType -> BodyType -> Bool) -> Eq BodyType
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: BodyType -> BodyType -> Bool
    == :: BodyType -> BodyType -> Bool
    $c/= :: BodyType -> BodyType -> Bool
    /= :: BodyType -> BodyType -> Bool
    Eq)

    -- | Stores the name of an array, the nest of loops, kernels,
    -- conditionals in which it is constructed, and the existing layout of
    -- the array. The latter is currently largely unused and not
    -- trustworthy, but might be useful in the future.
    type ArrayName = (VName, [BodyType], [Int])

    -- | Tuple of patternName and nested `level` it index occurred at, as well as
    -- what the actual iteration type is.
    data Dependency = Dependency
      { Dependency -> Int
    lvl :: Int,
        Dependency -> VarType
    varType :: VarType
      }
      deriving (Dependency -> Dependency -> Bool
    (Dependency -> Dependency -> Bool)
    -> (Dependency -> Dependency -> Bool) -> Eq Dependency
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: Dependency -> Dependency -> Bool
    == :: Dependency -> Dependency -> Bool
    $c/= :: Dependency -> Dependency -> Bool
    /= :: Dependency -> Dependency -> Bool
    Eq, Int -> Dependency -> ShowS
    [Dependency] -> ShowS
    Dependency -> [Char]
    (Int -> Dependency -> ShowS)
    -> (Dependency -> [Char])
    -> ([Dependency] -> ShowS)
    -> Show Dependency
    forall a.
    (Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> Dependency -> ShowS
    showsPrec :: Int -> Dependency -> ShowS
    $cshow :: Dependency -> [Char]
    show :: Dependency -> [Char]
    $cshowList :: [Dependency] -> ShowS
    showList :: [Dependency] -> ShowS
    Show)

    -- | Collect all features of access to a specific dimension of an array.
    data DimAccess rep = DimAccess
      { -- | Set of VNames of iteration variables (gtids, loop counters, etc.)
        -- that some access is variant to.
        -- An empty set indicates that the access is invariant.
        forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
    dependencies :: M.Map VName Dependency,
        -- | Used to store the name of the original expression from which `dependencies`
        -- was computed. `Nothing` if it is a constant.
        forall {k} (rep :: k). DimAccess rep -> Maybe VName
    originalVar :: Maybe VName
      }
      deriving (DimAccess rep -> DimAccess rep -> Bool
    (DimAccess rep -> DimAccess rep -> Bool)
    -> (DimAccess rep -> DimAccess rep -> Bool) -> Eq (DimAccess rep)
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    forall k (rep :: k). DimAccess rep -> DimAccess rep -> Bool
    $c== :: forall k (rep :: k). DimAccess rep -> DimAccess rep -> Bool
    == :: DimAccess rep -> DimAccess rep -> Bool
    $c/= :: forall k (rep :: k). DimAccess rep -> DimAccess rep -> Bool
    /= :: DimAccess rep -> DimAccess rep -> Bool
    Eq, Int -> DimAccess rep -> ShowS
    [DimAccess rep] -> ShowS
    DimAccess rep -> [Char]
    (Int -> DimAccess rep -> ShowS)
    -> (DimAccess rep -> [Char])
    -> ([DimAccess rep] -> ShowS)
    -> Show (DimAccess rep)
    forall a.
    (Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
    forall k (rep :: k). Int -> DimAccess rep -> ShowS
    forall k (rep :: k). [DimAccess rep] -> ShowS
    forall k (rep :: k). DimAccess rep -> [Char]
    $cshowsPrec :: forall k (rep :: k). Int -> DimAccess rep -> ShowS
    showsPrec :: Int -> DimAccess rep -> ShowS
    $cshow :: forall k (rep :: k). DimAccess rep -> [Char]
    show :: DimAccess rep -> [Char]
    $cshowList :: forall k (rep :: k). [DimAccess rep] -> ShowS
    showList :: [DimAccess rep] -> ShowS
    Show)

    instance Semigroup (DimAccess rep) where
      DimAccess rep
    adeps <> :: DimAccess rep -> DimAccess rep -> DimAccess rep
    <> DimAccess rep
    bdeps =
        Map VName Dependency -> Maybe VName -> DimAccess rep
    forall {k} (rep :: k).
    Map VName Dependency -> Maybe VName -> DimAccess rep
    DimAccess
          (DimAccess rep -> Map VName Dependency
    forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
    dependencies DimAccess rep
    adeps Map VName Dependency
    -> Map VName Dependency -> Map VName Dependency
    forall a. Semigroup a => a -> a -> a
    <> DimAccess rep -> Map VName Dependency
    forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
    dependencies DimAccess rep
    bdeps)
          ( case DimAccess rep -> Maybe VName
    forall {k} (rep :: k). DimAccess rep -> Maybe VName
    originalVar DimAccess rep
    adeps of
              Maybe VName
    Nothing -> DimAccess rep -> Maybe VName
    forall {k} (rep :: k). DimAccess rep -> Maybe VName
    originalVar DimAccess rep
    bdeps
              Maybe VName
    _ -> DimAccess rep -> Maybe VName
    forall {k} (rep :: k). DimAccess rep -> Maybe VName
    originalVar DimAccess rep
    adeps
          )

    instance Monoid (DimAccess rep) where
      mempty :: DimAccess rep
    mempty = Map VName Dependency -> Maybe VName -> DimAccess rep
    forall {k} (rep :: k).
    Map VName Dependency -> Maybe VName -> DimAccess rep
    DimAccess Map VName Dependency
    forall a. Monoid a => a
    mempty Maybe VName
    forall a. Maybe a
    Nothing

    isInvariant :: DimAccess rep -> Bool
    isInvariant :: forall {k} (rep :: k). DimAccess rep -> Bool
    isInvariant = Map VName Dependency -> Bool
    forall a. Map VName a -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null (Map VName Dependency -> Bool)
    -> (DimAccess rep -> Map VName Dependency) -> DimAccess rep -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . DimAccess rep -> Map VName Dependency
    forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
    dependencies

    -- | For each array access in a program, this data structure stores the
    -- dependencies of each dimension in the access, the array name, and the
    -- name of the SegOp that the access is contained in.
    -- Each DimAccess element corresponds to an access to a given dimension
    -- in the given array, in the same order of the dimensions.
    type IndexTable rep =
      M.Map SegOpName (M.Map ArrayName (M.Map IndexExprName [DimAccess rep]))

    unionIndexTables :: IndexTable rep -> IndexTable rep -> IndexTable rep
    unionIndexTables :: forall {k} (rep :: k).
    IndexTable rep -> IndexTable rep -> IndexTable rep
    unionIndexTables = (Map ArrayName (Map VName [DimAccess rep])
     -> Map ArrayName (Map VName [DimAccess rep])
     -> Map ArrayName (Map VName [DimAccess rep]))
    -> Map SegOpName (Map ArrayName (Map VName [DimAccess rep]))
    -> Map SegOpName (Map ArrayName (Map VName [DimAccess rep]))
    -> Map SegOpName (Map ArrayName (Map VName [DimAccess rep]))
    forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
    M.unionWith ((Map VName [DimAccess rep]
     -> Map VName [DimAccess rep] -> Map VName [DimAccess rep])
    -> Map ArrayName (Map VName [DimAccess rep])
    -> Map ArrayName (Map VName [DimAccess rep])
    -> Map ArrayName (Map VName [DimAccess rep])
    forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
    M.unionWith Map VName [DimAccess rep]
    -> Map VName [DimAccess rep] -> Map VName [DimAccess rep]
    forall k a. Ord k => Map k a -> Map k a -> Map k a
    M.union)

    -- | Make segops on arrays transitive, ie. if
    -- > let A = segmap (..) xs -- A indexes into xs
    -- > let B = segmap (..) A  -- B indexes into A
    -- Then B also derives all A's array-accesses, like xs.
    -- Runs in n²
    analysisPropagateByTransitivity :: IndexTable rep -> IndexTable rep
    analysisPropagateByTransitivity :: forall {k} (rep :: k). IndexTable rep -> IndexTable rep
    analysisPropagateByTransitivity IndexTable rep
    idx_table =
      (Map ArrayName (Map VName [DimAccess rep])
     -> Map ArrayName (Map VName [DimAccess rep]))
    -> IndexTable rep -> IndexTable rep
    forall a b k. (a -> b) -> Map k a -> Map k b
    M.map Map ArrayName (Map VName [DimAccess rep])
    -> Map ArrayName (Map VName [DimAccess rep])
    foldlArrayNameMap IndexTable rep
    idx_table
      where
        aggregateResults :: VName -> Map ArrayName (Map VName [DimAccess rep])
    aggregateResults VName
    arr_name =
          Map ArrayName (Map VName [DimAccess rep])
    -> (Map ArrayName (Map VName [DimAccess rep])
        -> Map ArrayName (Map VName [DimAccess rep]))
    -> Maybe (Map ArrayName (Map VName [DimAccess rep]))
    -> Map ArrayName (Map VName [DimAccess rep])
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe
            Map ArrayName (Map VName [DimAccess rep])
    forall a. Monoid a => a
    mempty
            Map ArrayName (Map VName [DimAccess rep])
    -> Map ArrayName (Map VName [DimAccess rep])
    foldlArrayNameMap
            ((SegOpName -> VName)
    -> IndexTable rep
    -> Map VName (Map ArrayName (Map VName [DimAccess rep]))
    forall k2 k1 a. Ord k2 => (k1 -> k2) -> Map k1 a -> Map k2 a
    M.mapKeys SegOpName -> VName
    vnameFromSegOp IndexTable rep
    idx_table Map VName (Map ArrayName (Map VName [DimAccess rep]))
    -> VName -> Maybe (Map ArrayName (Map VName [DimAccess rep]))
    forall k a. Ord k => Map k a -> k -> Maybe a
    M.!? VName
    arr_name)

        foldlArrayNameMap :: Map ArrayName (Map VName [DimAccess rep])
    -> Map ArrayName (Map VName [DimAccess rep])
    foldlArrayNameMap Map ArrayName (Map VName [DimAccess rep])
    aMap =
          (Map ArrayName (Map VName [DimAccess rep])
     -> Map ArrayName (Map VName [DimAccess rep])
     -> Map ArrayName (Map VName [DimAccess rep]))
    -> Map ArrayName (Map VName [DimAccess rep])
    -> [Map ArrayName (Map VName [DimAccess rep])]
    -> Map ArrayName (Map VName [DimAccess rep])
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl ((Map VName [DimAccess rep]
     -> Map VName [DimAccess rep] -> Map VName [DimAccess rep])
    -> Map ArrayName (Map VName [DimAccess rep])
    -> Map ArrayName (Map VName [DimAccess rep])
    -> Map ArrayName (Map VName [DimAccess rep])
    forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
    M.unionWith Map VName [DimAccess rep]
    -> Map VName [DimAccess rep] -> Map VName [DimAccess rep]
    forall k a. Ord k => Map k a -> Map k a -> Map k a
    M.union) Map ArrayName (Map VName [DimAccess rep])
    aMap ([Map ArrayName (Map VName [DimAccess rep])]
     -> Map ArrayName (Map VName [DimAccess rep]))
    -> [Map ArrayName (Map VName [DimAccess rep])]
    -> Map ArrayName (Map VName [DimAccess rep])
    forall a b. (a -> b) -> a -> b
    $
            (ArrayName -> Map ArrayName (Map VName [DimAccess rep]))
    -> [ArrayName] -> [Map ArrayName (Map VName [DimAccess rep])]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName -> Map ArrayName (Map VName [DimAccess rep])
    aggregateResults (VName -> Map ArrayName (Map VName [DimAccess rep]))
    -> (ArrayName -> VName)
    -> ArrayName
    -> Map ArrayName (Map VName [DimAccess rep])
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . \(VName
    a, [BodyType]
    _, [Int]
    _) -> VName
    a) ([ArrayName] -> [Map ArrayName (Map VName [DimAccess rep])])
    -> [ArrayName] -> [Map ArrayName (Map VName [DimAccess rep])]
    forall a b. (a -> b) -> a -> b
    $
              Map ArrayName (Map VName [DimAccess rep]) -> [ArrayName]
    forall k a. Map k a -> [k]
    M.keys Map ArrayName (Map VName [DimAccess rep])
    aMap

    --
    -- Helper types and functions to perform the analysis.
    --

    -- | Used during the analysis to keep track of the dependencies of patterns
    -- encountered so far.
    data Context rep = Context
      { -- | A mapping from patterns occuring in Let expressions to their dependencies
        --  and iteration types.
        forall {k} (rep :: k). Context rep -> Map VName (VariableInfo rep)
    assignments :: M.Map VName (VariableInfo rep),
        -- | Maps from sliced arrays to their respective access patterns.
        forall {k} (rep :: k).
    Context rep -> Map VName (ArrayName, [VName], [DimAccess rep])
    slices :: M.Map IndexExprName (ArrayName, [VName], [DimAccess rep]),
        -- | A list of the segMaps encountered during the analysis in the order they
        -- were encountered.
        forall {k} (rep :: k). Context rep -> [BodyType]
    parents :: [BodyType],
        -- | Current level of recursion, also just `length parents`
        forall {k} (rep :: k). Context rep -> Int
    currentLevel :: Int
      }
      deriving (Int -> Context rep -> ShowS
    [Context rep] -> ShowS
    Context rep -> [Char]
    (Int -> Context rep -> ShowS)
    -> (Context rep -> [Char])
    -> ([Context rep] -> ShowS)
    -> Show (Context rep)
    forall a.
    (Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
    forall k (rep :: k). Int -> Context rep -> ShowS
    forall k (rep :: k). [Context rep] -> ShowS
    forall k (rep :: k). Context rep -> [Char]
    $cshowsPrec :: forall k (rep :: k). Int -> Context rep -> ShowS
    showsPrec :: Int -> Context rep -> ShowS
    $cshow :: forall k (rep :: k). Context rep -> [Char]
    show :: Context rep -> [Char]
    $cshowList :: forall k (rep :: k). [Context rep] -> ShowS
    showList :: [Context rep] -> ShowS
    Show, Context rep -> Context rep -> Bool
    (Context rep -> Context rep -> Bool)
    -> (Context rep -> Context rep -> Bool) -> Eq (Context rep)
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    forall k (rep :: k). Context rep -> Context rep -> Bool
    $c== :: forall k (rep :: k). Context rep -> Context rep -> Bool
    == :: Context rep -> Context rep -> Bool
    $c/= :: forall k (rep :: k). Context rep -> Context rep -> Bool
    /= :: Context rep -> Context rep -> Bool
    Eq)

    instance Monoid (Context rep) where
      mempty :: Context rep
    mempty =
        Context
          { assignments :: Map VName (VariableInfo rep)
    assignments = Map VName (VariableInfo rep)
    forall a. Monoid a => a
    mempty,
            slices :: Map VName (ArrayName, [VName], [DimAccess rep])
    slices = Map VName (ArrayName, [VName], [DimAccess rep])
    forall a. Monoid a => a
    mempty,
            parents :: [BodyType]
    parents = [],
            currentLevel :: Int
    currentLevel = Int
    0
          }

    instance Semigroup (Context rep) where
      Context Map VName (VariableInfo rep)
    ass0 Map VName (ArrayName, [VName], [DimAccess rep])
    slices0 [BodyType]
    lastBody0 Int
    lvl0 <> :: Context rep -> Context rep -> Context rep
    <> Context Map VName (VariableInfo rep)
    ass1 Map VName (ArrayName, [VName], [DimAccess rep])
    slices1 [BodyType]
    lastBody1 Int
    lvl1 =
        Map VName (VariableInfo rep)
    -> Map VName (ArrayName, [VName], [DimAccess rep])
    -> [BodyType]
    -> Int
    -> Context rep
    forall {k} (rep :: k).
    Map VName (VariableInfo rep)
    -> Map VName (ArrayName, [VName], [DimAccess rep])
    -> [BodyType]
    -> Int
    -> Context rep
    Context
          (Map VName (VariableInfo rep)
    ass0 Map VName (VariableInfo rep)
    -> Map VName (VariableInfo rep) -> Map VName (VariableInfo rep)
    forall a. Semigroup a => a -> a -> a
    <> Map VName (VariableInfo rep)
    ass1)
          (Map VName (ArrayName, [VName], [DimAccess rep])
    slices0 Map VName (ArrayName, [VName], [DimAccess rep])
    -> Map VName (ArrayName, [VName], [DimAccess rep])
    -> Map VName (ArrayName, [VName], [DimAccess rep])
    forall a. Semigroup a => a -> a -> a
    <> Map VName (ArrayName, [VName], [DimAccess rep])
    slices1)
          ([BodyType]
    lastBody0 [BodyType] -> [BodyType] -> [BodyType]
    forall a. Semigroup a => a -> a -> a
    <> [BodyType]
    lastBody1)
          (Int -> Int -> Int
    forall a. Ord a => a -> a -> a
    max Int
    lvl0 Int
    lvl1)

    -- | Extend a context with another context.
    -- We never have to consider the case where VNames clash in the context, since
    -- they are unique.
    extend :: Context rep -> Context rep -> Context rep
    extend :: forall k (rep :: k). Context rep -> Context rep -> Context rep
    extend = Context rep -> Context rep -> Context rep
    forall a. Semigroup a => a -> a -> a
    (<>)

    allSegMap :: Context rep -> [SegOpName]
    allSegMap :: forall {k} (rep :: k). Context rep -> [SegOpName]
    allSegMap (Context Map VName (VariableInfo rep)
    _ Map VName (ArrayName, [VName], [DimAccess rep])
    _ [BodyType]
    parents Int
    _) = (BodyType -> Maybe SegOpName) -> [BodyType] -> [SegOpName]
    forall a b. (a -> Maybe b) -> [a] -> [b]
    mapMaybe BodyType -> Maybe SegOpName
    f [BodyType]
    parents
      where
        f :: BodyType -> Maybe SegOpName
    f (SegOpName SegOpName
    o) = SegOpName -> Maybe SegOpName
    forall a. a -> Maybe a
    Just SegOpName
    o
        f BodyType
    _ = Maybe SegOpName
    forall a. Maybe a
    Nothing

    -- | Context Value (VariableInfo) is the type used in the context to categorize
    -- assignments. For example, a pattern might depend on a function parameter, a
    -- gtid, or some other pattern.
    data VariableInfo rep = VariableInfo
      { forall {k} (rep :: k). VariableInfo rep -> Names
    deps :: Names,
        forall {k} (rep :: k). VariableInfo rep -> Int
    level :: Int,
        forall {k} (rep :: k). VariableInfo rep -> [BodyType]
    parents_nest :: [BodyType],
        forall {k} (rep :: k). VariableInfo rep -> VarType
    variableType :: VarType
      }
      deriving (Int -> VariableInfo rep -> ShowS
    [VariableInfo rep] -> ShowS
    VariableInfo rep -> [Char]
    (Int -> VariableInfo rep -> ShowS)
    -> (VariableInfo rep -> [Char])
    -> ([VariableInfo rep] -> ShowS)
    -> Show (VariableInfo rep)
    forall a.
    (Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
    forall k (rep :: k). Int -> VariableInfo rep -> ShowS
    forall k (rep :: k). [VariableInfo rep] -> ShowS
    forall k (rep :: k). VariableInfo rep -> [Char]
    $cshowsPrec :: forall k (rep :: k). Int -> VariableInfo rep -> ShowS
    showsPrec :: Int -> VariableInfo rep -> ShowS
    $cshow :: forall k (rep :: k). VariableInfo rep -> [Char]
    show :: VariableInfo rep -> [Char]
    $cshowList :: forall k (rep :: k). [VariableInfo rep] -> ShowS
    showList :: [VariableInfo rep] -> ShowS
    Show, VariableInfo rep -> VariableInfo rep -> Bool
    (VariableInfo rep -> VariableInfo rep -> Bool)
    -> (VariableInfo rep -> VariableInfo rep -> Bool)
    -> Eq (VariableInfo rep)
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    forall k (rep :: k). VariableInfo rep -> VariableInfo rep -> Bool
    $c== :: forall k (rep :: k). VariableInfo rep -> VariableInfo rep -> Bool
    == :: VariableInfo rep -> VariableInfo rep -> Bool
    $c/= :: forall k (rep :: k). VariableInfo rep -> VariableInfo rep -> Bool
    /= :: VariableInfo rep -> VariableInfo rep -> Bool
    Eq)

    data VarType
      = ConstType
      | Variable
      | ThreadID
      | LoopVar
      deriving (Int -> VarType -> ShowS
    [VarType] -> ShowS
    VarType -> [Char]
    (Int -> VarType -> ShowS)
    -> (VarType -> [Char]) -> ([VarType] -> ShowS) -> Show VarType
    forall a.
    (Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> VarType -> ShowS
    showsPrec :: Int -> VarType -> ShowS
    $cshow :: VarType -> [Char]
    show :: VarType -> [Char]
    $cshowList :: [VarType] -> ShowS
    showList :: [VarType] -> ShowS
    Show, VarType -> VarType -> Bool
    (VarType -> VarType -> Bool)
    -> (VarType -> VarType -> Bool) -> Eq VarType
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: VarType -> VarType -> Bool
    == :: VarType -> VarType -> Bool
    $c/= :: VarType -> VarType -> Bool
    /= :: VarType -> VarType -> Bool
    Eq)

    isCounter :: VarType -> Bool
    isCounter :: VarType -> Bool
    isCounter VarType
    LoopVar = Bool
    True
    isCounter VarType
    ThreadID = Bool
    True
    isCounter VarType
    _ = Bool
    False

    varInfoFromNames :: Context rep -> Names -> VariableInfo rep
    varInfoFromNames :: forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
    varInfoFromNames Context rep
    ctx Names
    names = do
      Names -> Int -> [BodyType] -> VarType -> VariableInfo rep
    forall {k} (rep :: k).
    Names -> Int -> [BodyType] -> VarType -> VariableInfo rep
    VariableInfo Names
    names (Context rep -> Int
    forall {k} (rep :: k). Context rep -> Int
    currentLevel Context rep
    ctx) (Context rep -> [BodyType]
    forall {k} (rep :: k). Context rep -> [BodyType]
    parents Context rep
    ctx) VarType
    Variable

    -- | Wrapper around the constructur of Context.
    oneContext :: VName -> VariableInfo rep -> Context rep
    oneContext :: forall {k} (rep :: k). VName -> VariableInfo rep -> Context rep
    oneContext VName
    name VariableInfo rep
    var_info =
      Context
        { assignments :: Map VName (VariableInfo rep)
    assignments = VName -> VariableInfo rep -> Map VName (VariableInfo rep)
    forall k a. k -> a -> Map k a
    M.singleton VName
    name VariableInfo rep
    var_info,
          slices :: Map VName (ArrayName, [VName], [DimAccess rep])
    slices = Map VName (ArrayName, [VName], [DimAccess rep])
    forall a. Monoid a => a
    mempty,
          parents :: [BodyType]
    parents = [],
          currentLevel :: Int
    currentLevel = Int
    0
        }

    -- | Create a singular varInfo with no dependencies.
    varInfoZeroDeps :: Context rep -> VariableInfo rep
    varInfoZeroDeps :: forall {k} (rep :: k). Context rep -> VariableInfo rep
    varInfoZeroDeps Context rep
    ctx =
      Names -> Int -> [BodyType] -> VarType -> VariableInfo rep
    forall {k} (rep :: k).
    Names -> Int -> [BodyType] -> VarType -> VariableInfo rep
    VariableInfo Names
    forall a. Monoid a => a
    mempty (Context rep -> Int
    forall {k} (rep :: k). Context rep -> Int
    currentLevel Context rep
    ctx) (Context rep -> [BodyType]
    forall {k} (rep :: k). Context rep -> [BodyType]
    parents Context rep
    ctx) VarType
    Variable

    -- | Create a singular context from a segspace
    contextFromNames :: Context rep -> VariableInfo rep -> [VName] -> Context rep
    contextFromNames :: forall {k} (rep :: k).
    Context rep -> VariableInfo rep -> [VName] -> Context rep
    contextFromNames Context rep
    ctx VariableInfo rep
    var_info = (Context rep -> Context rep -> Context rep)
    -> Context rep -> [Context rep] -> Context rep
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl' Context rep -> Context rep -> Context rep
    forall k (rep :: k). Context rep -> Context rep -> Context rep
    extend Context rep
    ctx ([Context rep] -> Context rep)
    -> ([VName] -> [Context rep]) -> [VName] -> Context rep
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VName -> Context rep) -> [VName] -> [Context rep]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName -> VariableInfo rep -> Context rep
    forall {k} (rep :: k). VName -> VariableInfo rep -> Context rep
    `oneContext` VariableInfo rep
    var_info)

    -- | A representation where we can analyse access patterns.
    class Analyse rep where
      -- | Analyse the op for this representation.
      analyseOp :: Op rep -> Context rep -> [VName] -> (Context rep, IndexTable rep)

    -- | Analyse each entry and accumulate the results.
    analyseDimAccesses :: (Analyse rep) => Prog rep -> IndexTable rep
    analyseDimAccesses :: forall rep. Analyse rep => Prog rep -> IndexTable rep
    analyseDimAccesses = (FunDef rep -> IndexTable rep) -> [FunDef rep] -> IndexTable rep
    forall m a. Monoid m => (a -> m) -> [a] -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap' FunDef rep -> IndexTable rep
    forall rep. Analyse rep => FunDef rep -> IndexTable rep
    analyseFunction ([FunDef rep] -> IndexTable rep)
    -> (Prog rep -> [FunDef rep]) -> Prog rep -> IndexTable rep
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Prog rep -> [FunDef rep]
    forall rep. Prog rep -> [FunDef rep]
    progFuns

    -- | Analyse each statement in a function body.
    analyseFunction :: (Analyse rep) => FunDef rep -> IndexTable rep
    analyseFunction :: forall rep. Analyse rep => FunDef rep -> IndexTable rep
    analyseFunction FunDef rep
    func =
      let stms :: [Stm rep]
    stms = Stms rep -> [Stm rep]
    forall rep. Stms rep -> [Stm rep]
    stmsToList (Stms rep -> [Stm rep])
    -> (Body rep -> Stms rep) -> Body rep -> [Stm rep]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Body rep -> Stms rep
    forall rep. Body rep -> Stms rep
    bodyStms (Body rep -> [Stm rep]) -> Body rep -> [Stm rep]
    forall a b. (a -> b) -> a -> b
    $ FunDef rep -> Body rep
    forall rep. FunDef rep -> Body rep
    funDefBody FunDef rep
    func
          -- Create a context containing the function parameters
          ctx :: Context rep
    ctx = Context rep -> VariableInfo rep -> [VName] -> Context rep
    forall {k} (rep :: k).
    Context rep -> VariableInfo rep -> [VName] -> Context rep
    contextFromNames Context rep
    forall a. Monoid a => a
    mempty (Context rep -> VariableInfo rep
    forall {k} (rep :: k). Context rep -> VariableInfo rep
    varInfoZeroDeps Context rep
    ctx) ([VName] -> Context rep) -> [VName] -> Context rep
    forall a b. (a -> b) -> a -> b
    $ (Param (FParamInfo rep) -> VName)
    -> [Param (FParamInfo rep)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map Param (FParamInfo rep) -> VName
    forall dec. Param dec -> VName
    paramName ([Param (FParamInfo rep)] -> [VName])
    -> [Param (FParamInfo rep)] -> [VName]
    forall a b. (a -> b) -> a -> b
    $ FunDef rep -> [Param (FParamInfo rep)]
    forall rep. FunDef rep -> [FParam rep]
    funDefParams FunDef rep
    func
       in (Context rep, IndexTable rep) -> IndexTable rep
    forall a b. (a, b) -> b
    snd ((Context rep, IndexTable rep) -> IndexTable rep)
    -> (Context rep, IndexTable rep) -> IndexTable rep
    forall a b. (a -> b) -> a -> b
    $ Context rep -> [Stm rep] -> (Context rep, IndexTable rep)
    forall rep.
    Analyse rep =>
    Context rep -> [Stm rep] -> (Context rep, IndexTable rep)
    analyseStmsPrimitive Context rep
    forall k (rep :: k). Context rep
    ctx [Stm rep]
    stms

    -- | Analyse each statement in a list of statements.
    analyseStmsPrimitive :: (Analyse rep) => Context rep -> [Stm rep] -> (Context rep, IndexTable rep)
    analyseStmsPrimitive :: forall rep.
    Analyse rep =>
    Context rep -> [Stm rep] -> (Context rep, IndexTable rep)
    analyseStmsPrimitive Context rep
    ctx =
      -- Fold over statements in body
      ((Context rep, IndexTable rep)
     -> Stm rep -> (Context rep, IndexTable rep))
    -> (Context rep, IndexTable rep)
    -> [Stm rep]
    -> (Context rep, IndexTable rep)
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl'
        (\(Context rep
    c, IndexTable rep
    r) Stm rep
    stm -> (IndexTable rep -> IndexTable rep)
    -> (Context rep, IndexTable rep) -> (Context rep, IndexTable rep)
    forall b c a. (b -> c) -> (a, b) -> (a, c)
    forall (p :: * -> * -> *) b c a.
    Bifunctor p =>
    (b -> c) -> p a b -> p a c
    second (IndexTable rep -> IndexTable rep -> IndexTable rep
    forall {k} (rep :: k).
    IndexTable rep -> IndexTable rep -> IndexTable rep
    unionIndexTables IndexTable rep
    r) ((Context rep, IndexTable rep) -> (Context rep, IndexTable rep))
    -> (Context rep, IndexTable rep) -> (Context rep, IndexTable rep)
    forall a b. (a -> b) -> a -> b
    $ Context rep -> Stm rep -> (Context rep, IndexTable rep)
    forall rep.
    Analyse rep =>
    Context rep -> Stm rep -> (Context rep, IndexTable rep)
    analyseStm Context rep
    c Stm rep
    stm)
        (Context rep
    ctx, IndexTable rep
    forall a. Monoid a => a
    mempty)

    -- | Same as analyseStmsPrimitive, but change the resulting context into
    -- a varInfo, mapped to pattern.
    analyseStms :: (Analyse rep) => Context rep -> (VName -> BodyType) -> [VName] -> [Stm rep] -> (Context rep, IndexTable rep)
    analyseStms :: forall rep.
    Analyse rep =>
    Context rep
    -> (VName -> BodyType)
    -> [VName]
    -> [Stm rep]
    -> (Context rep, IndexTable rep)
    analyseStms Context rep
    ctx VName -> BodyType
    body_constructor [VName]
    pats [Stm rep]
    body = do
      -- 0. Recurse into body with ctx
      let (Context rep
    ctx'', IndexTable rep
    indexTable) = Context rep -> [Stm rep] -> (Context rep, IndexTable rep)
    forall rep.
    Analyse rep =>
    Context rep -> [Stm rep] -> (Context rep, IndexTable rep)
    analyseStmsPrimitive Context rep
    recContext [Stm rep]
    body

      -- 0.1 Get all new slices
      let slices_new :: Map VName (ArrayName, [VName], [DimAccess rep])
    slices_new = Map VName (ArrayName, [VName], [DimAccess rep])
    -> Map VName (ArrayName, [VName], [DimAccess rep])
    -> Map VName (ArrayName, [VName], [DimAccess rep])
    forall k a b. Ord k => Map k a -> Map k b -> Map k a
    M.difference (Context rep -> Map VName (ArrayName, [VName], [DimAccess rep])
    forall {k} (rep :: k).
    Context rep -> Map VName (ArrayName, [VName], [DimAccess rep])
    slices Context rep
    ctx'') (Context rep -> Map VName (ArrayName, [VName], [DimAccess rep])
    forall {k} (rep :: k).
    Context rep -> Map VName (ArrayName, [VName], [DimAccess rep])
    slices Context rep
    ctx)
      -- 0.2 Make "IndexExpressions" of the slices
      let slices_indices :: IndexTable rep
    slices_indices =
            (IndexTable rep -> IndexTable rep -> IndexTable rep)
    -> IndexTable rep -> [IndexTable rep] -> IndexTable rep
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl IndexTable rep -> IndexTable rep -> IndexTable rep
    forall {k} (rep :: k).
    IndexTable rep -> IndexTable rep -> IndexTable rep
    unionIndexTables IndexTable rep
    indexTable
              ([IndexTable rep] -> IndexTable rep)
    -> [IndexTable rep] -> IndexTable rep
    forall a b. (a -> b) -> a -> b
    $ ((VName, (ArrayName, [VName], [DimAccess rep]))
     -> Maybe (IndexTable rep))
    -> [(VName, (ArrayName, [VName], [DimAccess rep]))]
    -> [IndexTable rep]
    forall a b. (a -> Maybe b) -> [a] -> [b]
    mapMaybe
                ( (VName
     -> (ArrayName, [VName], [DimAccess rep]) -> Maybe (IndexTable rep))
    -> (VName, (ArrayName, [VName], [DimAccess rep]))
    -> Maybe (IndexTable rep)
    forall a b c. (a -> b -> c) -> (a, b) -> c
    uncurry ((VName
      -> (ArrayName, [VName], [DimAccess rep]) -> Maybe (IndexTable rep))
     -> (VName, (ArrayName, [VName], [DimAccess rep]))
     -> Maybe (IndexTable rep))
    -> (VName
        -> (ArrayName, [VName], [DimAccess rep]) -> Maybe (IndexTable rep))
    -> (VName, (ArrayName, [VName], [DimAccess rep]))
    -> Maybe (IndexTable rep)
    forall a b. (a -> b) -> a -> b
    $ \VName
    _idx_expression (ArrayName
    array_name, [VName]
    patterns, [DimAccess rep]
    dim_indices) ->
                    IndexTable rep -> Maybe (IndexTable rep)
    forall a. a -> Maybe a
    Just (IndexTable rep -> Maybe (IndexTable rep))
    -> ((Context rep, IndexTable rep) -> IndexTable rep)
    -> (Context rep, IndexTable rep)
    -> Maybe (IndexTable rep)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Context rep, IndexTable rep) -> IndexTable rep
    forall a b. (a, b) -> b
    snd ((Context rep, IndexTable rep) -> Maybe (IndexTable rep))
    -> (Context rep, IndexTable rep) -> Maybe (IndexTable rep)
    forall a b. (a -> b) -> a -> b
    $
                      -- Should we use recContex instead of ctx''?
                      Context rep
    -> [VName]
    -> ArrayName
    -> [DimAccess rep]
    -> (Context rep, IndexTable rep)
    forall {k} (rep :: k).
    Context rep
    -> [VName]
    -> ArrayName
    -> [DimAccess rep]
    -> (Context rep, IndexTable rep)
    analyseIndex' Context rep
    ctx'' [VName]
    patterns ArrayName
    array_name [DimAccess rep]
    dim_indices
                )
              ([(VName, (ArrayName, [VName], [DimAccess rep]))]
     -> [IndexTable rep])
    -> [(VName, (ArrayName, [VName], [DimAccess rep]))]
    -> [IndexTable rep]
    forall a b. (a -> b) -> a -> b
    $ Map VName (ArrayName, [VName], [DimAccess rep])
    -> [(VName, (ArrayName, [VName], [DimAccess rep]))]
    forall k a. Map k a -> [(k, a)]
    M.toList Map VName (ArrayName, [VName], [DimAccess rep])
    slices_new

      -- 1. We do not want the returned context directly.
      --    however, we do want pat to map to the names what was hit in body.
      --    therefore we need to subtract the old context from the returned one,
      --    and discard all the keys within it.

      -- assignments :: M.Map VName (VariableInfo rep),
      let in_scope_dependencies_from_body :: Names
    in_scope_dependencies_from_body =
            Context rep -> Map VName (VariableInfo rep) -> Names
    forall {k} (rep :: k).
    Context rep -> Map VName (VariableInfo rep) -> Names
    rmOutOfScopeDeps Context rep
    ctx'' (Map VName (VariableInfo rep) -> Names)
    -> Map VName (VariableInfo rep) -> Names
    forall a b. (a -> b) -> a -> b
    $
              Map VName (VariableInfo rep)
    -> Map VName (VariableInfo rep) -> Map VName (VariableInfo rep)
    forall k a b. Ord k => Map k a -> Map k b -> Map k a
    M.difference (Context rep -> Map VName (VariableInfo rep)
    forall {k} (rep :: k). Context rep -> Map VName (VariableInfo rep)
    assignments Context rep
    ctx'') (Context rep -> Map VName (VariableInfo rep)
    forall {k} (rep :: k). Context rep -> Map VName (VariableInfo rep)
    assignments Context rep
    recContext)

      -- 2. We are ONLY interested in the rhs of assignments (ie. the
      --    dependencies of pat :) )
      let ctx' :: Context rep
    ctx' = (Context rep -> Context rep -> Context rep)
    -> Context rep -> [Context rep] -> Context rep
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl Context rep -> Context rep -> Context rep
    forall k (rep :: k). Context rep -> Context rep -> Context rep
    extend Context rep
    ctx ([Context rep] -> Context rep) -> [Context rep] -> Context rep
    forall a b. (a -> b) -> a -> b
    $ Names -> [Context rep]
    concatVariableInfo Names
    in_scope_dependencies_from_body -- . map snd $ M.toList varInfos
      -- 3. Now we have the correct context and result
      (Context rep
    ctx' {parents = parents ctx, currentLevel = currentLevel ctx, slices = slices ctx}, IndexTable rep
    slices_indices)
      where
        -- Extracts and merges `Names` in `VariableInfo`s, and makes a new VariableInfo. This
        -- MAY throw away needed information, but it was my best guess at a solution
        -- at the time of writing.
        concatVariableInfo :: Names -> [Context rep]
    concatVariableInfo Names
    dependencies =
          (VName -> Context rep) -> [VName] -> [Context rep]
    forall a b. (a -> b) -> [a] -> [b]
    map (\VName
    pat -> VName -> VariableInfo rep -> Context rep
    forall {k} (rep :: k). VName -> VariableInfo rep -> Context rep
    oneContext VName
    pat (Context rep -> Names -> VariableInfo rep
    forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
    varInfoFromNames Context rep
    ctx Names
    dependencies)) [VName]
    pats

        -- Context used for "recursion" into analyseStmsPrimitive
        recContext :: Context rep
    recContext =
          Context rep
    ctx
            { parents = parents ctx <> concatMap (\VName
    pat -> [VName -> BodyType
    body_constructor VName
    pat]) pats,
              currentLevel = currentLevel ctx + 1
            }

        -- Recursively looks up dependencies, until they're in scope or empty set.
        rmOutOfScopeDeps :: Context rep -> M.Map VName (VariableInfo rep) -> Names
        rmOutOfScopeDeps :: forall {k} (rep :: k).
    Context rep -> Map VName (VariableInfo rep) -> Names
    rmOutOfScopeDeps Context rep
    ctx' Map VName (VariableInfo rep)
    new_assignments =
          let throwaway_assignments :: Map VName (VariableInfo rep)
    throwaway_assignments = Context rep -> Map VName (VariableInfo rep)
    forall {k} (rep :: k). Context rep -> Map VName (VariableInfo rep)
    assignments Context rep
    ctx'
              local_assignments :: Map VName (VariableInfo rep)
    local_assignments = Context rep -> Map VName (VariableInfo rep)
    forall {k} (rep :: k). Context rep -> Map VName (VariableInfo rep)
    assignments Context rep
    ctx
              f :: Names -> VName -> VariableInfo rep -> Names
    f Names
    result VName
    a VariableInfo rep
    var_info =
                -- if the VName of the assignment exists in the context, we are good
                if VName
    a VName -> Map VName (VariableInfo rep) -> Bool
    forall k a. Ord k => k -> Map k a -> Bool
    `M.member` Map VName (VariableInfo rep)
    local_assignments
                  then Names
    result Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> VName -> Names
    oneName VName
    a
                  else -- Otherwise, recurse on its dependencies;
                  -- 0. Add dependencies in ctx to result
                    let ([VName]
    deps_in_ctx, [VName]
    deps_not_in_ctx) =
                          (VName -> Bool) -> [VName] -> ([VName], [VName])
    forall a. (a -> Bool) -> [a] -> ([a], [a])
    L.partition (VName -> Map VName (VariableInfo rep) -> Bool
    forall k a. Ord k => k -> Map k a -> Bool
    `M.member` Map VName (VariableInfo rep)
    local_assignments) ([VName] -> ([VName], [VName])) -> [VName] -> ([VName], [VName])
    forall a b. (a -> b) -> a -> b
    $
                            Names -> [VName]
    namesToList (VariableInfo rep -> Names
    forall {k} (rep :: k). VariableInfo rep -> Names
    deps VariableInfo rep
    var_info)
                        deps_not_in_ctx' :: Map VName (VariableInfo rep)
    deps_not_in_ctx' =
                          [(VName, VariableInfo rep)] -> Map VName (VariableInfo rep)
    forall k a. Ord k => [(k, a)] -> Map k a
    M.fromList ([(VName, VariableInfo rep)] -> Map VName (VariableInfo rep))
    -> [(VName, VariableInfo rep)] -> Map VName (VariableInfo rep)
    forall a b. (a -> b) -> a -> b
    $
                            (VName -> Maybe (VName, VariableInfo rep))
    -> [VName] -> [(VName, VariableInfo rep)]
    forall a b. (a -> Maybe b) -> [a] -> [b]
    mapMaybe
                              (\VName
    d -> (VName
    d,) (VariableInfo rep -> (VName, VariableInfo rep))
    -> Maybe (VariableInfo rep) -> Maybe (VName, VariableInfo rep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> VName -> Map VName (VariableInfo rep) -> Maybe (VariableInfo rep)
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    d Map VName (VariableInfo rep)
    throwaway_assignments)
                              [VName]
    deps_not_in_ctx
                     in Names
    result
                          Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> [VName] -> Names
    namesFromList [VName]
    deps_in_ctx
                          Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Context rep -> Map VName (VariableInfo rep) -> Names
    forall {k} (rep :: k).
    Context rep -> Map VName (VariableInfo rep) -> Names
    rmOutOfScopeDeps Context rep
    ctx' Map VName (VariableInfo rep)
    deps_not_in_ctx'
           in (Names -> VName -> VariableInfo rep -> Names)
    -> Names -> Map VName (VariableInfo rep) -> Names
    forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
    M.foldlWithKey Names -> VName -> VariableInfo rep -> Names
    forall {k} {rep :: k}. Names -> VName -> VariableInfo rep -> Names
    f Names
    forall a. Monoid a => a
    mempty Map VName (VariableInfo rep)
    new_assignments

    -- | Analyse a rep statement and return the updated context and array index
    -- descriptors.
    analyseStm :: (Analyse rep) => Context rep -> Stm rep -> (Context rep, IndexTable rep)
    analyseStm :: forall rep.
    Analyse rep =>
    Context rep -> Stm rep -> (Context rep, IndexTable rep)
    analyseStm Context rep
    ctx (Let Pat (LetDec rep)
    pats StmAux (ExpDec rep)
    _ Exp rep
    e) = do
      -- Get the name of the first element in a pattern
      let pattern_names :: [VName]
    pattern_names = (PatElem (LetDec rep) -> VName)
    -> [PatElem (LetDec rep)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map PatElem (LetDec rep) -> VName
    forall dec. PatElem dec -> VName
    patElemName ([PatElem (LetDec rep)] -> [VName])
    -> [PatElem (LetDec rep)] -> [VName]
    forall a b. (a -> b) -> a -> b
    $ Pat (LetDec rep) -> [PatElem (LetDec rep)]
    forall dec. Pat dec -> [PatElem dec]
    patElems Pat (LetDec rep)
    pats

      -- Construct the result and Context from the subexpression. If the
      -- subexpression is a body, we recurse into it.
      case Exp rep
    e of
        BasicOp (Index VName
    name (Slice [DimIndex SubExp]
    dim_subexp)) ->
          Context rep
    -> [VName]
    -> VName
    -> [DimIndex SubExp]
    -> (Context rep, IndexTable rep)
    forall {k} (rep :: k).
    Context rep
    -> [VName]
    -> VName
    -> [DimIndex SubExp]
    -> (Context rep, IndexTable rep)
    analyseIndex Context rep
    ctx [VName]
    pattern_names VName
    name [DimIndex SubExp]
    dim_subexp
        BasicOp (Update Safety
    _ VName
    name (Slice [DimIndex SubExp]
    dim_subexp) SubExp
    _subexp) ->
          Context rep
    -> [VName]
    -> VName
    -> [DimIndex SubExp]
    -> (Context rep, IndexTable rep)
    forall {k} (rep :: k).
    Context rep
    -> [VName]
    -> VName
    -> [DimIndex SubExp]
    -> (Context rep, IndexTable rep)
    analyseIndex Context rep
    ctx [VName]
    pattern_names VName
    name [DimIndex SubExp]
    dim_subexp
        BasicOp BasicOp
    op ->
          Context rep -> BasicOp -> [VName] -> (Context rep, IndexTable rep)
    forall {k} (rep :: k).
    Context rep -> BasicOp -> [VName] -> (Context rep, IndexTable rep)
    analyseBasicOp Context rep
    ctx BasicOp
    op [VName]
    pattern_names
        Match [SubExp]
    conds [Case (Body rep)]
    cases Body rep
    default_body MatchDec (BranchType rep)
    _ ->
          Context rep
    -> [VName]
    -> Body rep
    -> [Body rep]
    -> (Context rep, IndexTable rep)
    forall rep.
    Analyse rep =>
    Context rep
    -> [VName]
    -> Body rep
    -> [Body rep]
    -> (Context rep, IndexTable rep)
    analyseMatch Context rep
    ctx' [VName]
    pattern_names Body rep
    default_body ([Body rep] -> (Context rep, IndexTable rep))
    -> [Body rep] -> (Context rep, IndexTable rep)
    forall a b. (a -> b) -> a -> b
    $ (Case (Body rep) -> Body rep) -> [Case (Body rep)] -> [Body rep]
    forall a b. (a -> b) -> [a] -> [b]
    map Case (Body rep) -> Body rep
    forall body. Case body -> body
    caseBody [Case (Body rep)]
    cases
          where
            ctx' :: Context rep
    ctx' =
              Context rep -> VariableInfo rep -> [VName] -> Context rep
    forall {k} (rep :: k).
    Context rep -> VariableInfo rep -> [VName] -> Context rep
    contextFromNames Context rep
    ctx (Context rep -> VariableInfo rep
    forall {k} (rep :: k). Context rep -> VariableInfo rep
    varInfoZeroDeps Context rep
    ctx) ([VName] -> Context rep) -> [VName] -> Context rep
    forall a b. (a -> b) -> a -> b
    $
                (SubExp -> [VName]) -> [SubExp] -> [VName]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap (Names -> [VName]
    namesToList (Names -> [VName]) -> (SubExp -> Names) -> SubExp -> [VName]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SubExp -> Names
    forall a. FreeIn a => a -> Names
    freeIn) [SubExp]
    conds
        Loop [(FParam rep, SubExp)]
    bindings LoopForm
    loop Body rep
    body ->
          Context rep
    -> [(FParam rep, SubExp)]
    -> LoopForm
    -> Body rep
    -> [VName]
    -> (Context rep, IndexTable rep)
    forall rep.
    Analyse rep =>
    Context rep
    -> [(FParam rep, SubExp)]
    -> LoopForm
    -> Body rep
    -> [VName]
    -> (Context rep, IndexTable rep)
    analyseLoop Context rep
    ctx [(FParam rep, SubExp)]
    bindings LoopForm
    loop Body rep
    body [VName]
    pattern_names
        Apply Name
    _name [(SubExp, Diet)]
    diets [(RetType rep, RetAls)]
    _ Safety
    _ ->
          Context rep
    -> [VName] -> [(SubExp, Diet)] -> (Context rep, IndexTable rep)
    forall {k} (rep :: k).
    Context rep
    -> [VName] -> [(SubExp, Diet)] -> (Context rep, IndexTable rep)
    analyseApply Context rep
    ctx [VName]
    pattern_names [(SubExp, Diet)]
    diets
        WithAcc [WithAccInput rep]
    _ Lambda rep
    _ ->
          (Context rep
    ctx, IndexTable rep
    forall a. Monoid a => a
    mempty) -- ignored
        Op Op rep
    op ->
          Op rep -> Context rep -> [VName] -> (Context rep, IndexTable rep)
    forall rep.
    Analyse rep =>
    Op rep -> Context rep -> [VName] -> (Context rep, IndexTable rep)
    analyseOp Op rep
    op Context rep
    ctx [VName]
    pattern_names

    -- If left, this is just a regular index. If right, a slice happened.
    getIndexDependencies :: Context rep -> [DimIndex SubExp] -> Either [DimAccess rep] [DimAccess rep]
    getIndexDependencies :: forall {k} (rep :: k).
    Context rep
    -> [DimIndex SubExp] -> Either [DimAccess rep] [DimAccess rep]
    getIndexDependencies Context rep
    ctx [DimIndex SubExp]
    dims =
      (Either [DimAccess rep] [DimAccess rep], Int)
    -> Either [DimAccess rep] [DimAccess rep]
    forall a b. (a, b) -> a
    fst ((Either [DimAccess rep] [DimAccess rep], Int)
     -> Either [DimAccess rep] [DimAccess rep])
    -> (Either [DimAccess rep] [DimAccess rep], Int)
    -> Either [DimAccess rep] [DimAccess rep]
    forall a b. (a -> b) -> a -> b
    $
        (DimIndex SubExp
     -> (Either [DimAccess rep] [DimAccess rep], Int)
     -> (Either [DimAccess rep] [DimAccess rep], Int))
    -> (Either [DimAccess rep] [DimAccess rep], Int)
    -> [DimIndex SubExp]
    -> (Either [DimAccess rep] [DimAccess rep], Int)
    forall a b. (a -> b -> b) -> b -> [a] -> b
    forall (t :: * -> *) a b.
    Foldable t =>
    (a -> b -> b) -> b -> t a -> b
    foldr
          ( \DimIndex SubExp
    idx (Either [DimAccess rep] [DimAccess rep]
    a, Int
    i) ->
              ( ([DimAccess rep] -> Either [DimAccess rep] [DimAccess rep])
    -> ([DimAccess rep] -> Either [DimAccess rep] [DimAccess rep])
    -> Either [DimAccess rep] [DimAccess rep]
    -> Either [DimAccess rep] [DimAccess rep]
    forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
    either (DimIndex SubExp
    -> [DimAccess rep] -> Either [DimAccess rep] [DimAccess rep]
    matchDimIndex DimIndex SubExp
    idx) (([DimAccess rep] -> Either [DimAccess rep] [DimAccess rep])
    -> ([DimAccess rep] -> Either [DimAccess rep] [DimAccess rep])
    -> Either [DimAccess rep] [DimAccess rep]
    -> Either [DimAccess rep] [DimAccess rep]
    forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
    either [DimAccess rep] -> Either [DimAccess rep] [DimAccess rep]
    forall a b. b -> Either a b
    Right [DimAccess rep] -> Either [DimAccess rep] [DimAccess rep]
    forall a b. b -> Either a b
    Right (Either [DimAccess rep] [DimAccess rep]
     -> Either [DimAccess rep] [DimAccess rep])
    -> ([DimAccess rep] -> Either [DimAccess rep] [DimAccess rep])
    -> [DimAccess rep]
    -> Either [DimAccess rep] [DimAccess rep]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . DimIndex SubExp
    -> [DimAccess rep] -> Either [DimAccess rep] [DimAccess rep]
    matchDimIndex DimIndex SubExp
    idx) Either [DimAccess rep] [DimAccess rep]
    a,
                Int
    i Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    1
              )
          )
          ([DimAccess rep] -> Either [DimAccess rep] [DimAccess rep]
    forall a b. a -> Either a b
    Left [], [DimIndex SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [DimIndex SubExp]
    dims Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    1)
          [DimIndex SubExp]
    dims
      where
        matchDimIndex :: DimIndex SubExp
    -> [DimAccess rep] -> Either [DimAccess rep] [DimAccess rep]
    matchDimIndex (DimFix SubExp
    subExpression) [DimAccess rep]
    accumulator =
          [DimAccess rep] -> Either [DimAccess rep] [DimAccess rep]
    forall a b. a -> Either a b
    Left ([DimAccess rep] -> Either [DimAccess rep] [DimAccess rep])
    -> [DimAccess rep] -> Either [DimAccess rep] [DimAccess rep]
    forall a b. (a -> b) -> a -> b
    $ Context rep -> SubExp -> DimAccess rep
    forall {k} (rep :: k). Context rep -> SubExp -> DimAccess rep
    consolidate Context rep
    ctx SubExp
    subExpression DimAccess rep -> [DimAccess rep] -> [DimAccess rep]
    forall a. a -> [a] -> [a]
    : [DimAccess rep]
    accumulator
        -- If we encounter a DimSlice, add it to a map of `DimSlice`s and check
        -- result later.
        matchDimIndex (DimSlice SubExp
    offset SubExp
    num_elems SubExp
    stride) [DimAccess rep]
    accumulator =
          -- We assume that a slice is iterated sequentially, so we have to
          -- create a fake dependency for the slice.
          let dimAccess' :: DimAccess rep
    dimAccess' = Map VName Dependency -> Maybe VName -> DimAccess rep
    forall {k} (rep :: k).
    Map VName Dependency -> Maybe VName -> DimAccess rep
    DimAccess (VName -> Dependency -> Map VName Dependency
    forall k a. k -> a -> Map k a
    M.singleton (Name -> Int -> VName
    VName Name
    "slice" Int
    0) (Dependency -> Map VName Dependency)
    -> Dependency -> Map VName Dependency
    forall a b. (a -> b) -> a -> b
    $ Int -> VarType -> Dependency
    Dependency (Context rep -> Int
    forall {k} (rep :: k). Context rep -> Int
    currentLevel Context rep
    ctx) VarType
    LoopVar) (VName -> Maybe VName
    forall a. a -> Maybe a
    Just (VName -> Maybe VName) -> VName -> Maybe VName
    forall a b. (a -> b) -> a -> b
    $ Name -> Int -> VName
    VName Name
    "slice" Int
    0)
              cons :: SubExp -> DimAccess rep
    cons = Context rep -> SubExp -> DimAccess rep
    forall {k} (rep :: k). Context rep -> SubExp -> DimAccess rep
    consolidate Context rep
    ctx
              dimAccess :: DimAccess rep
    dimAccess = DimAccess rep
    forall k (rep :: k). DimAccess rep
    dimAccess' DimAccess rep -> DimAccess rep -> DimAccess rep
    forall a. Semigroup a => a -> a -> a
    <> SubExp -> DimAccess rep
    cons SubExp
    offset DimAccess rep -> DimAccess rep -> DimAccess rep
    forall a. Semigroup a => a -> a -> a
    <> SubExp -> DimAccess rep
    cons SubExp
    num_elems DimAccess rep -> DimAccess rep -> DimAccess rep
    forall a. Semigroup a => a -> a -> a
    <> SubExp -> DimAccess rep
    cons SubExp
    stride
           in [DimAccess rep] -> Either [DimAccess rep] [DimAccess rep]
    forall a b. b -> Either a b
    Right ([DimAccess rep] -> Either [DimAccess rep] [DimAccess rep])
    -> [DimAccess rep] -> Either [DimAccess rep] [DimAccess rep]
    forall a b. (a -> b) -> a -> b
    $ DimAccess rep
    dimAccess DimAccess rep -> [DimAccess rep] -> [DimAccess rep]
    forall a. a -> [a] -> [a]
    : [DimAccess rep]
    accumulator

    -- | Gets the dependencies of each dimension and either returns a result, or
    -- adds a slice to the context.
    analyseIndex :: Context rep -> [VName] -> VName -> [DimIndex SubExp] -> (Context rep, IndexTable rep)
    analyseIndex :: forall {k} (rep :: k).
    Context rep
    -> [VName]
    -> VName
    -> [DimIndex SubExp]
    -> (Context rep, IndexTable rep)
    analyseIndex Context rep
    ctx [VName]
    pats VName
    arr_name [DimIndex SubExp]
    dim_indices =
      -- Get the dependendencies of each dimension
      let dependencies :: Either [DimAccess rep] [DimAccess rep]
    dependencies = Context rep
    -> [DimIndex SubExp] -> Either [DimAccess rep] [DimAccess rep]
    forall {k} (rep :: k).
    Context rep
    -> [DimIndex SubExp] -> Either [DimAccess rep] [DimAccess rep]
    getIndexDependencies Context rep
    ctx [DimIndex SubExp]
    dim_indices
          -- Extend the current context with current pattern(s) and its deps
          ctx' :: Context rep
    ctx' = Context rep -> [DimIndex SubExp] -> [VName] -> Context rep
    forall {k} (rep :: k).
    Context rep -> [DimIndex SubExp] -> [VName] -> Context rep
    analyseIndexContextFromIndices Context rep
    ctx [DimIndex SubExp]
    dim_indices [VName]
    pats

          -- The bodytype(s) are used in the result construction
          array_name' :: ArrayName
    array_name' =
            -- For now, we assume the array is in row-major-order, hence the
            -- identity permutation. In the future, we might want to infer its
            -- layout, for example, if the array is the result of a transposition.
            let layout :: [Int]
    layout = [Int
    0 .. [DimIndex SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [DimIndex SubExp]
    dim_indices Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    1]
             in -- 2. If the arrayname was not in assignments, it was not an immediately
                --    allocated array.
                ArrayName -> Maybe ArrayName -> ArrayName
    forall a. a -> Maybe a -> a
    fromMaybe (VName
    arr_name, [], [Int]
    layout)
                  -- 1. Maybe find the array name, and the "stack" of body types that the
                  -- array was allocated in.
                  (Maybe ArrayName -> ArrayName)
    -> ([ArrayName] -> Maybe ArrayName) -> [ArrayName] -> ArrayName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (ArrayName -> Bool) -> [ArrayName] -> Maybe ArrayName
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
    L.find (\(VName
    n, [BodyType]
    _, [Int]
    _) -> VName
    n VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    == VName
    arr_name)
                  -- 0. Get the "stack" of bodytypes for each assignment
                  ([ArrayName] -> ArrayName) -> [ArrayName] -> ArrayName
    forall a b. (a -> b) -> a -> b
    $ ((VName, VariableInfo rep) -> ArrayName)
    -> [(VName, VariableInfo rep)] -> [ArrayName]
    forall a b. (a -> b) -> [a] -> [b]
    map (\(VName
    n, VariableInfo rep
    vi) -> (VName
    n, VariableInfo rep -> [BodyType]
    forall {k} (rep :: k). VariableInfo rep -> [BodyType]
    parents_nest VariableInfo rep
    vi, [Int]
    layout)) (Map VName (VariableInfo rep) -> [(VName, VariableInfo rep)]
    forall k a. Map k a -> [(k, a)]
    M.toList (Map VName (VariableInfo rep) -> [(VName, VariableInfo rep)])
    -> Map VName (VariableInfo rep) -> [(VName, VariableInfo rep)]
    forall a b. (a -> b) -> a -> b
    $ Context rep -> Map VName (VariableInfo rep)
    forall {k} (rep :: k). Context rep -> Map VName (VariableInfo rep)
    assignments Context rep
    ctx')
       in ([DimAccess rep] -> (Context rep, IndexTable rep))
    -> ([DimAccess rep] -> (Context rep, IndexTable rep))
    -> Either [DimAccess rep] [DimAccess rep]
    -> (Context rep, IndexTable rep)
    forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
    either (Context rep
    -> ArrayName -> [DimAccess rep] -> (Context rep, IndexTable rep)
    forall {k} (rep :: k).
    Context rep
    -> ArrayName -> [DimAccess rep] -> (Context rep, IndexTable rep)
    index Context rep
    ctx' ArrayName
    array_name') (Context rep
    -> ArrayName -> [DimAccess rep] -> (Context rep, IndexTable rep)
    forall {k} (rep :: k).
    Context rep
    -> ArrayName -> [DimAccess rep] -> (Context rep, IndexTable rep)
    slice Context rep
    ctx' ArrayName
    array_name') Either [DimAccess rep] [DimAccess rep]
    dependencies
      where
        slice :: Context rep -> ArrayName -> [DimAccess rep] -> (Context rep, IndexTable rep)
        slice :: forall {k} (rep :: k).
    Context rep
    -> ArrayName -> [DimAccess rep] -> (Context rep, IndexTable rep)
    slice Context rep
    context ArrayName
    array_name [DimAccess rep]
    dims =
          (Context rep
    context {slices = M.insert (head pats) (array_name, pats, dims) $ slices context}, IndexTable rep
    forall a. Monoid a => a
    mempty)

        index :: Context rep -> ArrayName -> [DimAccess rep] -> (Context rep, IndexTable rep)
        index :: forall {k} (rep :: k).
    Context rep
    -> ArrayName -> [DimAccess rep] -> (Context rep, IndexTable rep)
    index Context rep
    context array_name :: ArrayName
    array_name@(VName
    name, [BodyType]
    _, [Int]
    _) [DimAccess rep]
    dim_access =
          -- If the arrayname is a `DimSlice` we want to fixup the access
          case VName
    -> Map VName (ArrayName, [VName], [DimAccess rep])
    -> Maybe (ArrayName, [VName], [DimAccess rep])
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    name (Map VName (ArrayName, [VName], [DimAccess rep])
     -> Maybe (ArrayName, [VName], [DimAccess rep]))
    -> Map VName (ArrayName, [VName], [DimAccess rep])
    -> Maybe (ArrayName, [VName], [DimAccess rep])
    forall a b. (a -> b) -> a -> b
    $ Context rep -> Map VName (ArrayName, [VName], [DimAccess rep])
    forall {k} (rep :: k).
    Context rep -> Map VName (ArrayName, [VName], [DimAccess rep])
    slices Context rep
    context of
            Maybe (ArrayName, [VName], [DimAccess rep])
    Nothing -> Context rep
    -> [VName]
    -> ArrayName
    -> [DimAccess rep]
    -> (Context rep, IndexTable rep)
    forall {k} (rep :: k).
    Context rep
    -> [VName]
    -> ArrayName
    -> [DimAccess rep]
    -> (Context rep, IndexTable rep)
    analyseIndex' Context rep
    context [VName]
    pats ArrayName
    array_name [DimAccess rep]
    dim_access
            Just (ArrayName
    arr_name', [VName]
    pats', [DimAccess rep]
    slice_access) ->
              Context rep
    -> [VName]
    -> ArrayName
    -> [DimAccess rep]
    -> (Context rep, IndexTable rep)
    forall {k} (rep :: k).
    Context rep
    -> [VName]
    -> ArrayName
    -> [DimAccess rep]
    -> (Context rep, IndexTable rep)
    analyseIndex'
                Context rep
    context
                [VName]
    pats'
                ArrayName
    arr_name'
                ([DimAccess rep] -> [DimAccess rep]
    forall a. HasCallStack => [a] -> [a]
    init [DimAccess rep]
    slice_access [DimAccess rep] -> [DimAccess rep] -> [DimAccess rep]
    forall a. [a] -> [a] -> [a]
    ++ [[DimAccess rep] -> DimAccess rep
    forall a. HasCallStack => [a] -> a
    head [DimAccess rep]
    dim_access DimAccess rep -> DimAccess rep -> DimAccess rep
    forall a. Semigroup a => a -> a -> a
    <> [DimAccess rep] -> DimAccess rep
    forall a. HasCallStack => [a] -> a
    last [DimAccess rep]
    slice_access] [DimAccess rep] -> [DimAccess rep] -> [DimAccess rep]
    forall a. [a] -> [a] -> [a]
    ++ Int -> [DimAccess rep] -> [DimAccess rep]
    forall a. Int -> [a] -> [a]
    drop Int
    1 [DimAccess rep]
    dim_access)

    analyseIndexContextFromIndices :: Context rep -> [DimIndex SubExp] -> [VName] -> Context rep
    analyseIndexContextFromIndices :: forall {k} (rep :: k).
    Context rep -> [DimIndex SubExp] -> [VName] -> Context rep
    analyseIndexContextFromIndices Context rep
    ctx [DimIndex SubExp]
    dim_accesses [VName]
    pats =
      let subexprs :: [VName]
    subexprs =
            (DimIndex SubExp -> Maybe VName) -> [DimIndex SubExp] -> [VName]
    forall a b. (a -> Maybe b) -> [a] -> [b]
    mapMaybe
              ( \case
                  DimFix (Var VName
    v) -> VName -> Maybe VName
    forall a. a -> Maybe a
    Just VName
    v
                  DimFix (Constant PrimValue
    _) -> Maybe VName
    forall a. Maybe a
    Nothing
                  DimSlice SubExp
    _offs SubExp
    _n SubExp
    _stride -> Maybe VName
    forall a. Maybe a
    Nothing
              )
              [DimIndex SubExp]
    dim_accesses

          -- Add each non-constant DimIndex as a dependency to the index expression
          var_info :: VariableInfo rep
    var_info = Context rep -> Names -> VariableInfo rep
    forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
    varInfoFromNames Context rep
    ctx (Names -> VariableInfo rep) -> Names -> VariableInfo rep
    forall a b. (a -> b) -> a -> b
    $ [VName] -> Names
    namesFromList [VName]
    subexprs
       in -- Extend context with the dependencies index expression
          (Context rep -> Context rep -> Context rep)
    -> Context rep -> [Context rep] -> Context rep
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl' Context rep -> Context rep -> Context rep
    forall k (rep :: k). Context rep -> Context rep -> Context rep
    extend Context rep
    ctx ([Context rep] -> Context rep) -> [Context rep] -> Context rep
    forall a b. (a -> b) -> a -> b
    $ (VName -> Context rep) -> [VName] -> [Context rep]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName -> VariableInfo rep -> Context rep
    forall {k} (rep :: k). VName -> VariableInfo rep -> Context rep
    `oneContext` VariableInfo rep
    var_info) [VName]
    pats

    analyseIndex' ::
      Context rep ->
      [VName] ->
      ArrayName ->
      [DimAccess rep] ->
      (Context rep, IndexTable rep)
    analyseIndex' :: forall {k} (rep :: k).
    Context rep
    -> [VName]
    -> ArrayName
    -> [DimAccess rep]
    -> (Context rep, IndexTable rep)
    analyseIndex' Context rep
    ctx [VName]
    _ ArrayName
    _ [] = (Context rep
    ctx, IndexTable rep
    forall a. Monoid a => a
    mempty)
    analyseIndex' Context rep
    ctx [VName]
    _ ArrayName
    _ [DimAccess rep
    _] = (Context rep
    ctx, IndexTable rep
    forall a. Monoid a => a
    mempty)
    analyseIndex' Context rep
    ctx [VName]
    pats ArrayName
    arr_name [DimAccess rep]
    dim_accesses =
      -- Get the name of all segmaps in the current "callstack"
      let segmaps :: [SegOpName]
    segmaps = Context rep -> [SegOpName]
    forall {k} (rep :: k). Context rep -> [SegOpName]
    allSegMap Context rep
    ctx
          idx_expr_name :: [VName]
    idx_expr_name = [VName]
    pats --                                                IndexExprName
          -- For each pattern, create a mapping to the dimensional indices
          map_ixd_expr :: [Map VName [DimAccess rep]]
    map_ixd_expr = (VName -> Map VName [DimAccess rep])
    -> [VName] -> [Map VName [DimAccess rep]]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName -> [DimAccess rep] -> Map VName [DimAccess rep]
    forall k a. k -> a -> Map k a
    `M.singleton` [DimAccess rep]
    dim_accesses) [VName]
    idx_expr_name --       IndexExprName |-> [DimAccess]
          -- For each pattern -> [DimAccess] mapping, create a mapping from the array
          -- name that was indexed.
          map_array :: [Map ArrayName (Map VName [DimAccess rep])]
    map_array = (Map VName [DimAccess rep]
     -> Map ArrayName (Map VName [DimAccess rep]))
    -> [Map VName [DimAccess rep]]
    -> [Map ArrayName (Map VName [DimAccess rep])]
    forall a b. (a -> b) -> [a] -> [b]
    map (ArrayName
    -> Map VName [DimAccess rep]
    -> Map ArrayName (Map VName [DimAccess rep])
    forall k a. k -> a -> Map k a
    M.singleton ArrayName
    arr_name) [Map VName [DimAccess rep]]
    map_ixd_expr --   ArrayName |-> IndexExprName |-> [DimAccess]
          -- ∀ (arr_name -> IdxExp -> [DimAccess]) mappings, create a mapping from all
          -- segmaps in current callstack (segThread & segGroups alike).
          results :: [IndexTable rep]
    results = (Map ArrayName (Map VName [DimAccess rep]) -> [IndexTable rep])
    -> [Map ArrayName (Map VName [DimAccess rep])] -> [IndexTable rep]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap (\Map ArrayName (Map VName [DimAccess rep])
    ma -> (SegOpName -> IndexTable rep) -> [SegOpName] -> [IndexTable rep]
    forall a b. (a -> b) -> [a] -> [b]
    map (SegOpName
    -> Map ArrayName (Map VName [DimAccess rep]) -> IndexTable rep
    forall k a. k -> a -> Map k a
    `M.singleton` Map ArrayName (Map VName [DimAccess rep])
    ma) [SegOpName]
    segmaps) [Map ArrayName (Map VName [DimAccess rep])]
    map_array

          res :: IndexTable rep
    res = (IndexTable rep -> IndexTable rep -> IndexTable rep)
    -> IndexTable rep -> [IndexTable rep] -> IndexTable rep
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl' IndexTable rep -> IndexTable rep -> IndexTable rep
    forall {k} (rep :: k).
    IndexTable rep -> IndexTable rep -> IndexTable rep
    unionIndexTables IndexTable rep
    forall a. Monoid a => a
    mempty [IndexTable rep]
    results
       in (Context rep
    ctx, IndexTable rep
    res)

    analyseBasicOp :: Context rep -> BasicOp -> [VName] -> (Context rep, IndexTable rep)
    analyseBasicOp :: forall {k} (rep :: k).
    Context rep -> BasicOp -> [VName] -> (Context rep, IndexTable rep)
    analyseBasicOp Context rep
    ctx BasicOp
    expression [VName]
    pats =
      -- Construct a VariableInfo from the subexpressions
      let ctx_val :: VariableInfo rep
    ctx_val = case BasicOp
    expression of
            SubExp SubExp
    se -> SubExp -> VariableInfo rep
    forall {k} {rep :: k}. SubExp -> VariableInfo rep
    varInfoFromSubExp SubExp
    se
            Opaque OpaqueOp
    _ SubExp
    se -> SubExp -> VariableInfo rep
    forall {k} {rep :: k}. SubExp -> VariableInfo rep
    varInfoFromSubExp SubExp
    se
            ArrayVal [PrimValue]
    _ PrimType
    _ -> (Context rep -> Names -> VariableInfo rep
    forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
    varInfoFromNames Context rep
    ctx Names
    forall a. Monoid a => a
    mempty) {variableType = ConstType}
            ArrayLit [SubExp]
    ses Type
    _t -> Names -> [SubExp] -> VariableInfo rep
    concatVariableInfos Names
    forall a. Monoid a => a
    mempty [SubExp]
    ses
            UnOp UnOp
    _ SubExp
    se -> SubExp -> VariableInfo rep
    forall {k} {rep :: k}. SubExp -> VariableInfo rep
    varInfoFromSubExp SubExp
    se
            BinOp BinOp
    _ SubExp
    lsubexp SubExp
    rsubexp -> Names -> [SubExp] -> VariableInfo rep
    concatVariableInfos Names
    forall a. Monoid a => a
    mempty [SubExp
    lsubexp, SubExp
    rsubexp]
            CmpOp CmpOp
    _ SubExp
    lsubexp SubExp
    rsubexp -> Names -> [SubExp] -> VariableInfo rep
    concatVariableInfos Names
    forall a. Monoid a => a
    mempty [SubExp
    lsubexp, SubExp
    rsubexp]
            ConvOp ConvOp
    _ SubExp
    se -> SubExp -> VariableInfo rep
    forall {k} {rep :: k}. SubExp -> VariableInfo rep
    varInfoFromSubExp SubExp
    se
            Assert SubExp
    se ErrorMsg SubExp
    _ -> SubExp -> VariableInfo rep
    forall {k} {rep :: k}. SubExp -> VariableInfo rep
    varInfoFromSubExp SubExp
    se
            Index VName
    name Slice SubExp
    _ ->
              [Char] -> VariableInfo rep
    forall a. HasCallStack => [Char] -> a
    error ([Char] -> VariableInfo rep) -> [Char] -> VariableInfo rep
    forall a b. (a -> b) -> a -> b
    $ [Char]
    "unhandled: Index (This should NEVER happen) into " [Char] -> ShowS
    forall a. [a] -> [a] -> [a]
    ++ VName -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString VName
    name
            Update Safety
    _ VName
    name Slice SubExp
    _slice SubExp
    _subexp ->
              [Char] -> VariableInfo rep
    forall a. HasCallStack => [Char] -> a
    error ([Char] -> VariableInfo rep) -> [Char] -> VariableInfo rep
    forall a b. (a -> b) -> a -> b
    $ [Char]
    "unhandled: Update (This should NEVER happen) onto " [Char] -> ShowS
    forall a. [a] -> [a] -> [a]
    ++ VName -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString VName
    name
            -- Technically, do we need this case?
            Concat Int
    _ NonEmpty VName
    _ SubExp
    length_subexp -> SubExp -> VariableInfo rep
    forall {k} {rep :: k}. SubExp -> VariableInfo rep
    varInfoFromSubExp SubExp
    length_subexp
            Manifest VName
    name [Int]
    _dim -> Context rep -> Names -> VariableInfo rep
    forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
    varInfoFromNames Context rep
    ctx (Names -> VariableInfo rep) -> Names -> VariableInfo rep
    forall a b. (a -> b) -> a -> b
    $ VName -> Names
    oneName VName
    name
            Iota SubExp
    end SubExp
    start SubExp
    stride IntType
    _ -> Names -> [SubExp] -> VariableInfo rep
    concatVariableInfos Names
    forall a. Monoid a => a
    mempty [SubExp
    end, SubExp
    start, SubExp
    stride]
            Replicate (Shape [SubExp]
    shape) SubExp
    value' -> Names -> [SubExp] -> VariableInfo rep
    concatVariableInfos Names
    forall a. Monoid a => a
    mempty (SubExp
    value' SubExp -> [SubExp] -> [SubExp]
    forall a. a -> [a] -> [a]
    : [SubExp]
    shape)
            Scratch PrimType
    _ [SubExp]
    sers -> Names -> [SubExp] -> VariableInfo rep
    concatVariableInfos Names
    forall a. Monoid a => a
    mempty [SubExp]
    sers
            Reshape VName
    name NewShape SubExp
    newshape ->
              Names -> [SubExp] -> VariableInfo rep
    concatVariableInfos (VName -> Names
    oneName VName
    name) (ShapeBase SubExp -> [SubExp]
    forall d. ShapeBase d -> [d]
    shapeDims (NewShape SubExp -> ShapeBase SubExp
    forall d. NewShape d -> ShapeBase d
    newShape NewShape SubExp
    newshape))
            Rearrange VName
    name [Int]
    _ -> Context rep -> Names -> VariableInfo rep
    forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
    varInfoFromNames Context rep
    ctx (Names -> VariableInfo rep) -> Names -> VariableInfo rep
    forall a b. (a -> b) -> a -> b
    $ VName -> Names
    oneName VName
    name
            UpdateAcc Safety
    _ VName
    name [SubExp]
    lsubexprs [SubExp]
    rsubexprs ->
              Names -> [SubExp] -> VariableInfo rep
    concatVariableInfos (VName -> Names
    oneName VName
    name) ([SubExp]
    lsubexprs [SubExp] -> [SubExp] -> [SubExp]
    forall a. [a] -> [a] -> [a]
    ++ [SubExp]
    rsubexprs)
            FlatIndex VName
    name FlatSlice SubExp
    _ -> Context rep -> Names -> VariableInfo rep
    forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
    varInfoFromNames Context rep
    ctx (Names -> VariableInfo rep) -> Names -> VariableInfo rep
    forall a b. (a -> b) -> a -> b
    $ VName -> Names
    oneName VName
    name
            FlatUpdate VName
    name FlatSlice SubExp
    _ VName
    source -> Context rep -> Names -> VariableInfo rep
    forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
    varInfoFromNames Context rep
    ctx (Names -> VariableInfo rep) -> Names -> VariableInfo rep
    forall a b. (a -> b) -> a -> b
    $ [VName] -> Names
    namesFromList [VName
    name, VName
    source]
          ctx' :: Context rep
    ctx' = (Context rep -> Context rep -> Context rep)
    -> Context rep -> [Context rep] -> Context rep
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl' Context rep -> Context rep -> Context rep
    forall k (rep :: k). Context rep -> Context rep -> Context rep
    extend Context rep
    ctx ([Context rep] -> Context rep) -> [Context rep] -> Context rep
    forall a b. (a -> b) -> a -> b
    $ (VName -> Context rep) -> [VName] -> [Context rep]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName -> VariableInfo rep -> Context rep
    forall {k} (rep :: k). VName -> VariableInfo rep -> Context rep
    `oneContext` VariableInfo rep
    ctx_val) [VName]
    pats
       in (Context rep
    ctx', IndexTable rep
    forall a. Monoid a => a
    mempty)
      where
        concatVariableInfos :: Names -> [SubExp] -> VariableInfo rep
    concatVariableInfos Names
    ne [SubExp]
    nn =
          Context rep -> Names -> VariableInfo rep
    forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
    varInfoFromNames Context rep
    ctx (Names
    ne Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> [Names] -> Names
    forall a. Monoid a => [a] -> a
    mconcat ((SubExp -> Names) -> [SubExp] -> [Names]
    forall a b. (a -> b) -> [a] -> [b]
    map ([VName] -> Context rep -> SubExp -> Names
    forall {k} (rep :: k). [VName] -> Context rep -> SubExp -> Names
    analyseSubExp [VName]
    pats Context rep
    ctx) [SubExp]
    nn))

        varInfoFromSubExp :: SubExp -> VariableInfo rep
    varInfoFromSubExp (Constant PrimValue
    _) = (Context rep -> Names -> VariableInfo rep
    forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
    varInfoFromNames Context rep
    ctx Names
    forall a. Monoid a => a
    mempty) {variableType = ConstType}
        varInfoFromSubExp (Var VName
    v) =
          case VName -> Map VName (VariableInfo rep) -> Maybe (VariableInfo rep)
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    v (Context rep -> Map VName (VariableInfo rep)
    forall {k} (rep :: k). Context rep -> Map VName (VariableInfo rep)
    assignments Context rep
    ctx) of
            Just VariableInfo rep
    _ -> (Context rep -> Names -> VariableInfo rep
    forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
    varInfoFromNames Context rep
    ctx (Names -> VariableInfo rep) -> Names -> VariableInfo rep
    forall a b. (a -> b) -> a -> b
    $ VName -> Names
    oneName VName
    v) {variableType = Variable}
            Maybe (VariableInfo rep)
    Nothing -> (Context rep -> Names -> VariableInfo rep
    forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
    varInfoFromNames Context rep
    ctx Names
    forall a. Monoid a => a
    mempty) {variableType = Variable} -- Means a global.

    analyseMatch :: (Analyse rep) => Context rep -> [VName] -> Body rep -> [Body rep] -> (Context rep, IndexTable rep)
    analyseMatch :: forall rep.
    Analyse rep =>
    Context rep
    -> [VName]
    -> Body rep
    -> [Body rep]
    -> (Context rep, IndexTable rep)
    analyseMatch Context rep
    ctx [VName]
    pats Body rep
    body [Body rep]
    parents =
      let ctx'' :: Context rep
    ctx'' = Context rep
    ctx {currentLevel = currentLevel ctx - 1}
       in ((Context rep, IndexTable rep)
     -> Body rep -> (Context rep, IndexTable rep))
    -> (Context rep, IndexTable rep)
    -> [Body rep]
    -> (Context rep, IndexTable rep)
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl
            ( \(Context rep
    ctx', IndexTable rep
    res) Body rep
    b ->
                -- This Little Maneuver's Gonna Cost Us 51 Years
                (Context rep -> Context rep)
    -> (IndexTable rep -> IndexTable rep)
    -> (Context rep, IndexTable rep)
    -> (Context rep, IndexTable rep)
    forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
    forall (p :: * -> * -> *) a b c d.
    Bifunctor p =>
    (a -> b) -> (c -> d) -> p a c -> p b d
    bimap Context rep -> Context rep
    forall {k} {rep :: k}. Context rep -> Context rep
    constLevel (IndexTable rep -> IndexTable rep -> IndexTable rep
    forall {k} (rep :: k).
    IndexTable rep -> IndexTable rep -> IndexTable rep
    unionIndexTables IndexTable rep
    res)
                  ((Context rep, IndexTable rep) -> (Context rep, IndexTable rep))
    -> (Stms rep -> (Context rep, IndexTable rep))
    -> Stms rep
    -> (Context rep, IndexTable rep)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Context rep
    -> (VName -> BodyType)
    -> [VName]
    -> [Stm rep]
    -> (Context rep, IndexTable rep)
    forall rep.
    Analyse rep =>
    Context rep
    -> (VName -> BodyType)
    -> [VName]
    -> [Stm rep]
    -> (Context rep, IndexTable rep)
    analyseStms Context rep
    ctx' VName -> BodyType
    CondBodyName [VName]
    pats
                  ([Stm rep] -> (Context rep, IndexTable rep))
    -> (Stms rep -> [Stm rep])
    -> Stms rep
    -> (Context rep, IndexTable rep)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Stms rep -> [Stm rep]
    forall rep. Stms rep -> [Stm rep]
    stmsToList
                  (Stms rep -> (Context rep, IndexTable rep))
    -> Stms rep -> (Context rep, IndexTable rep)
    forall a b. (a -> b) -> a -> b
    $ Body rep -> Stms rep
    forall rep. Body rep -> Stms rep
    bodyStms Body rep
    b
            )
            (Context rep
    ctx'', IndexTable rep
    forall a. Monoid a => a
    mempty)
            (Body rep
    body Body rep -> [Body rep] -> [Body rep]
    forall a. a -> [a] -> [a]
    : [Body rep]
    parents)
      where
        constLevel :: Context rep -> Context rep
    constLevel Context rep
    context = Context rep
    context {currentLevel = currentLevel ctx - 1}

    analyseLoop :: (Analyse rep) => Context rep -> [(FParam rep, SubExp)] -> LoopForm -> Body rep -> [VName] -> (Context rep, IndexTable rep)
    analyseLoop :: forall rep.
    Analyse rep =>
    Context rep
    -> [(FParam rep, SubExp)]
    -> LoopForm
    -> Body rep
    -> [VName]
    -> (Context rep, IndexTable rep)
    analyseLoop Context rep
    ctx [(FParam rep, SubExp)]
    bindings LoopForm
    loop Body rep
    body [VName]
    pats = do
      let next_level :: Int
    next_level = Context rep -> Int
    forall {k} (rep :: k). Context rep -> Int
    currentLevel Context rep
    ctx
      let ctx'' :: Context rep
    ctx'' = Context rep
    ctx {currentLevel = next_level}
      let ctx' :: Context rep
    ctx' =
            Context rep -> VariableInfo rep -> [VName] -> Context rep
    forall {k} (rep :: k).
    Context rep -> VariableInfo rep -> [VName] -> Context rep
    contextFromNames Context rep
    ctx'' ((Context rep -> VariableInfo rep
    forall {k} (rep :: k). Context rep -> VariableInfo rep
    varInfoZeroDeps Context rep
    ctx) {variableType = LoopVar}) ([VName] -> Context rep) -> [VName] -> Context rep
    forall a b. (a -> b) -> a -> b
    $
              case LoopForm
    loop of
                WhileLoop VName
    iv -> VName
    iv VName -> [VName] -> [VName]
    forall a. a -> [a] -> [a]
    : ((FParam rep, SubExp) -> VName)
    -> [(FParam rep, SubExp)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (FParam rep -> VName
    forall dec. Param dec -> VName
    paramName (FParam rep -> VName)
    -> ((FParam rep, SubExp) -> FParam rep)
    -> (FParam rep, SubExp)
    -> VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (FParam rep, SubExp) -> FParam rep
    forall a b. (a, b) -> a
    fst) [(FParam rep, SubExp)]
    bindings
                ForLoop VName
    iv IntType
    _ SubExp
    _ -> VName
    iv VName -> [VName] -> [VName]
    forall a. a -> [a] -> [a]
    : ((FParam rep, SubExp) -> VName)
    -> [(FParam rep, SubExp)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (FParam rep -> VName
    forall dec. Param dec -> VName
    paramName (FParam rep -> VName)
    -> ((FParam rep, SubExp) -> FParam rep)
    -> (FParam rep, SubExp)
    -> VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (FParam rep, SubExp) -> FParam rep
    forall a b. (a, b) -> a
    fst) [(FParam rep, SubExp)]
    bindings

      -- Extend context with the loop expression
      Context rep
    -> (VName -> BodyType)
    -> [VName]
    -> [Stm rep]
    -> (Context rep, IndexTable rep)
    forall rep.
    Analyse rep =>
    Context rep
    -> (VName -> BodyType)
    -> [VName]
    -> [Stm rep]
    -> (Context rep, IndexTable rep)
    analyseStms Context rep
    ctx' VName -> BodyType
    LoopBodyName [VName]
    pats ([Stm rep] -> (Context rep, IndexTable rep))
    -> [Stm rep] -> (Context rep, IndexTable rep)
    forall a b. (a -> b) -> a -> b
    $ Stms rep -> [Stm rep]
    forall rep. Stms rep -> [Stm rep]
    stmsToList (Stms rep -> [Stm rep]) -> Stms rep -> [Stm rep]
    forall a b. (a -> b) -> a -> b
    $ Body rep -> Stms rep
    forall rep. Body rep -> Stms rep
    bodyStms Body rep
    body

    analyseApply :: Context rep -> [VName] -> [(SubExp, Diet)] -> (Context rep, IndexTable rep)
    analyseApply :: forall {k} (rep :: k).
    Context rep
    -> [VName] -> [(SubExp, Diet)] -> (Context rep, IndexTable rep)
    analyseApply Context rep
    ctx [VName]
    pats [(SubExp, Diet)]
    diets =
      ( (Context rep -> Context rep -> Context rep)
    -> Context rep -> [Context rep] -> Context rep
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl' Context rep -> Context rep -> Context rep
    forall k (rep :: k). Context rep -> Context rep -> Context rep
    extend Context rep
    ctx ([Context rep] -> Context rep) -> [Context rep] -> Context rep
    forall a b. (a -> b) -> a -> b
    $ (VName -> Context rep) -> [VName] -> [Context rep]
    forall a b. (a -> b) -> [a] -> [b]
    map (\VName
    pat -> VName -> VariableInfo rep -> Context rep
    forall {k} (rep :: k). VName -> VariableInfo rep -> Context rep
    oneContext VName
    pat (VariableInfo rep -> Context rep)
    -> VariableInfo rep -> Context rep
    forall a b. (a -> b) -> a -> b
    $ Context rep -> Names -> VariableInfo rep
    forall {k} (rep :: k). Context rep -> Names -> VariableInfo rep
    varInfoFromNames Context rep
    ctx (Names -> VariableInfo rep) -> Names -> VariableInfo rep
    forall a b. (a -> b) -> a -> b
    $ [Names] -> Names
    forall a. Monoid a => [a] -> a
    mconcat ([Names] -> Names) -> [Names] -> Names
    forall a b. (a -> b) -> a -> b
    $ ((SubExp, Diet) -> Names) -> [(SubExp, Diet)] -> [Names]
    forall a b. (a -> b) -> [a] -> [b]
    map (SubExp -> Names
    forall a. FreeIn a => a -> Names
    freeIn (SubExp -> Names)
    -> ((SubExp, Diet) -> SubExp) -> (SubExp, Diet) -> Names
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (SubExp, Diet) -> SubExp
    forall a b. (a, b) -> a
    fst) [(SubExp, Diet)]
    diets) [VName]
    pats,
        IndexTable rep
    forall a. Monoid a => a
    mempty
      )

    segOpType :: SegOp lvl rep -> VName -> SegOpName
    segOpType :: forall lvl rep. SegOp lvl rep -> VName -> SegOpName
    segOpType (SegMap {}) = VName -> SegOpName
    SegmentedMap
    segOpType (SegRed {}) = VName -> SegOpName
    SegmentedRed
    segOpType (SegScan {}) = VName -> SegOpName
    SegmentedScan
    segOpType (SegHist {}) = VName -> SegOpName
    SegmentedHist

    analyseSegOp :: (Analyse rep) => SegOp lvl rep -> Context rep -> [VName] -> (Context rep, IndexTable rep)
    analyseSegOp :: forall rep lvl.
    Analyse rep =>
    SegOp lvl rep
    -> Context rep -> [VName] -> (Context rep, IndexTable rep)
    analyseSegOp SegOp lvl rep
    op Context rep
    ctx [VName]
    pats =
      let next_level :: Int
    next_level = Context rep -> Int
    forall {k} (rep :: k). Context rep -> Int
    currentLevel Context rep
    ctx Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + [(VName, SubExp)] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length (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) Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    1
          ctx' :: Context rep
    ctx' = Context rep
    ctx {currentLevel = next_level}
          segspace_context :: Context rep
    segspace_context =
            (Context rep -> Context rep -> Context rep)
    -> Context rep -> [Context rep] -> Context rep
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl' Context rep -> Context rep -> Context rep
    forall k (rep :: k). Context rep -> Context rep -> Context rep
    extend Context rep
    ctx'
              ([Context rep] -> Context rep)
    -> (SegSpace -> [Context rep]) -> SegSpace -> Context rep
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ((VName, Int) -> Context rep) -> [(VName, Int)] -> [Context rep]
    forall a b. (a -> b) -> [a] -> [b]
    map (\(VName
    n, Int
    i) -> VName -> VariableInfo rep -> Context rep
    forall {k} (rep :: k). VName -> VariableInfo rep -> Context rep
    oneContext VName
    n (VariableInfo rep -> Context rep)
    -> VariableInfo rep -> Context rep
    forall a b. (a -> b) -> a -> b
    $ Names -> Int -> [BodyType] -> VarType -> VariableInfo rep
    forall {k} (rep :: k).
    Names -> Int -> [BodyType] -> VarType -> VariableInfo rep
    VariableInfo Names
    forall a. Monoid a => a
    mempty (Context rep -> Int
    forall {k} (rep :: k). Context rep -> Int
    currentLevel Context rep
    ctx Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    i) (Context rep -> [BodyType]
    forall {k} (rep :: k). Context rep -> [BodyType]
    parents Context rep
    ctx') VarType
    ThreadID)
              ([(VName, Int)] -> [Context rep])
    -> (SegSpace -> [(VName, Int)]) -> SegSpace -> [Context rep]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (\[VName]
    segspace_params -> [VName] -> [Int] -> [(VName, Int)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [VName]
    segspace_params [Int
    0 ..])
              -- contextFromNames ctx' Parallel
              ([VName] -> [(VName, Int)])
    -> (SegSpace -> [VName]) -> SegSpace -> [(VName, Int)]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ((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])
    -> (SegSpace -> [(VName, SubExp)]) -> SegSpace -> [VName]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SegSpace -> [(VName, SubExp)]
    unSegSpace
              (SegSpace -> Context rep) -> SegSpace -> Context rep
    forall a b. (a -> b) -> a -> b
    $ SegOp lvl rep -> SegSpace
    forall lvl rep. SegOp lvl rep -> SegSpace
    segSpace SegOp lvl rep
    op
       in -- Analyse statements in the SegOp body
          Context rep
    -> (VName -> BodyType)
    -> [VName]
    -> [Stm rep]
    -> (Context rep, IndexTable rep)
    forall rep.
    Analyse rep =>
    Context rep
    -> (VName -> BodyType)
    -> [VName]
    -> [Stm rep]
    -> (Context rep, IndexTable rep)
    analyseStms Context rep
    segspace_context (SegOpName -> BodyType
    SegOpName (SegOpName -> BodyType)
    -> (VName -> SegOpName) -> VName -> BodyType
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SegOp lvl rep -> VName -> SegOpName
    forall lvl rep. SegOp lvl rep -> VName -> SegOpName
    segOpType SegOp lvl rep
    op) [VName]
    pats ([Stm rep] -> (Context rep, IndexTable rep))
    -> (KernelBody rep -> [Stm rep])
    -> KernelBody rep
    -> (Context rep, IndexTable rep)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Stms rep -> [Stm rep]
    forall rep. Stms rep -> [Stm rep]
    stmsToList (Stms rep -> [Stm rep])
    -> (KernelBody rep -> Stms rep) -> KernelBody rep -> [Stm rep]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . KernelBody rep -> Stms rep
    forall rep. KernelBody rep -> Stms rep
    kernelBodyStms (KernelBody rep -> (Context rep, IndexTable rep))
    -> KernelBody rep -> (Context rep, IndexTable rep)
    forall a b. (a -> b) -> a -> b
    $ SegOp lvl rep -> KernelBody rep
    forall lvl rep. SegOp lvl rep -> KernelBody rep
    segBody SegOp lvl rep
    op

    analyseSizeOp :: SizeOp -> Context rep -> [VName] -> (Context rep, IndexTable rep)
    analyseSizeOp :: forall {k} (rep :: k).
    SizeOp -> Context rep -> [VName] -> (Context rep, IndexTable rep)
    analyseSizeOp SizeOp
    op Context rep
    ctx [VName]
    pats =
      let ctx' :: Context rep
    ctx' = case SizeOp
    op of
            CmpSizeLe Name
    _name SizeClass
    _class SubExp
    subexp -> [SubExp] -> Context rep
    subexprsToContext [SubExp
    subexp]
            CalcNumBlocks SubExp
    lsubexp Name
    _name SubExp
    rsubexp -> [SubExp] -> Context rep
    subexprsToContext [SubExp
    lsubexp, SubExp
    rsubexp]
            SizeOp
    _ -> Context rep
    ctx
          -- Add sizeOp to context
          ctx'' :: Context rep
    ctx'' =
            (Context rep -> Context rep -> Context rep)
    -> Context rep -> [Context rep] -> Context rep
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl' Context rep -> Context rep -> Context rep
    forall k (rep :: k). Context rep -> Context rep -> Context rep
    extend Context rep
    ctx' ([Context rep] -> Context rep) -> [Context rep] -> Context rep
    forall a b. (a -> b) -> a -> b
    $
              (VName -> Context rep) -> [VName] -> [Context rep]
    forall a b. (a -> b) -> [a] -> [b]
    map
                (\VName
    pat -> VName -> VariableInfo rep -> Context rep
    forall {k} (rep :: k). VName -> VariableInfo rep -> Context rep
    oneContext VName
    pat (VariableInfo rep -> Context rep)
    -> VariableInfo rep -> Context rep
    forall a b. (a -> b) -> a -> b
    $ (Context rep -> VariableInfo rep
    forall {k} (rep :: k). Context rep -> VariableInfo rep
    varInfoZeroDeps Context rep
    ctx) {parents_nest = parents ctx'})
                [VName]
    pats
       in (Context rep
    ctx'', IndexTable rep
    forall a. Monoid a => a
    mempty)
      where
        subexprsToContext :: [SubExp] -> Context rep
    subexprsToContext =
          Context rep -> VariableInfo rep -> [VName] -> Context rep
    forall {k} (rep :: k).
    Context rep -> VariableInfo rep -> [VName] -> Context rep
    contextFromNames Context rep
    ctx (Context rep -> VariableInfo rep
    forall {k} (rep :: k). Context rep -> VariableInfo rep
    varInfoZeroDeps Context rep
    ctx)
            ([VName] -> Context rep)
    -> ([SubExp] -> [VName]) -> [SubExp] -> Context rep
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (SubExp -> [VName]) -> [SubExp] -> [VName]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap (Names -> [VName]
    namesToList (Names -> [VName]) -> (SubExp -> Names) -> SubExp -> [VName]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [VName] -> Context rep -> SubExp -> Names
    forall {k} (rep :: k). [VName] -> Context rep -> SubExp -> Names
    analyseSubExp [VName]
    pats Context rep
    ctx)

    -- | Analyse statements in a rep body.
    analyseGPUBody :: (Analyse rep) => Body rep -> Context rep -> (Context rep, IndexTable rep)
    analyseGPUBody :: forall rep.
    Analyse rep =>
    Body rep -> Context rep -> (Context rep, IndexTable rep)
    analyseGPUBody Body rep
    body Context rep
    ctx =
      Context rep -> [Stm rep] -> (Context rep, IndexTable rep)
    forall rep.
    Analyse rep =>
    Context rep -> [Stm rep] -> (Context rep, IndexTable rep)
    analyseStmsPrimitive Context rep
    ctx ([Stm rep] -> (Context rep, IndexTable rep))
    -> [Stm rep] -> (Context rep, IndexTable rep)
    forall a b. (a -> b) -> a -> b
    $ Stms rep -> [Stm rep]
    forall rep. Stms rep -> [Stm rep]
    stmsToList (Stms rep -> [Stm rep]) -> Stms rep -> [Stm rep]
    forall a b. (a -> b) -> a -> b
    $ Body rep -> Stms rep
    forall rep. Body rep -> Stms rep
    bodyStms Body rep
    body

    analyseOtherOp :: Context rep -> [VName] -> (Context rep, IndexTable rep)
    analyseOtherOp :: forall {k} (rep :: k).
    Context rep -> [VName] -> (Context rep, IndexTable rep)
    analyseOtherOp Context rep
    ctx [VName]
    _ = (Context rep
    ctx, IndexTable rep
    forall a. Monoid a => a
    mempty)

    -- | Returns an intmap of names, to be used as dependencies in construction of
    -- VariableInfos.
    analyseSubExp :: [VName] -> Context rep -> SubExp -> Names
    analyseSubExp :: forall {k} (rep :: k). [VName] -> Context rep -> SubExp -> Names
    analyseSubExp [VName]
    _ Context rep
    _ (Constant PrimValue
    _) = Names
    forall a. Monoid a => a
    mempty
    analyseSubExp [VName]
    _ Context rep
    _ (Var VName
    v) = VName -> Names
    oneName VName
    v

    -- | Reduce a DimFix into its set of dependencies
    consolidate :: Context rep -> SubExp -> DimAccess rep
    consolidate :: forall {k} (rep :: k). Context rep -> SubExp -> DimAccess rep
    consolidate Context rep
    _ (Constant PrimValue
    _) = DimAccess rep
    forall a. Monoid a => a
    mempty
    consolidate Context rep
    ctx (Var VName
    v) = Map VName Dependency -> Maybe VName -> DimAccess rep
    forall {k} (rep :: k).
    Map VName Dependency -> Maybe VName -> DimAccess rep
    DimAccess (Context rep -> VName -> Map VName Dependency
    forall {k} (rep :: k). Context rep -> VName -> Map VName Dependency
    reduceDependencies Context rep
    ctx VName
    v) (VName -> Maybe VName
    forall a. a -> Maybe a
    Just VName
    v)

    -- | Recursively lookup vnames until vars with no deps are reached.
    reduceDependencies :: Context rep -> VName -> M.Map VName Dependency
    reduceDependencies :: forall {k} (rep :: k). Context rep -> VName -> Map VName Dependency
    reduceDependencies Context rep
    ctx VName
    v =
      case VName -> Map VName (VariableInfo rep) -> Maybe (VariableInfo rep)
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    v (Context rep -> Map VName (VariableInfo rep)
    forall {k} (rep :: k). Context rep -> Map VName (VariableInfo rep)
    assignments Context rep
    ctx) of
        Maybe (VariableInfo rep)
    Nothing -> Map VName Dependency
    forall a. Monoid a => a
    mempty -- Means a global.
        Just (VariableInfo Names
    deps Int
    lvl [BodyType]
    _parents VarType
    t) ->
          -- We detect whether it is a threadID or loop counter by checking
          -- whether or not it has any dependencies
          case VarType
    t of
            VarType
    ThreadID -> [(VName, Dependency)] -> Map VName Dependency
    forall k a. Ord k => [(k, a)] -> Map k a
    M.fromList [(VName
    v, Int -> VarType -> Dependency
    Dependency Int
    lvl VarType
    t)]
            VarType
    LoopVar -> [(VName, Dependency)] -> Map VName Dependency
    forall k a. Ord k => [(k, a)] -> Map k a
    M.fromList [(VName
    v, Int -> VarType -> Dependency
    Dependency Int
    lvl VarType
    t)]
            VarType
    Variable -> [Map VName Dependency] -> Map VName Dependency
    forall a. Monoid a => [a] -> a
    mconcat ([Map VName Dependency] -> Map VName Dependency)
    -> [Map VName Dependency] -> Map VName Dependency
    forall a b. (a -> b) -> a -> b
    $ (VName -> Map VName Dependency)
    -> [VName] -> [Map VName Dependency]
    forall a b. (a -> b) -> [a] -> [b]
    map (Context rep -> VName -> Map VName Dependency
    forall {k} (rep :: k). Context rep -> VName -> Map VName Dependency
    reduceDependencies Context rep
    ctx) ([VName] -> [Map VName Dependency])
    -> [VName] -> [Map VName Dependency]
    forall a b. (a -> b) -> a -> b
    $ Names -> [VName]
    namesToList Names
    deps
            VarType
    ConstType -> Map VName Dependency
    forall a. Monoid a => a
    mempty

    -- Misc functions

    -- Instances for AST types that we actually support
    instance Analyse GPU where
      analyseOp :: Op GPU -> Context GPU -> [VName] -> (Context GPU, IndexTable GPU)
    analyseOp Op GPU
    gpu_op
        | (SegOp SegOp SegLevel GPU
    op) <- Op GPU
    gpu_op = SegOp SegLevel GPU
    -> Context GPU -> [VName] -> (Context GPU, IndexTable GPU)
    forall rep lvl.
    Analyse rep =>
    SegOp lvl rep
    -> Context rep -> [VName] -> (Context rep, IndexTable rep)
    analyseSegOp SegOp SegLevel GPU
    op
        | (SizeOp SizeOp
    op) <- Op GPU
    gpu_op = SizeOp -> Context GPU -> [VName] -> (Context GPU, IndexTable GPU)
    forall {k} (rep :: k).
    SizeOp -> Context rep -> [VName] -> (Context rep, IndexTable rep)
    analyseSizeOp SizeOp
    op
        | (GPUBody [Type]
    _ Body GPU
    body) <- Op GPU
    gpu_op = (Context GPU, IndexTable GPU)
    -> [VName] -> (Context GPU, IndexTable GPU)
    forall a. a -> [VName] -> a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ((Context GPU, IndexTable GPU)
     -> [VName] -> (Context GPU, IndexTable GPU))
    -> (Context GPU -> (Context GPU, IndexTable GPU))
    -> Context GPU
    -> [VName]
    -> (Context GPU, IndexTable GPU)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Body GPU -> Context GPU -> (Context GPU, IndexTable GPU)
    forall rep.
    Analyse rep =>
    Body rep -> Context rep -> (Context rep, IndexTable rep)
    analyseGPUBody Body GPU
    body
        | (Futhark.IR.GPU.OtherOp SOAC GPU
    _) <- Op GPU
    gpu_op = Context GPU -> [VName] -> (Context GPU, IndexTable GPU)
    forall {k} (rep :: k).
    Context rep -> [VName] -> (Context rep, IndexTable rep)
    analyseOtherOp

    instance Analyse MC where
      analyseOp :: Op MC -> Context MC -> [VName] -> (Context MC, IndexTable MC)
    analyseOp Op MC
    mc_op
        | ParOp Maybe (SegOp () MC)
    Nothing SegOp () MC
    seq_segop <- Op MC
    mc_op = SegOp () MC -> Context MC -> [VName] -> (Context MC, IndexTable MC)
    forall rep lvl.
    Analyse rep =>
    SegOp lvl rep
    -> Context rep -> [VName] -> (Context rep, IndexTable rep)
    analyseSegOp SegOp () MC
    seq_segop
        | ParOp (Just SegOp () MC
    segop) SegOp () MC
    seq_segop <- Op MC
    mc_op = \Context MC
    ctx [VName]
    name -> do
            let (Context MC
    ctx', IndexTable MC
    res') = SegOp () MC -> Context MC -> [VName] -> (Context MC, IndexTable MC)
    forall rep lvl.
    Analyse rep =>
    SegOp lvl rep
    -> Context rep -> [VName] -> (Context rep, IndexTable rep)
    analyseSegOp SegOp () MC
    segop Context MC
    ctx [VName]
    name
            let (Context MC
    ctx'', IndexTable MC
    res'') = SegOp () MC -> Context MC -> [VName] -> (Context MC, IndexTable MC)
    forall rep lvl.
    Analyse rep =>
    SegOp lvl rep
    -> Context rep -> [VName] -> (Context rep, IndexTable rep)
    analyseSegOp SegOp () MC
    seq_segop Context MC
    ctx' [VName]
    name
            (Context MC
    ctx'', IndexTable MC -> IndexTable MC -> IndexTable MC
    forall {k} (rep :: k).
    IndexTable rep -> IndexTable rep -> IndexTable rep
    unionIndexTables IndexTable MC
    res' IndexTable MC
    res'')
        | Futhark.IR.MC.OtherOp SOAC MC
    _ <- Op MC
    mc_op = Context MC -> [VName] -> (Context MC, IndexTable MC)
    forall {k} (rep :: k).
    Context rep -> [VName] -> (Context rep, IndexTable rep)
    analyseOtherOp

    -- Unfortunately we need these instances, even though they may never appear.
    instance Analyse GPUMem where
      analyseOp :: Op GPUMem
    -> Context GPUMem -> [VName] -> (Context GPUMem, IndexTable GPUMem)
    analyseOp Op GPUMem
    _ = [Char]
    -> Context GPUMem -> [VName] -> (Context GPUMem, IndexTable GPUMem)
    forall a. HasCallStack => [Char] -> a
    error ([Char]
     -> Context GPUMem
     -> [VName]
     -> (Context GPUMem, IndexTable GPUMem))
    -> [Char]
    -> Context GPUMem
    -> [VName]
    -> (Context GPUMem, IndexTable GPUMem)
    forall a b. (a -> b) -> a -> b
    $ ShowS
    notImplementedYet [Char]
    "GPUMem"

    instance Analyse MCMem where
      analyseOp :: Op MCMem
    -> Context MCMem -> [VName] -> (Context MCMem, IndexTable MCMem)
    analyseOp Op MCMem
    _ = [Char]
    -> Context MCMem -> [VName] -> (Context MCMem, IndexTable MCMem)
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "Unexpected?"

    instance Analyse Seq where
      analyseOp :: Op Seq -> Context Seq -> [VName] -> (Context Seq, IndexTable Seq)
    analyseOp Op Seq
    _ = [Char] -> Context Seq -> [VName] -> (Context Seq, IndexTable Seq)
    forall a. HasCallStack => [Char] -> a
    error ([Char] -> Context Seq -> [VName] -> (Context Seq, IndexTable Seq))
    -> [Char]
    -> Context Seq
    -> [VName]
    -> (Context Seq, IndexTable Seq)
    forall a b. (a -> b) -> a -> b
    $ ShowS
    notImplementedYet [Char]
    "Seq"

    instance Analyse SeqMem where
      analyseOp :: Op SeqMem
    -> Context SeqMem -> [VName] -> (Context SeqMem, IndexTable SeqMem)
    analyseOp Op SeqMem
    _ = [Char]
    -> Context SeqMem -> [VName] -> (Context SeqMem, IndexTable SeqMem)
    forall a. HasCallStack => [Char] -> a
    error ([Char]
     -> Context SeqMem
     -> [VName]
     -> (Context SeqMem, IndexTable SeqMem))
    -> [Char]
    -> Context SeqMem
    -> [VName]
    -> (Context SeqMem, IndexTable SeqMem)
    forall a b. (a -> b) -> a -> b
    $ ShowS
    notImplementedYet [Char]
    "SeqMem"

    instance Analyse SOACS where
      analyseOp :: Op SOACS
    -> Context SOACS -> [VName] -> (Context SOACS, IndexTable SOACS)
    analyseOp Op SOACS
    _ = [Char]
    -> Context SOACS -> [VName] -> (Context SOACS, IndexTable SOACS)
    forall a. HasCallStack => [Char] -> a
    error ([Char]
     -> Context SOACS -> [VName] -> (Context SOACS, IndexTable SOACS))
    -> [Char]
    -> Context SOACS
    -> [VName]
    -> (Context SOACS, IndexTable SOACS)
    forall a b. (a -> b) -> a -> b
    $ ShowS
    notImplementedYet [Char]
    "SOACS"

    notImplementedYet :: String -> String
    notImplementedYet :: ShowS
    notImplementedYet [Char]
    s = [Char]
    "Access pattern analysis for the " [Char] -> ShowS
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    s [Char] -> ShowS
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    " backend is not implemented."

    instance Pretty (IndexTable rep) where
      pretty :: forall ann. IndexTable rep -> Doc ann
    pretty = [Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    stack ([Doc ann] -> Doc ann)
    -> (IndexTable rep -> [Doc ann]) -> IndexTable rep -> Doc ann
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ((SegOpName, Map ArrayName (Map VName [DimAccess rep])) -> Doc ann)
    -> [(SegOpName, Map ArrayName (Map VName [DimAccess rep]))]
    -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map (SegOpName, Map ArrayName (Map VName [DimAccess rep])) -> Doc ann
    forall {a} {ann}.
    Pretty a =>
    (a, Map ArrayName (Map VName [DimAccess rep])) -> Doc ann
    f ([(SegOpName, Map ArrayName (Map VName [DimAccess rep]))]
     -> [Doc ann])
    -> (IndexTable rep
        -> [(SegOpName, Map ArrayName (Map VName [DimAccess rep]))])
    -> IndexTable rep
    -> [Doc ann]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . IndexTable rep
    -> [(SegOpName, Map ArrayName (Map VName [DimAccess rep]))]
    forall k a. Map k a -> [(k, a)]
    M.toList :: IndexTable rep -> Doc ann
        where
          f :: (a, Map ArrayName (Map VName [DimAccess rep])) -> Doc ann
    f (a
    segop, Map ArrayName (Map VName [DimAccess rep])
    arrNameToIdxExprMap) = a -> Doc ann
    forall ann. a -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty a
    segop Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    forall ann. Doc ann
    colon Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Map ArrayName (Map VName [DimAccess rep]) -> Doc ann
    forall {a}. Map ArrayName (Map VName [DimAccess rep]) -> Doc a
    g Map ArrayName (Map VName [DimAccess rep])
    arrNameToIdxExprMap

          g :: Map ArrayName (Map VName [DimAccess rep]) -> Doc a
    g Map ArrayName (Map VName [DimAccess rep])
    maps = Doc a
    forall ann. Doc ann
    lbrace Doc a -> Doc a -> Doc a
    forall ann. Doc ann -> Doc ann -> Doc ann
    </> Int -> Doc a -> Doc a
    forall ann. Int -> Doc ann -> Doc ann
    indent Int
    4 ([(ArrayName, Map VName [DimAccess rep])] -> Doc a
    forall ann. [(ArrayName, Map VName [DimAccess rep])] -> Doc ann
    mapprintArray ([(ArrayName, Map VName [DimAccess rep])] -> Doc a)
    -> [(ArrayName, Map VName [DimAccess rep])] -> Doc a
    forall a b. (a -> b) -> a -> b
    $ Map ArrayName (Map VName [DimAccess rep])
    -> [(ArrayName, Map VName [DimAccess rep])]
    forall k a. Map k a -> [(k, a)]
    M.toList Map ArrayName (Map VName [DimAccess rep])
    maps) Doc a -> Doc a -> Doc a
    forall ann. Doc ann -> Doc ann -> Doc ann
    </> Doc a
    forall ann. Doc ann
    rbrace

          mapprintArray :: [(ArrayName, M.Map IndexExprName [DimAccess rep])] -> Doc ann
          mapprintArray :: forall ann. [(ArrayName, Map VName [DimAccess rep])] -> Doc ann
    mapprintArray [] = Doc ann
    ""
          mapprintArray [(ArrayName, Map VName [DimAccess rep])
    m] = (ArrayName, Map VName [DimAccess rep]) -> Doc ann
    forall ann. (ArrayName, Map VName [DimAccess rep]) -> Doc ann
    printArrayMap (ArrayName, Map VName [DimAccess rep])
    m
          mapprintArray ((ArrayName, Map VName [DimAccess rep])
    m : [(ArrayName, Map VName [DimAccess rep])]
    mm) = (ArrayName, Map VName [DimAccess rep]) -> Doc ann
    forall ann. (ArrayName, Map VName [DimAccess rep]) -> Doc ann
    printArrayMap (ArrayName, Map VName [DimAccess rep])
    m Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    </> [(ArrayName, Map VName [DimAccess rep])] -> Doc ann
    forall ann. [(ArrayName, Map VName [DimAccess rep])] -> Doc ann
    mapprintArray [(ArrayName, Map VName [DimAccess rep])]
    mm

          printArrayMap :: (ArrayName, M.Map IndexExprName [DimAccess rep]) -> Doc ann
          printArrayMap :: forall ann. (ArrayName, Map VName [DimAccess rep]) -> Doc ann
    printArrayMap ((VName
    name, [BodyType]
    _, [Int]
    layout), Map VName [DimAccess rep]
    maps) =
            Doc ann
    "(arr)"
              Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> VName -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. VName -> Doc ann
    pretty VName
    name
              Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    forall ann. Doc ann
    colon
              Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> [Int] -> Doc ann
    forall ann. [Int] -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty [Int]
    layout
              Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    forall ann. Doc ann
    lbrace
              Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    </> Int -> Doc ann -> Doc ann
    forall ann. Int -> Doc ann -> Doc ann
    indent Int
    4 ([(VName, [DimAccess rep])] -> Doc ann
    forall ann. [(VName, [DimAccess rep])] -> Doc ann
    mapprintIdxExpr (Map VName [DimAccess rep] -> [(VName, [DimAccess rep])]
    forall k a. Map k a -> [(k, a)]
    M.toList Map VName [DimAccess rep]
    maps))
              Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    </> Doc ann
    forall ann. Doc ann
    rbrace

          mapprintIdxExpr :: [(IndexExprName, [DimAccess rep])] -> Doc ann
          mapprintIdxExpr :: forall ann. [(VName, [DimAccess rep])] -> Doc ann
    mapprintIdxExpr [] = Doc ann
    ""
          mapprintIdxExpr [(VName, [DimAccess rep])
    m] = (VName, [DimAccess rep]) -> Doc ann
    forall {a} {ann}. Pretty a => (a, [DimAccess rep]) -> Doc ann
    printIdxExpMap (VName, [DimAccess rep])
    m
          mapprintIdxExpr ((VName, [DimAccess rep])
    m : [(VName, [DimAccess rep])]
    mm) = (VName, [DimAccess rep]) -> Doc ann
    forall {a} {ann}. Pretty a => (a, [DimAccess rep]) -> Doc ann
    printIdxExpMap (VName, [DimAccess rep])
    m Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    </> [(VName, [DimAccess rep])] -> Doc ann
    forall ann. [(VName, [DimAccess rep])] -> Doc ann
    mapprintIdxExpr [(VName, [DimAccess rep])]
    mm

          printIdxExpMap :: (a, [DimAccess rep]) -> Doc ann
    printIdxExpMap (a
    name, [DimAccess rep]
    mems) = Doc ann
    "(idx)" Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> a -> Doc ann
    forall ann. a -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty a
    name Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    ":" Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    </> Int -> Doc ann -> Doc ann
    forall ann. Int -> Doc ann -> Doc ann
    indent Int
    4 ([DimAccess rep] -> Doc ann
    forall ann. [DimAccess rep] -> Doc ann
    printDimAccess [DimAccess rep]
    mems)

          printDimAccess :: [DimAccess rep] -> Doc ann
          printDimAccess :: forall ann. [DimAccess rep] -> Doc ann
    printDimAccess [DimAccess rep]
    dim_accesses = [Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    stack ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
    forall a b. (a -> b) -> a -> b
    $ (Int -> DimAccess rep -> Doc ann)
    -> [Int] -> [DimAccess rep] -> [Doc ann]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith (((Int, DimAccess rep) -> Doc ann)
    -> Int -> DimAccess rep -> Doc ann
    forall a b c. ((a, b) -> c) -> a -> b -> c
    curry (Int, DimAccess rep) -> Doc ann
    forall ann. (Int, DimAccess rep) -> Doc ann
    printDim) [Int
    0 ..] [DimAccess rep]
    dim_accesses

          printDim :: (Int, DimAccess rep) -> Doc ann
          printDim :: forall ann. (Int, DimAccess rep) -> Doc ann
    printDim (Int
    i, DimAccess rep
    m) = Int -> Doc ann
    forall ann. Int -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty Int
    i Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    ":" Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Int -> Doc ann -> Doc ann
    forall ann. Int -> Doc ann -> Doc ann
    indent Int
    0 (DimAccess rep -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. DimAccess rep -> Doc ann
    pretty DimAccess rep
    m)

    instance Pretty (DimAccess rep) where
      pretty :: forall ann. DimAccess rep -> Doc ann
    pretty DimAccess rep
    dim_access =
        -- Instead of using `brackets $` we manually enclose with `[`s, to add
        -- spacing between the enclosed elements
        if case DimAccess rep -> Maybe VName
    forall {k} (rep :: k). DimAccess rep -> Maybe VName
    originalVar DimAccess rep
    dim_access of
          Maybe VName
    Nothing -> Bool
    True
          Just VName
    n ->
            Map VName Dependency -> Int
    forall a. Map VName a -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length (DimAccess rep -> Map VName Dependency
    forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
    dependencies DimAccess rep
    dim_access) Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    1 Bool -> Bool -> Bool
    && VName
    n VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    == [VName] -> VName
    forall a. HasCallStack => [a] -> a
    head (((VName, Dependency) -> VName) -> [(VName, Dependency)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName, Dependency) -> VName
    forall a b. (a, b) -> a
    fst ([(VName, Dependency)] -> [VName])
    -> [(VName, Dependency)] -> [VName]
    forall a b. (a -> b) -> a -> b
    $ Map VName Dependency -> [(VName, Dependency)]
    forall k a. Map k a -> [(k, a)]
    M.toList (Map VName Dependency -> [(VName, Dependency)])
    -> Map VName Dependency -> [(VName, Dependency)]
    forall a b. (a -> b) -> a -> b
    $ DimAccess rep -> Map VName Dependency
    forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
    dependencies DimAccess rep
    dim_access)
            -- Only print the original name if it is different from the first (and single) dependency
          then
            Doc ann
    "dependencies"
              Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    forall ann. Doc ann
    equals
              Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    align (Map VName Dependency -> Doc ann
    forall {ann}. Map VName Dependency -> Doc ann
    prettyDeps (Map VName Dependency -> Doc ann)
    -> Map VName Dependency -> Doc ann
    forall a b. (a -> b) -> a -> b
    $ DimAccess rep -> Map VName Dependency
    forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
    dependencies DimAccess rep
    dim_access)
          else
            Doc ann
    "dependencies"
              Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    forall ann. Doc ann
    equals
              Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Maybe VName -> Doc ann
    forall ann. Maybe VName -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty (DimAccess rep -> Maybe VName
    forall {k} (rep :: k). DimAccess rep -> Maybe VName
    originalVar DimAccess rep
    dim_access)
              Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    "->"
              Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    align (Map VName Dependency -> Doc ann
    forall {ann}. Map VName Dependency -> Doc ann
    prettyDeps (Map VName Dependency -> Doc ann)
    -> Map VName Dependency -> Doc ann
    forall a b. (a -> b) -> a -> b
    $ DimAccess rep -> Map VName Dependency
    forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
    dependencies DimAccess rep
    dim_access)
        where
          prettyDeps :: Map VName Dependency -> Doc ann
    prettyDeps = Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    braces (Doc ann -> Doc ann)
    -> (Map VName Dependency -> Doc ann)
    -> Map VName Dependency
    -> Doc ann
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    commasep ([Doc ann] -> Doc ann)
    -> (Map VName Dependency -> [Doc ann])
    -> Map VName Dependency
    -> Doc ann
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ((VName, Dependency) -> Doc ann)
    -> [(VName, Dependency)] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName, Dependency) -> Doc ann
    forall {a} {ann}. Pretty a => (a, Dependency) -> Doc ann
    printPair ([(VName, Dependency)] -> [Doc ann])
    -> (Map VName Dependency -> [(VName, Dependency)])
    -> Map VName Dependency
    -> [Doc ann]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Map VName Dependency -> [(VName, Dependency)]
    forall k a. Map k a -> [(k, a)]
    M.toList
          printPair :: (a, Dependency) -> Doc ann
    printPair (a
    name, Dependency Int
    lvl VarType
    vtype) = a -> Doc ann
    forall ann. a -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty a
    name Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Int -> Doc ann
    forall ann. Int -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty Int
    lvl Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> VarType -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. VarType -> Doc ann
    pretty VarType
    vtype

    instance Pretty SegOpName where
      pretty :: forall ann. SegOpName -> Doc ann
    pretty (SegmentedMap VName
    name) = Doc ann
    "(segmap)" Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> VName -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. VName -> Doc ann
    pretty VName
    name
      pretty (SegmentedRed VName
    name) = Doc ann
    "(segred)" Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> VName -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. VName -> Doc ann
    pretty VName
    name
      pretty (SegmentedScan VName
    name) = Doc ann
    "(segscan)" Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> VName -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. VName -> Doc ann
    pretty VName
    name
      pretty (SegmentedHist VName
    name) = Doc ann
    "(seghist)" Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> VName -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. VName -> Doc ann
    pretty VName
    name

    instance Pretty BodyType where
      pretty :: forall ann. BodyType -> Doc ann
    pretty (SegOpName (SegmentedMap VName
    name)) = VName -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. VName -> Doc ann
    pretty VName
    name Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    forall ann. Doc ann
    colon Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    "segmap"
      pretty (SegOpName (SegmentedRed VName
    name)) = VName -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. VName -> Doc ann
    pretty VName
    name Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    forall ann. Doc ann
    colon Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    "segred"
      pretty (SegOpName (SegmentedScan VName
    name)) = VName -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. VName -> Doc ann
    pretty VName
    name Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    forall ann. Doc ann
    colon Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    "segscan"
      pretty (SegOpName (SegmentedHist VName
    name)) = VName -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. VName -> Doc ann
    pretty VName
    name Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    forall ann. Doc ann
    colon Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    "seghist"
      pretty (LoopBodyName VName
    name) = VName -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. VName -> Doc ann
    pretty VName
    name Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    forall ann. Doc ann
    colon Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    "loop"
      pretty (CondBodyName VName
    name) = VName -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. VName -> Doc ann
    pretty VName
    name Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    forall ann. Doc ann
    colon Doc ann -> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann -> Doc ann
    <+> Doc ann
    "cond"

    instance Pretty VarType where
      pretty :: forall ann. VarType -> Doc ann
    pretty VarType
    ConstType = Doc ann
    "const"
      pretty VarType
    Variable = Doc ann
    "var"
      pretty VarType
    ThreadID = Doc ann
    "tid"
      pretty VarType
    LoopVar = Doc ann
    "iter"