    {-# LANGUAGE FlexibleContexts #-}
    {-# LANGUAGE FlexibleInstances #-}
    {-# LANGUAGE TypeFamilies #-}

    module Futhark.Optimise.ArrayShortCircuiting.TopdownAnalysis
      ( TopdownEnv (..),
        ScopeTab,
        TopDownHelper,
        InhibitTab,
        updateTopdownEnv,
        updateTopdownEnvLoop,
        getDirAliasedIxfn,
        getDirAliasedIxfn',
        addInvAliasesVarTab,
        areAnyAliased,
        isInScope,
        nonNegativesInPat,
      )
    where

    import Data.Map.Strict qualified as M
    import Data.Maybe
    import Futhark.Analysis.PrimExp.Convert
    import Futhark.IR.Aliases
    import Futhark.IR.GPUMem as GPU
    import Futhark.IR.MCMem as MC
    import Futhark.IR.Mem.LMAD qualified as LMAD
    import Futhark.Optimise.ArrayShortCircuiting.DataStructs

    type DirAlias = LMAD -> Maybe LMAD
    -- ^ A direct aliasing transformation

    type InvAlias = Maybe (LMAD -> LMAD)
    -- ^ An inverse aliasing transformation

    type VarAliasTab = M.Map VName (VName, DirAlias, InvAlias)

    type MemAliasTab = M.Map VName Names

    data TopdownEnv rep = TopdownEnv
      { -- | contains the already allocated memory blocks
        forall rep. TopdownEnv rep -> AllocTab
    alloc :: AllocTab,
        -- | variable info, including var-to-memblock assocs
        forall rep. TopdownEnv rep -> ScopeTab rep
    scope :: ScopeTab rep,
        -- | the inherited inhibitions from the previous try
        forall rep. TopdownEnv rep -> InhibitTab
    inhibited :: InhibitTab,
        -- | for statements such as transpose, reshape, index, etc., that alias
        --   an array variable: maps var-names to pair of aliased var name
        --   and index function transformation. For example, for
        --   @let b = a[slc]@ it should add the binding
        --   @ b |-> (a, `slice` slc )@
        forall rep. TopdownEnv rep -> VarAliasTab
    v_alias :: VarAliasTab,
        -- | keeps track of memory block aliasing.
        --   this needs to be implemented
        forall rep. TopdownEnv rep -> InhibitTab
    m_alias :: MemAliasTab,
        -- | Contains symbol information about the variables in the program. Used to
        -- determine if a variable is non-negative.
        forall rep. TopdownEnv rep -> Names
    nonNegatives :: Names,
        forall rep. TopdownEnv rep -> Map VName (PrimExp VName)
    scalarTable :: M.Map VName (PrimExp VName),
        -- | A list of known relations of the form 'VName' @<@ 'SubExp', typically
        -- gotten from 'LoopForm' and 'SegSpace'.
        forall rep. TopdownEnv rep -> [(VName, PrimExp VName)]
    knownLessThan :: [(VName, PrimExp VName)],
        -- | A list of the asserts encountered so far
        forall rep. TopdownEnv rep -> [SubExp]
    td_asserts :: [SubExp]
      }

    isInScope :: TopdownEnv rep -> VName -> Bool
    isInScope :: forall rep. TopdownEnv rep -> VName -> Bool
    isInScope TopdownEnv rep
    td_env VName
    m = VName
    m VName -> Map VName (NameInfo (Aliases rep)) -> Bool
    forall k a. Ord k => k -> Map k a -> Bool
    `M.member` TopdownEnv rep -> Map VName (NameInfo (Aliases rep))
    forall rep. TopdownEnv rep -> ScopeTab rep
    scope TopdownEnv rep
    td_env

    -- | Get alias and (direct) index function mapping from expression
    getDirAliasFromExp :: Exp (Aliases rep) -> Maybe (VName, DirAlias)
    getDirAliasFromExp :: forall rep. Exp (Aliases rep) -> Maybe (VName, DirAlias)
    getDirAliasFromExp (BasicOp (SubExp (Var VName
    x))) = (VName, DirAlias) -> Maybe (VName, DirAlias)
    forall a. a -> Maybe a
    Just (VName
    x, DirAlias
    forall a. a -> Maybe a
    Just)
    getDirAliasFromExp (BasicOp (Opaque OpaqueOp
    _ (Var VName
    x))) = (VName, DirAlias) -> Maybe (VName, DirAlias)
    forall a. a -> Maybe a
    Just (VName
    x, DirAlias
    forall a. a -> Maybe a
    Just)
    getDirAliasFromExp (BasicOp (Reshape VName
    x NewShape SubExp
    shp)) =
      case NewShape SubExp -> ReshapeKind
    reshapeKind NewShape SubExp
    shp of
        ReshapeKind
    ReshapeCoerce ->
          (VName, DirAlias) -> Maybe (VName, DirAlias)
    forall a. a -> Maybe a
    Just (VName
    x, DirAlias
    forall a. a -> Maybe a
    Just DirAlias -> (LMAD -> LMAD) -> DirAlias
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (LMAD -> Shape (TPrimExp Int64 VName) -> LMAD
    forall num. LMAD num -> Shape num -> LMAD num
    `LMAD.coerce` (SubExp -> TPrimExp Int64 VName)
    -> [SubExp] -> Shape (TPrimExp Int64 VName)
    forall a b. (a -> b) -> [a] -> [b]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap SubExp -> TPrimExp Int64 VName
    pe64 (ShapeBase SubExp -> [SubExp]
    forall d. ShapeBase d -> [d]
    shapeDims (ShapeBase SubExp -> [SubExp]) -> ShapeBase SubExp -> [SubExp]
    forall a b. (a -> b) -> a -> b
    $ NewShape SubExp -> ShapeBase SubExp
    forall d. NewShape d -> ShapeBase d
    newShape NewShape SubExp
    shp)))
        ReshapeKind
    ReshapeArbitrary ->
          (VName, DirAlias) -> Maybe (VName, DirAlias)
    forall a. a -> Maybe a
    Just (VName
    x, (LMAD -> Shape (TPrimExp Int64 VName) -> Maybe LMAD
    forall num.
    (Eq num, IntegralExp num) =>
    LMAD num -> Shape num -> Maybe (LMAD num)
    `LMAD.reshape` (SubExp -> TPrimExp Int64 VName)
    -> [SubExp] -> Shape (TPrimExp Int64 VName)
    forall a b. (a -> b) -> [a] -> [b]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap SubExp -> TPrimExp Int64 VName
    pe64 (ShapeBase SubExp -> [SubExp]
    forall d. ShapeBase d -> [d]
    shapeDims (ShapeBase SubExp -> [SubExp]) -> ShapeBase SubExp -> [SubExp]
    forall a b. (a -> b) -> a -> b
    $ NewShape SubExp -> ShapeBase SubExp
    forall d. NewShape d -> ShapeBase d
    newShape NewShape SubExp
    shp)))
    getDirAliasFromExp (BasicOp (Rearrange VName
    _ [Int]
    _)) =
      Maybe (VName, DirAlias)
    forall a. Maybe a
    Nothing
    getDirAliasFromExp (BasicOp (Index VName
    x Slice SubExp
    slc)) =
      (VName, DirAlias) -> Maybe (VName, DirAlias)
    forall a. a -> Maybe a
    Just (VName
    x, DirAlias
    forall a. a -> Maybe a
    Just DirAlias -> (LMAD -> LMAD) -> DirAlias
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (LMAD -> Slice (TPrimExp Int64 VName) -> LMAD
    forall num.
    (Eq num, IntegralExp num) =>
    LMAD num -> Slice num -> LMAD num
    `LMAD.slice` ([DimIndex (TPrimExp Int64 VName)] -> Slice (TPrimExp Int64 VName)
    forall d. [DimIndex d] -> Slice d
    Slice ([DimIndex (TPrimExp Int64 VName)] -> Slice (TPrimExp Int64 VName))
    -> [DimIndex (TPrimExp Int64 VName)]
    -> Slice (TPrimExp Int64 VName)
    forall a b. (a -> b) -> a -> b
    $ (DimIndex SubExp -> DimIndex (TPrimExp Int64 VName))
    -> [DimIndex SubExp] -> [DimIndex (TPrimExp Int64 VName)]
    forall a b. (a -> b) -> [a] -> [b]
    map ((SubExp -> TPrimExp Int64 VName)
    -> DimIndex SubExp -> DimIndex (TPrimExp Int64 VName)
    forall a b. (a -> b) -> DimIndex a -> DimIndex b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap SubExp -> TPrimExp Int64 VName
    pe64) ([DimIndex SubExp] -> [DimIndex (TPrimExp Int64 VName)])
    -> [DimIndex SubExp] -> [DimIndex (TPrimExp Int64 VName)]
    forall a b. (a -> b) -> a -> b
    $ Slice SubExp -> [DimIndex SubExp]
    forall d. Slice d -> [DimIndex d]
    unSlice Slice SubExp
    slc)))
    getDirAliasFromExp (BasicOp (Update Safety
    _ VName
    x Slice SubExp
    _ SubExp
    _elm)) = (VName, DirAlias) -> Maybe (VName, DirAlias)
    forall a. a -> Maybe a
    Just (VName
    x, DirAlias
    forall a. a -> Maybe a
    Just)
    getDirAliasFromExp (BasicOp (FlatIndex VName
    x (FlatSlice SubExp
    offset [FlatDimIndex SubExp]
    idxs))) =
      (VName, DirAlias) -> Maybe (VName, DirAlias)
    forall a. a -> Maybe a
    Just
        ( VName
    x,
          DirAlias
    forall a. a -> Maybe a
    Just DirAlias -> (LMAD -> LMAD) -> DirAlias
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (LMAD -> FlatSlice (TPrimExp Int64 VName) -> LMAD
    forall num.
    IntegralExp num =>
    LMAD num -> FlatSlice num -> LMAD num
    `LMAD.flatSlice` TPrimExp Int64 VName
    -> [FlatDimIndex (TPrimExp Int64 VName)]
    -> FlatSlice (TPrimExp Int64 VName)
    forall d. d -> [FlatDimIndex d] -> FlatSlice d
    FlatSlice (SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    offset) ((FlatDimIndex SubExp -> FlatDimIndex (TPrimExp Int64 VName))
    -> [FlatDimIndex SubExp] -> [FlatDimIndex (TPrimExp Int64 VName)]
    forall a b. (a -> b) -> [a] -> [b]
    map ((SubExp -> TPrimExp Int64 VName)
    -> FlatDimIndex SubExp -> FlatDimIndex (TPrimExp Int64 VName)
    forall a b. (a -> b) -> FlatDimIndex a -> FlatDimIndex b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap SubExp -> TPrimExp Int64 VName
    pe64) [FlatDimIndex SubExp]
    idxs))
        )
    getDirAliasFromExp (BasicOp (FlatUpdate VName
    x FlatSlice SubExp
    _ VName
    _)) = (VName, DirAlias) -> Maybe (VName, DirAlias)
    forall a. a -> Maybe a
    Just (VName
    x, DirAlias
    forall a. a -> Maybe a
    Just)
    getDirAliasFromExp Exp (Aliases rep)
    _ = Maybe (VName, DirAlias)
    forall a. Maybe a
    Nothing

    -- | This was former @createsAliasedArrOK@ from DataStructs
    --   While Rearrange creates aliased arrays, we
    --   do not yet support them because it would mean we have
    --   to "reverse" the index function, for example to support
    --   coalescing in the case below,
    --       @let a = map f a0   @
    --       @let b = transpose a@
    --       @let y[4] = copy(b) @
    --   we would need to assign to @a@ as index function, the
    --   inverse of the transpose, such that, when creating @b@
    --   by transposition we get a directly-mapped array, which
    --   is expected by the copying in y[4].
    --   For the moment we support only transposition and VName-expressions,
    --     but rotations and full slices could also be supported.
    --
    -- This function complements 'getDirAliasFromExp' by returning a function that
    -- applies the inverse index function transformation.
    getInvAliasFromExp :: Exp (Aliases rep) -> InvAlias
    getInvAliasFromExp :: forall rep. Exp (Aliases rep) -> InvAlias
    getInvAliasFromExp (BasicOp (SubExp (Var VName
    _))) = (LMAD -> LMAD) -> InvAlias
    forall a. a -> Maybe a
    Just LMAD -> LMAD
    forall a. a -> a
    id
    getInvAliasFromExp (BasicOp (Opaque OpaqueOp
    _ (Var VName
    _))) = (LMAD -> LMAD) -> InvAlias
    forall a. a -> Maybe a
    Just LMAD -> LMAD
    forall a. a -> a
    id
    getInvAliasFromExp (BasicOp Update {}) = (LMAD -> LMAD) -> InvAlias
    forall a. a -> Maybe a
    Just LMAD -> LMAD
    forall a. a -> a
    id
    getInvAliasFromExp (BasicOp (Rearrange VName
    _ [Int]
    perm)) =
      (LMAD -> LMAD) -> InvAlias
    forall a. a -> Maybe a
    Just (LMAD -> [Int] -> LMAD
    forall num. LMAD num -> [Int] -> LMAD num
    `LMAD.permute` [Int] -> [Int]
    rearrangeInverse [Int]
    perm)
    getInvAliasFromExp Exp (Aliases rep)
    _ = InvAlias
    forall a. Maybe a
    Nothing

    class TopDownHelper inner where
      innerNonNegatives :: [VName] -> inner -> Names

      innerKnownLessThan :: inner -> [(VName, PrimExp VName)]

      scopeHelper :: inner -> Scope rep

    instance TopDownHelper (SegOp lvl rep) where
      innerNonNegatives :: [VName] -> SegOp lvl rep -> Names
    innerNonNegatives [VName]
    _ SegOp lvl rep
    op =
        ((VName, SubExp) -> Names) -> [(VName, SubExp)] -> Names
    forall m a. Monoid m => (a -> m) -> [a] -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap (VName -> Names
    oneName (VName -> Names)
    -> ((VName, SubExp) -> VName) -> (VName, SubExp) -> Names
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VName, SubExp) -> VName
    forall a b. (a, b) -> a
    fst) ([(VName, SubExp)] -> Names) -> [(VName, SubExp)] -> Names
    forall a b. (a -> b) -> a -> b
    $ SegSpace -> [(VName, SubExp)]
    unSegSpace (SegSpace -> [(VName, SubExp)]) -> SegSpace -> [(VName, SubExp)]
    forall a b. (a -> b) -> a -> b
    $ SegOp lvl rep -> SegSpace
    forall lvl rep. SegOp lvl rep -> SegSpace
    segSpace SegOp lvl rep
    op

      innerKnownLessThan :: SegOp lvl rep -> [(VName, PrimExp VName)]
    innerKnownLessThan SegOp lvl rep
    op =
        ((VName, SubExp) -> (VName, PrimExp VName))
    -> [(VName, SubExp)] -> [(VName, PrimExp VName)]
    forall a b. (a -> b) -> [a] -> [b]
    map ((SubExp -> PrimExp VName)
    -> (VName, SubExp) -> (VName, PrimExp VName)
    forall a b. (a -> b) -> (VName, a) -> (VName, b)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap ((SubExp -> PrimExp VName)
     -> (VName, SubExp) -> (VName, PrimExp VName))
    -> (SubExp -> PrimExp VName)
    -> (VName, SubExp)
    -> (VName, PrimExp VName)
    forall a b. (a -> b) -> a -> b
    $ PrimType -> SubExp -> PrimExp VName
    primExpFromSubExp (PrimType -> SubExp -> PrimExp VName)
    -> PrimType -> SubExp -> PrimExp VName
    forall a b. (a -> b) -> a -> b
    $ IntType -> PrimType
    IntType IntType
    Int64) ([(VName, SubExp)] -> [(VName, PrimExp VName)])
    -> [(VName, SubExp)] -> [(VName, PrimExp VName)]
    forall a b. (a -> b) -> a -> b
    $ SegSpace -> [(VName, SubExp)]
    unSegSpace (SegSpace -> [(VName, SubExp)]) -> SegSpace -> [(VName, SubExp)]
    forall a b. (a -> b) -> a -> b
    $ SegOp lvl rep -> SegSpace
    forall lvl rep. SegOp lvl rep -> SegSpace
    segSpace SegOp lvl rep
    op

      scopeHelper :: forall rep. SegOp lvl rep -> Scope rep
    scopeHelper SegOp lvl rep
    op = SegSpace -> Scope rep
    forall rep. SegSpace -> Scope rep
    scopeOfSegSpace (SegSpace -> Scope rep) -> SegSpace -> Scope rep
    forall a b. (a -> b) -> a -> b
    $ SegOp lvl rep -> SegSpace
    forall lvl rep. SegOp lvl rep -> SegSpace
    segSpace SegOp lvl rep
    op

    instance TopDownHelper (HostOp NoOp (Aliases GPUMem)) where
      innerNonNegatives :: [VName] -> HostOp NoOp (Aliases GPUMem) -> Names
    innerNonNegatives [VName]
    vs (SegOp SegOp SegLevel (Aliases GPUMem)
    op) = [VName] -> SegOp SegLevel (Aliases GPUMem) -> Names
    forall inner. TopDownHelper inner => [VName] -> inner -> Names
    innerNonNegatives [VName]
    vs SegOp SegLevel (Aliases GPUMem)
    op
      innerNonNegatives [VName
    vname] (SizeOp (GetSize Name
    _ SizeClass
    _)) = VName -> Names
    oneName VName
    vname
      innerNonNegatives [VName
    vname] (SizeOp (GetSizeMax SizeClass
    _)) = VName -> Names
    oneName VName
    vname
      innerNonNegatives [VName]
    _ HostOp NoOp (Aliases GPUMem)
    _ = Names
    forall a. Monoid a => a
    mempty

      innerKnownLessThan :: HostOp NoOp (Aliases GPUMem) -> [(VName, PrimExp VName)]
    innerKnownLessThan (SegOp SegOp SegLevel (Aliases GPUMem)
    op) = SegOp SegLevel (Aliases GPUMem) -> [(VName, PrimExp VName)]
    forall inner.
    TopDownHelper inner =>
    inner -> [(VName, PrimExp VName)]
    innerKnownLessThan SegOp SegLevel (Aliases GPUMem)
    op
      innerKnownLessThan HostOp NoOp (Aliases GPUMem)
    _ = [(VName, PrimExp VName)]
    forall a. Monoid a => a
    mempty

      scopeHelper :: forall rep. HostOp NoOp (Aliases GPUMem) -> Scope rep
    scopeHelper (SegOp SegOp SegLevel (Aliases GPUMem)
    op) = SegOp SegLevel (Aliases GPUMem) -> Scope rep
    forall rep. SegOp SegLevel (Aliases GPUMem) -> Scope rep
    forall inner rep. TopDownHelper inner => inner -> Scope rep
    scopeHelper SegOp SegLevel (Aliases GPUMem)
    op
      scopeHelper HostOp NoOp (Aliases GPUMem)
    _ = Scope rep
    forall a. Monoid a => a
    mempty

    instance (TopDownHelper (inner (Aliases MCMem))) => TopDownHelper (MC.MCOp inner (Aliases MCMem)) where
      innerNonNegatives :: [VName] -> MCOp inner (Aliases MCMem) -> Names
    innerNonNegatives [VName]
    vs (ParOp Maybe (SegOp () (Aliases MCMem))
    par_op SegOp () (Aliases MCMem)
    op) =
        Names
    -> (SegOp () (Aliases MCMem) -> Names)
    -> Maybe (SegOp () (Aliases MCMem))
    -> Names
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe Names
    forall a. Monoid a => a
    mempty ([VName] -> SegOp () (Aliases MCMem) -> Names
    forall inner. TopDownHelper inner => [VName] -> inner -> Names
    innerNonNegatives [VName]
    vs) Maybe (SegOp () (Aliases MCMem))
    par_op
          Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> [VName] -> SegOp () (Aliases MCMem) -> Names
    forall inner. TopDownHelper inner => [VName] -> inner -> Names
    innerNonNegatives [VName]
    vs SegOp () (Aliases MCMem)
    op
      innerNonNegatives [VName]
    vs (MC.OtherOp inner (Aliases MCMem)
    op) =
        [VName] -> inner (Aliases MCMem) -> Names
    forall inner. TopDownHelper inner => [VName] -> inner -> Names
    innerNonNegatives [VName]
    vs inner (Aliases MCMem)
    op
      innerKnownLessThan :: MCOp inner (Aliases MCMem) -> [(VName, PrimExp VName)]
    innerKnownLessThan (ParOp Maybe (SegOp () (Aliases MCMem))
    par_op SegOp () (Aliases MCMem)
    op) =
        [(VName, PrimExp VName)]
    -> (SegOp () (Aliases MCMem) -> [(VName, PrimExp VName)])
    -> Maybe (SegOp () (Aliases MCMem))
    -> [(VName, PrimExp VName)]
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe [(VName, PrimExp VName)]
    forall a. Monoid a => a
    mempty SegOp () (Aliases MCMem) -> [(VName, PrimExp VName)]
    forall inner.
    TopDownHelper inner =>
    inner -> [(VName, PrimExp VName)]
    innerKnownLessThan Maybe (SegOp () (Aliases MCMem))
    par_op [(VName, PrimExp VName)]
    -> [(VName, PrimExp VName)] -> [(VName, PrimExp VName)]
    forall a. Semigroup a => a -> a -> a
    <> SegOp () (Aliases MCMem) -> [(VName, PrimExp VName)]
    forall inner.
    TopDownHelper inner =>
    inner -> [(VName, PrimExp VName)]
    innerKnownLessThan SegOp () (Aliases MCMem)
    op
      innerKnownLessThan (MC.OtherOp inner (Aliases MCMem)
    op) =
        inner (Aliases MCMem) -> [(VName, PrimExp VName)]
    forall inner.
    TopDownHelper inner =>
    inner -> [(VName, PrimExp VName)]
    innerKnownLessThan inner (Aliases MCMem)
    op
      scopeHelper :: forall rep. MCOp inner (Aliases MCMem) -> Scope rep
    scopeHelper (ParOp Maybe (SegOp () (Aliases MCMem))
    par_op SegOp () (Aliases MCMem)
    op) =
        Scope rep
    -> (SegOp () (Aliases MCMem) -> Scope rep)
    -> Maybe (SegOp () (Aliases MCMem))
    -> Scope rep
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe Scope rep
    forall a. Monoid a => a
    mempty SegOp () (Aliases MCMem) -> Scope rep
    forall rep. SegOp () (Aliases MCMem) -> Scope rep
    forall inner rep. TopDownHelper inner => inner -> Scope rep
    scopeHelper Maybe (SegOp () (Aliases MCMem))
    par_op Scope rep -> Scope rep -> Scope rep
    forall a. Semigroup a => a -> a -> a
    <> SegOp () (Aliases MCMem) -> Scope rep
    forall rep. SegOp () (Aliases MCMem) -> Scope rep
    forall inner rep. TopDownHelper inner => inner -> Scope rep
    scopeHelper SegOp () (Aliases MCMem)
    op
      scopeHelper MC.OtherOp {} = Scope rep
    forall a. Monoid a => a
    mempty

    instance TopDownHelper (NoOp rep) where
      innerNonNegatives :: [VName] -> NoOp rep -> Names
    innerNonNegatives [VName]
    _ NoOp rep
    NoOp = Names
    forall a. Monoid a => a
    mempty
      innerKnownLessThan :: NoOp rep -> [(VName, PrimExp VName)]
    innerKnownLessThan NoOp rep
    NoOp = [(VName, PrimExp VName)]
    forall a. Monoid a => a
    mempty
      scopeHelper :: forall rep. NoOp rep -> Scope rep
    scopeHelper NoOp rep
    NoOp = Scope rep
    forall a. Monoid a => a
    mempty

    -- | fills in the TopdownEnv table
    updateTopdownEnv ::
      (ASTRep rep, Op rep ~ MemOp inner rep, TopDownHelper (inner (Aliases rep))) =>
      TopdownEnv rep ->
      Stm (Aliases rep) ->
      TopdownEnv rep
    updateTopdownEnv :: forall rep (inner :: * -> *).
    (ASTRep rep, Op rep ~ MemOp inner rep,
     TopDownHelper (inner (Aliases rep))) =>
    TopdownEnv rep -> Stm (Aliases rep) -> TopdownEnv rep
    updateTopdownEnv TopdownEnv rep
    env stm :: Stm (Aliases rep)
    stm@(Let (Pat [PatElem (LetDec (Aliases rep))
    pe]) StmAux (ExpDec (Aliases rep))
    _ (Op (Alloc (Var VName
    vname) Space
    sp))) =
      TopdownEnv rep
    env
        { alloc = M.insert (patElemName pe) sp $ alloc env,
          scope = scope env <> scopeOf stm,
          nonNegatives = nonNegatives env <> oneName vname
        }
    updateTopdownEnv TopdownEnv rep
    env stm :: Stm (Aliases rep)
    stm@(Let Pat (LetDec (Aliases rep))
    pat StmAux (ExpDec (Aliases rep))
    _ (Op (Inner inner (Aliases rep)
    inner))) =
      TopdownEnv rep
    env
        { scope = scope env <> scopeOf stm <> scopeHelper inner,
          nonNegatives = nonNegatives env <> innerNonNegatives (patNames pat) inner,
          knownLessThan = knownLessThan env <> innerKnownLessThan inner
        }
    updateTopdownEnv TopdownEnv rep
    env stm :: Stm (Aliases rep)
    stm@(Let (Pat [PatElem (LetDec (Aliases rep))]
    _) StmAux (ExpDec (Aliases rep))
    _ (BasicOp (Assert SubExp
    se ErrorMsg SubExp
    _))) =
      TopdownEnv rep
    env
        { scope = scope env <> scopeOf stm,
          td_asserts = se : td_asserts env
        }
    updateTopdownEnv TopdownEnv rep
    env stm :: Stm (Aliases rep)
    stm@(Let (Pat [PatElem (LetDec (Aliases rep))
    pe]) StmAux (ExpDec (Aliases rep))
    _ Exp (Aliases rep)
    e)
      | Just (VName
    x, DirAlias
    ixfn) <- Exp (Aliases rep) -> Maybe (VName, DirAlias)
    forall rep. Exp (Aliases rep) -> Maybe (VName, DirAlias)
    getDirAliasFromExp Exp (Aliases rep)
    e =
          let ixfn_inv :: InvAlias
    ixfn_inv = Exp (Aliases rep) -> InvAlias
    forall rep. Exp (Aliases rep) -> InvAlias
    getInvAliasFromExp Exp (Aliases rep)
    e
           in TopdownEnv rep
    env
                { v_alias = M.insert (patElemName pe) (x, ixfn, ixfn_inv) (v_alias env),
                  scope = scope env <> scopeOf stm,
                  nonNegatives = nonNegatives env <> nonNegativesInPat (stmPat stm)
                }
    updateTopdownEnv TopdownEnv rep
    env Stm (Aliases rep)
    stm =
      TopdownEnv rep
    env
        { scope = scope env <> scopeOf stm,
          nonNegatives = nonNegatives env <> nonNegativesInPat (stmPat stm)
        }

    nonNegativesInPat :: (Typed rep) => Pat rep -> Names
    nonNegativesInPat :: forall rep. Typed rep => Pat rep -> Names
    nonNegativesInPat (Pat [PatElem rep]
    elems) =
      (PatElem rep -> Names) -> [PatElem rep] -> Names
    forall m a. Monoid m => (a -> m) -> [a] -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap ([VName] -> Names
    namesFromList ([VName] -> Names)
    -> (PatElem rep -> [VName]) -> PatElem rep -> Names
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (SubExp -> Maybe VName) -> [SubExp] -> [VName]
    forall a b. (a -> Maybe b) -> [a] -> [b]
    mapMaybe SubExp -> Maybe VName
    subExpVar ([SubExp] -> [VName])
    -> (PatElem rep -> [SubExp]) -> PatElem rep -> [VName]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . TypeBase (ShapeBase SubExp) NoUniqueness -> [SubExp]
    forall u. TypeBase (ShapeBase SubExp) u -> [SubExp]
    arrayDims (TypeBase (ShapeBase SubExp) NoUniqueness -> [SubExp])
    -> (PatElem rep -> TypeBase (ShapeBase SubExp) NoUniqueness)
    -> PatElem rep
    -> [SubExp]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . PatElem rep -> TypeBase (ShapeBase SubExp) NoUniqueness
    forall t. Typed t => t -> TypeBase (ShapeBase SubExp) NoUniqueness
    typeOf) [PatElem rep]
    elems

    -- | The topdown handler for loops.
    updateTopdownEnvLoop :: TopdownEnv rep -> [(FParam rep, SubExp)] -> LoopForm -> TopdownEnv rep
    updateTopdownEnvLoop :: forall rep.
    TopdownEnv rep
    -> [(FParam rep, SubExp)] -> LoopForm -> TopdownEnv rep
    updateTopdownEnvLoop TopdownEnv rep
    td_env [(FParam rep, SubExp)]
    arginis LoopForm
    lform =
      let scopetab :: ScopeTab rep
    scopetab =
            TopdownEnv rep -> ScopeTab rep
    forall rep. TopdownEnv rep -> ScopeTab rep
    scope TopdownEnv rep
    td_env
              ScopeTab rep -> ScopeTab rep -> ScopeTab rep
    forall a. Semigroup a => a -> a -> a
    <> [FParam rep] -> ScopeTab rep
    forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
    scopeOfFParams (((FParam rep, SubExp) -> FParam rep)
    -> [(FParam rep, SubExp)] -> [FParam rep]
    forall a b. (a -> b) -> [a] -> [b]
    map (FParam rep, SubExp) -> FParam rep
    forall a b. (a, b) -> a
    fst [(FParam rep, SubExp)]
    arginis)
              ScopeTab rep -> ScopeTab rep -> ScopeTab rep
    forall a. Semigroup a => a -> a -> a
    <> LoopForm -> ScopeTab rep
    forall rep. LoopForm -> Scope rep
    scopeOfLoopForm LoopForm
    lform
          non_negatives :: Names
    non_negatives =
            TopdownEnv rep -> Names
    forall rep. TopdownEnv rep -> Names
    nonNegatives TopdownEnv rep
    td_env Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> case LoopForm
    lform of
              ForLoop VName
    v IntType
    _ SubExp
    _ -> VName -> Names
    oneName VName
    v
              LoopForm
    _ -> Names
    forall a. Monoid a => a
    mempty
          less_than :: [(VName, PrimExp VName)]
    less_than =
            case LoopForm
    lform of
              ForLoop VName
    v IntType
    _ SubExp
    b -> [(VName
    v, PrimType -> SubExp -> PrimExp VName
    primExpFromSubExp (IntType -> PrimType
    IntType IntType
    Int64) SubExp
    b)]
              LoopForm
    _ -> [(VName, PrimExp VName)]
    forall a. Monoid a => a
    mempty
       in TopdownEnv rep
    td_env
            { scope = scopetab,
              nonNegatives = non_negatives,
              knownLessThan = less_than <> knownLessThan td_env
            }

    -- | Get direct aliased index function.  Returns a triple of current memory
    -- block to be coalesced, the destination memory block and the index function of
    -- the access in the space of the destination block.
    getDirAliasedIxfn :: (HasMemBlock (Aliases rep)) => TopdownEnv rep -> CoalsTab -> VName -> Maybe (VName, VName, LMAD)
    getDirAliasedIxfn :: forall rep.
    HasMemBlock (Aliases rep) =>
    TopdownEnv rep -> CoalsTab -> VName -> Maybe (VName, VName, LMAD)
    getDirAliasedIxfn TopdownEnv rep
    td_env CoalsTab
    coals_tab VName
    x =
      case VName -> Scope (Aliases rep) -> Maybe ArrayMemBound
    forall rep.
    HasMemBlock rep =>
    VName -> Scope rep -> Maybe ArrayMemBound
    getScopeMemInfo VName
    x (TopdownEnv rep -> Scope (Aliases rep)
    forall rep. TopdownEnv rep -> ScopeTab rep
    scope TopdownEnv rep
    td_env) of
        Just (MemBlock PrimType
    _ ShapeBase SubExp
    _ VName
    m_x LMAD
    orig_ixfun) ->
          case VName -> CoalsTab -> Maybe CoalsEntry
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    m_x CoalsTab
    coals_tab of
            Just CoalsEntry
    coal_etry -> do
              (Coalesced _ (MemBlock _ _ m ixf) _) <- VarAliasTab -> Map VName Coalesced -> VName -> Maybe Coalesced
    walkAliasTab (TopdownEnv rep -> VarAliasTab
    forall rep. TopdownEnv rep -> VarAliasTab
    v_alias TopdownEnv rep
    td_env) (CoalsEntry -> Map VName Coalesced
    vartab CoalsEntry
    coal_etry) VName
    x
              pure (m_x, m, ixf)
            Maybe CoalsEntry
    Nothing ->
              -- This value is not subject to coalescing at the moment. Just return the
              -- original index function
              (VName, VName, LMAD) -> Maybe (VName, VName, LMAD)
    forall a. a -> Maybe a
    Just (VName
    m_x, VName
    m_x, LMAD
    orig_ixfun)
        Maybe ArrayMemBound
    Nothing -> Maybe (VName, VName, LMAD)
    forall a. Maybe a
    Nothing

    -- | Like 'getDirAliasedIxfn', but this version returns 'Nothing' if the value
    -- is not currently subject to coalescing.
    getDirAliasedIxfn' :: (HasMemBlock (Aliases rep)) => TopdownEnv rep -> CoalsTab -> VName -> Maybe (VName, VName, LMAD)
    getDirAliasedIxfn' :: forall rep.
    HasMemBlock (Aliases rep) =>
    TopdownEnv rep -> CoalsTab -> VName -> Maybe (VName, VName, LMAD)
    getDirAliasedIxfn' TopdownEnv rep
    td_env CoalsTab
    coals_tab VName
    x =
      case VName -> Scope (Aliases rep) -> Maybe ArrayMemBound
    forall rep.
    HasMemBlock rep =>
    VName -> Scope rep -> Maybe ArrayMemBound
    getScopeMemInfo VName
    x (TopdownEnv rep -> Scope (Aliases rep)
    forall rep. TopdownEnv rep -> ScopeTab rep
    scope TopdownEnv rep
    td_env) of
        Just (MemBlock PrimType
    _ ShapeBase SubExp
    _ VName
    m_x LMAD
    _) ->
          case VName -> CoalsTab -> Maybe CoalsEntry
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    m_x CoalsTab
    coals_tab of
            Just CoalsEntry
    coal_etry -> do
              (Coalesced _ (MemBlock _ _ m ixf) _) <- VarAliasTab -> Map VName Coalesced -> VName -> Maybe Coalesced
    walkAliasTab (TopdownEnv rep -> VarAliasTab
    forall rep. TopdownEnv rep -> VarAliasTab
    v_alias TopdownEnv rep
    td_env) (CoalsEntry -> Map VName Coalesced
    vartab CoalsEntry
    coal_etry) VName
    x
              pure (m_x, m, ixf)
            Maybe CoalsEntry
    Nothing ->
              -- This value is not subject to coalescing at the moment. Just return the
              -- original index function
              Maybe (VName, VName, LMAD)
    forall a. Maybe a
    Nothing
        Maybe ArrayMemBound
    Nothing -> Maybe (VName, VName, LMAD)
    forall a. Maybe a
    Nothing

    -- | Given a 'VName', walk the 'VarAliasTab' until found in the 'Map'.
    walkAliasTab ::
      VarAliasTab ->
      M.Map VName Coalesced ->
      VName ->
      Maybe Coalesced
    walkAliasTab :: VarAliasTab -> Map VName Coalesced -> VName -> Maybe Coalesced
    walkAliasTab VarAliasTab
    _ Map VName Coalesced
    vtab VName
    x
      | Just Coalesced
    c <- VName -> Map VName Coalesced -> Maybe Coalesced
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    x Map VName Coalesced
    vtab =
          Coalesced -> Maybe Coalesced
    forall a. a -> Maybe a
    Just Coalesced
    c -- @x@ is in @vartab@ together with its new ixfun
    walkAliasTab VarAliasTab
    alias_tab Map VName Coalesced
    vtab VName
    x
      | Just (VName
    x0, DirAlias
    alias0, InvAlias
    _) <- VName -> VarAliasTab -> Maybe (VName, DirAlias, InvAlias)
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    x VarAliasTab
    alias_tab = do
          Coalesced knd (MemBlock pt shp vname ixf) substs <- VarAliasTab -> Map VName Coalesced -> VName -> Maybe Coalesced
    walkAliasTab VarAliasTab
    alias_tab Map VName Coalesced
    vtab VName
    x0
          ixf' <- alias0 ixf
          pure $ Coalesced knd (MemBlock pt shp vname ixf') substs
    walkAliasTab VarAliasTab
    _ Map VName Coalesced
    _ VName
    _ = Maybe Coalesced
    forall a. Maybe a
    Nothing

    -- | We assume @x@ is in @vartab@ and we add the variables that @x@ aliases
    --   for as long as possible following a chain of direct-aliasing operators,
    --   i.e., without considering aliasing of if-then-else, loops, etc. For example:
    --     @ x0 = if c then ... else ...@
    --     @ x1 = rearrange r1 x0 @
    --     @ x2 = reverse x1@
    --     @ y[slc] = x2 @
    --   We assume @vartab@ constains a binding for @x2@, and calling this function
    --     with @x2@ as argument should also insert entries for @x1@ and @x0@ to
    --     @vartab@, of course if their aliasing operations are invertible.
    --   We assume inverting aliases has been performed by the top-down pass.
    addInvAliasesVarTab ::
      (HasMemBlock (Aliases rep)) =>
      TopdownEnv rep ->
      M.Map VName Coalesced ->
      VName ->
      Maybe (M.Map VName Coalesced)
    addInvAliasesVarTab :: forall rep.
    HasMemBlock (Aliases rep) =>
    TopdownEnv rep
    -> Map VName Coalesced -> VName -> Maybe (Map VName Coalesced)
    addInvAliasesVarTab TopdownEnv rep
    td_env Map VName Coalesced
    vtab VName
    x
      | Just (Coalesced CoalescedKind
    _ (MemBlock PrimType
    _ ShapeBase SubExp
    _ VName
    m_y LMAD
    x_ixfun) FreeVarSubsts
    fv_subs) <- VName -> Map VName Coalesced -> Maybe Coalesced
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    x Map VName Coalesced
    vtab =
          case VName -> VarAliasTab -> Maybe (VName, DirAlias, InvAlias)
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    x (TopdownEnv rep -> VarAliasTab
    forall rep. TopdownEnv rep -> VarAliasTab
    v_alias TopdownEnv rep
    td_env) of
            Maybe (VName, DirAlias, InvAlias)
    Nothing -> Map VName Coalesced -> Maybe (Map VName Coalesced)
    forall a. a -> Maybe a
    Just Map VName Coalesced
    vtab
            Just (VName
    _, DirAlias
    _, InvAlias
    Nothing) -> Maybe (Map VName Coalesced)
    forall a. Maybe a
    Nothing -- can't invert ixfun, conservatively fail!
            Just (VName
    x0, DirAlias
    _, Just LMAD -> LMAD
    inv_alias0) ->
              let x_ixfn0 :: LMAD
    x_ixfn0 = LMAD -> LMAD
    inv_alias0 LMAD
    x_ixfun
               in case VName -> Scope (Aliases rep) -> Maybe ArrayMemBound
    forall rep.
    HasMemBlock rep =>
    VName -> Scope rep -> Maybe ArrayMemBound
    getScopeMemInfo VName
    x0 (TopdownEnv rep -> Scope (Aliases rep)
    forall rep. TopdownEnv rep -> ScopeTab rep
    scope TopdownEnv rep
    td_env) of
                    Maybe ArrayMemBound
    Nothing -> [Char] -> Maybe (Map VName Coalesced)
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "impossible"
                    Just (MemBlock PrimType
    ptp ShapeBase SubExp
    shp VName
    _ LMAD
    _) ->
                      let coal :: Coalesced
    coal = CoalescedKind -> ArrayMemBound -> FreeVarSubsts -> Coalesced
    Coalesced CoalescedKind
    TransitiveCoal (PrimType -> ShapeBase SubExp -> VName -> LMAD -> ArrayMemBound
    MemBlock PrimType
    ptp ShapeBase SubExp
    shp VName
    m_y LMAD
    x_ixfn0) FreeVarSubsts
    fv_subs
                          vartab' :: Map VName Coalesced
    vartab' = VName -> Coalesced -> Map VName Coalesced -> Map VName Coalesced
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    M.insert VName
    x0 Coalesced
    coal Map VName Coalesced
    vtab
                       in TopdownEnv rep
    -> Map VName Coalesced -> VName -> Maybe (Map VName Coalesced)
    forall rep.
    HasMemBlock (Aliases rep) =>
    TopdownEnv rep
    -> Map VName Coalesced -> VName -> Maybe (Map VName Coalesced)
    addInvAliasesVarTab TopdownEnv rep
    td_env Map VName Coalesced
    vartab' VName
    x0
    addInvAliasesVarTab TopdownEnv rep
    _ Map VName Coalesced
    _ VName
    _ = Maybe (Map VName Coalesced)
    forall a. Maybe a
    Nothing

    areAliased :: TopdownEnv rep -> VName -> VName -> Bool
    areAliased :: forall rep. TopdownEnv rep -> VName -> VName -> Bool
    areAliased TopdownEnv rep
    _ VName
    m_x VName
    m_y =
      -- this is a dummy implementation
      VName
    m_x VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    == VName
    m_y

    areAnyAliased :: TopdownEnv rep -> VName -> [VName] -> Bool
    areAnyAliased :: forall rep. TopdownEnv rep -> VName -> [VName] -> Bool
    areAnyAliased TopdownEnv rep
    td_env VName
    m_x =
      (VName -> Bool) -> [VName] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any (TopdownEnv rep -> VName -> VName -> Bool
    forall rep. TopdownEnv rep -> VName -> VName -> Bool
    areAliased TopdownEnv rep
    td_env VName
    m_x)