{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TypeFamilies #-}

-- | The bulk of the short-circuiting implementation.
module Futhark.Optimise.ArrayShortCircuiting.ArrayCoalescing
  ( mkCoalsTab,
    CoalsTab,
    mkCoalsTabGPU,
    mkCoalsTabMC,
  )
where

import Control.Exception.Base qualified as Exc
import Control.Monad
import Control.Monad.Reader
import Control.Monad.State.Strict
import Data.Function ((&))
import Data.List qualified as L
import Data.List.NonEmpty (NonEmpty (..))
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Sequence (Seq (..))
import Data.Set qualified as S
import Futhark.Analysis.LastUse
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.IR.SeqMem
import Futhark.MonadFreshNames
import Futhark.Optimise.ArrayShortCircuiting.DataStructs
import Futhark.Optimise.ArrayShortCircuiting.MemRefAggreg
import Futhark.Optimise.ArrayShortCircuiting.TopdownAnalysis
import Futhark.Util

-- | A helper type describing representations that can be short-circuited.
type Coalesceable rep inner =
  ( Mem rep inner,
    ASTRep rep,
    CanBeAliased inner,
    AliasableRep rep,
    Op rep ~ MemOp inner rep,
    HasMemBlock (Aliases rep),
    LetDec rep ~ LetDecMem,
    TopDownHelper (inner (Aliases rep))
  )

type ComputeScalarTable rep op =
  ScopeTab rep -> op -> ScalarTableM rep (M.Map VName (PrimExp VName))

-- Helper type for computing scalar tables on ops.
newtype ComputeScalarTableOnOp rep = ComputeScalarTableOnOp
  { scalarTableOnOp :: ComputeScalarTable rep (Op (Aliases rep))
  }

type ScalarTableM rep a = Reader (ComputeScalarTableOnOp rep) a

data ShortCircuitReader rep = ShortCircuitReader
  { onOp ::
      LUTabFun ->
      Pat (VarAliases, LetDecMem) ->
      Certs ->
      Op (Aliases rep) ->
      TopdownEnv rep ->
      BotUpEnv ->
      ShortCircuitM rep BotUpEnv,
    ssPointFromOp ::
      LUTabFun ->
      TopdownEnv rep ->
      ScopeTab rep ->
      Pat (VarAliases, LetDecMem) ->
      Certs ->
      Op (Aliases rep) ->
      Maybe [SSPointInfo]
  }

newtype ShortCircuitM rep a = ShortCircuitM (ReaderT (ShortCircuitReader rep) (State VNameSource) a)
  deriving (Functor, Applicative, Monad, MonadReader (ShortCircuitReader rep), MonadState VNameSource)

instance MonadFreshNames (ShortCircuitM rep) where
  putNameSource = put
  getNameSource = get

emptyTopdownEnv :: TopdownEnv rep
emptyTopdownEnv =
  TopdownEnv
    { alloc = mempty,
      scope = mempty,
      inhibited = mempty,
      v_alias = mempty,
      m_alias = mempty,
      nonNegatives = mempty,
      scalarTable = mempty,
      knownLessThan = mempty,
      td_asserts = mempty
    }

emptyBotUpEnv :: BotUpEnv
emptyBotUpEnv =
  BotUpEnv
    { scals = mempty,
      activeCoals = mempty,
      successCoals = mempty,
      inhibit = mempty
    }

--------------------------------------------------------------------------------
--- Main Coalescing Transformation computes a successful coalescing table    ---
--------------------------------------------------------------------------------

-- | Given a 'Prog' in 'SegMem' representation, compute the coalescing table
-- by folding over each function.
mkCoalsTab :: (MonadFreshNames m) => Prog (Aliases SeqMem) -> m (M.Map Name CoalsTab)
mkCoalsTab prog =
  mkCoalsTabProg
    (lastUseSeqMem prog)
    (ShortCircuitReader shortCircuitSeqMem genSSPointInfoSeqMem)
    (ComputeScalarTableOnOp $ const $ const $ pure mempty)
    prog

-- | Given a 'Prog' in 'GPUMem' representation, compute the coalescing table
-- by folding over each function.
mkCoalsTabGPU :: (MonadFreshNames m) => Prog (Aliases GPUMem) -> m (M.Map Name CoalsTab)
mkCoalsTabGPU prog =
  mkCoalsTabProg
    (lastUseGPUMem prog)
    (ShortCircuitReader shortCircuitGPUMem genSSPointInfoGPUMem)
    (ComputeScalarTableOnOp (computeScalarTableMemOp computeScalarTableGPUMem))
    prog

-- | Given a 'Prog' in 'MCMem' representation, compute the coalescing table
-- by folding over each function.
mkCoalsTabMC :: (MonadFreshNames m) => Prog (Aliases MCMem) -> m (M.Map Name CoalsTab)
mkCoalsTabMC prog =
  mkCoalsTabProg
    (lastUseMCMem prog)
    (ShortCircuitReader shortCircuitMCMem genSSPointInfoMCMem)
    (ComputeScalarTableOnOp (computeScalarTableMemOp computeScalarTableMCMem))
    prog

-- | Given a function, compute the coalescing table
mkCoalsTabProg ::
  (MonadFreshNames m, Coalesceable rep inner) =>
  LUTabProg ->
  ShortCircuitReader rep ->
  ComputeScalarTableOnOp rep ->
  Prog (Aliases rep) ->
  m (M.Map Name CoalsTab)
mkCoalsTabProg (_, lutab_prog) r computeScalarOnOp prog =
  fmap M.fromList . mapM onFun . progFuns $ prog
  where
    consts_scope = scopeOf (progConsts prog)
    onFun fun@(FunDef _ _ fname _ fpars body) = do
      -- First compute last-use information
      let unique_mems = getUniqueMemFParam fpars
          lutab = lutab_prog M.! fname
          scalar_table =
            runReader
              ( concatMapM
                  ( computeScalarTable $
                      consts_scope
                        <> scopeOf fun
                        <> scopeOf (bodyStms body)
                  )
                  (stmsToList $ bodyStms body)
              )
              computeScalarOnOp
          topenv =
            emptyTopdownEnv
              { scope = consts_scope <> scopeOfFParams fpars,
                alloc = unique_mems,
                scalarTable = scalar_table,
                nonNegatives = foldMap paramSizes fpars
              }
          ShortCircuitM m = fixPointCoalesce lutab fpars body topenv
      (fname,) <$> modifyNameSource (runState (runReaderT m r))

paramSizes :: Param FParamMem -> Names
paramSizes (Param _ _ (MemArray _ shp _ _)) = freeIn shp
paramSizes _ = mempty

-- | Short-circuit handler for a 'SeqMem' 'Op'.
--
-- Because 'SeqMem' don't have any special operation, simply return the input
-- 'BotUpEnv'.
shortCircuitSeqMem :: LUTabFun -> Pat (VarAliases, LetDecMem) -> Certs -> Op (Aliases SeqMem) -> TopdownEnv SeqMem -> BotUpEnv -> ShortCircuitM SeqMem BotUpEnv
shortCircuitSeqMem _ _ _ _ _ = pure

-- | Short-circuit handler for SegOp.
shortCircuitSegOp ::
  (Coalesceable rep inner) =>
  (lvl -> Bool) ->
  LUTabFun ->
  Pat (VarAliases, LetDecMem) ->
  Certs ->
  SegOp lvl (Aliases rep) ->
  TopdownEnv rep ->
  BotUpEnv ->
  ShortCircuitM rep BotUpEnv
shortCircuitSegOp lvlOK lutab pat pat_certs (SegMap lvl space _ kernel_body) td_env bu_env =
  -- No special handling necessary for 'SegMap'. Just call the helper-function.
  shortCircuitSegOpHelper 0 lvlOK lvl lutab pat pat_certs space kernel_body td_env bu_env
shortCircuitSegOp lvlOK lutab pat pat_certs (SegRed lvl space binops _ kernel_body) td_env bu_env =
  -- When handling 'SegRed', we we first invalidate all active coalesce-entries
  -- where any of the variables in 'vartab' are also free in the list of
  -- 'SegBinOp'. In other words, anything that is used as part of the reduction
  -- step should probably not be coalesced.
  let to_fail = M.filter (\entry -> namesFromList (M.keys $ vartab entry) `namesIntersect` foldMap (freeIn . segBinOpLambda) binops) $ activeCoals bu_env
      (active, inh) =
        foldl markFailedCoal (activeCoals bu_env, inhibit bu_env) $ M.keys to_fail
      bu_env' = bu_env {activeCoals = active, inhibit = inh}
      num_reds = length red_ts
   in shortCircuitSegOpHelper num_reds lvlOK lvl lutab pat pat_certs space kernel_body td_env bu_env'
  where
    segment_dims = init $ segSpaceDims space
    red_ts = do
      op <- binops
      let shp = Shape segment_dims <> segBinOpShape op
      map (`arrayOfShape` shp) (lambdaReturnType $ segBinOpLambda op)
shortCircuitSegOp lvlOK lutab pat pat_certs (SegScan lvl space binops _ kernel_body) td_env bu_env =
  -- Like in the handling of 'SegRed', we do not want to coalesce anything that
  -- is used in the 'SegBinOp'
  let to_fail = M.filter (\entry -> namesFromList (M.keys $ vartab entry) `namesIntersect` foldMap (freeIn . segBinOpLambda) binops) $ activeCoals bu_env
      (active, inh) = foldl markFailedCoal (activeCoals bu_env, inhibit bu_env) $ M.keys to_fail
      bu_env' = bu_env {activeCoals = active, inhibit = inh}
   in shortCircuitSegOpHelper 0 lvlOK lvl lutab pat pat_certs space kernel_body td_env bu_env'
shortCircuitSegOp lvlOK lutab pat pat_certs (SegHist lvl space histops _ kernel_body) td_env bu_env = do
  -- Need to take zipped patterns and histDest (flattened) and insert transitive coalesces
  let to_fail = M.filter (\entry -> namesFromList (M.keys $ vartab entry) `namesIntersect` foldMap (freeIn . histOp) histops) $ activeCoals bu_env
      (active, inh) = foldl markFailedCoal (activeCoals bu_env, inhibit bu_env) $ M.keys to_fail
      bu_env' = bu_env {activeCoals = active, inhibit = inh}
  bu_env'' <- shortCircuitSegOpHelper 0 lvlOK lvl lutab pat pat_certs space kernel_body td_env bu_env'
  pure $
    foldl insertHistCoals bu_env'' $
      zip (patElems pat) $
        concatMap histDest histops
  where
    insertHistCoals acc (PatElem p _, hist_dest) =
      case ( getScopeMemInfo p $ scope td_env,
             getScopeMemInfo hist_dest $ scope td_env
           ) of
        (Just (MemBlock _ _ p_mem _), Just (MemBlock _ _ dest_mem _)) ->
          case M.lookup p_mem $ successCoals acc of
            Just entry ->
              -- Update this entry with an optdep for the memory block of hist_dest
              let entry' = entry {optdeps = M.insert p p_mem $ optdeps entry}
               in acc
                    { successCoals = M.insert p_mem entry' $ successCoals acc,
                      activeCoals = M.insert dest_mem entry $ activeCoals acc
                    }
            Nothing -> acc
        _ -> acc

-- | Short-circuit handler for 'GPUMem' 'Op'.
--
-- When the 'Op' is a 'SegOp', we handle it accordingly, otherwise we do
-- nothing.
shortCircuitGPUMem ::
  LUTabFun ->
  Pat (VarAliases, LetDecMem) ->
  Certs ->
  Op (Aliases GPUMem) ->
  TopdownEnv GPUMem ->
  BotUpEnv ->
  ShortCircuitM GPUMem BotUpEnv
shortCircuitGPUMem _ _ _ (Alloc _ _) _ bu_env = pure bu_env
shortCircuitGPUMem lutab pat certs (Inner (GPU.SegOp op)) td_env bu_env =
  shortCircuitSegOp isSegThread lutab pat certs op td_env bu_env
shortCircuitGPUMem lutab pat certs (Inner (GPU.GPUBody _ body)) td_env bu_env = do
  fresh1 <- newNameFromString "gpubody"
  fresh2 <- newNameFromString "gpubody"
  shortCircuitSegOpHelper
    0
    isSegThread
    -- Construct a 'SegLevel' corresponding to a single thread
    ( GPU.SegThread GPU.SegNoVirt $
        Just $
          GPU.KernelGrid
            (GPU.Count $ Constant $ IntValue $ Int64Value 1)
            (GPU.Count $ Constant $ IntValue $ Int64Value 1)
    )
    lutab
    pat
    certs
    (SegSpace fresh1 [(fresh2, Constant $ IntValue $ Int64Value 1)])
    (bodyToKernelBody body)
    td_env
    bu_env
shortCircuitGPUMem _ _ _ (Inner (GPU.SizeOp _)) _ bu_env = pure bu_env
shortCircuitGPUMem _ _ _ (Inner (GPU.OtherOp NoOp)) _ bu_env = pure bu_env

shortCircuitMCMem ::
  LUTabFun ->
  Pat (VarAliases, LetDecMem) ->
  Certs ->
  Op (Aliases MCMem) ->
  TopdownEnv MCMem ->
  BotUpEnv ->
  ShortCircuitM MCMem BotUpEnv
shortCircuitMCMem _ _ _ (Alloc _ _) _ bu_env = pure bu_env
shortCircuitMCMem _ _ _ (Inner (MC.OtherOp NoOp)) _ bu_env = pure bu_env
shortCircuitMCMem lutab pat certs (Inner (MC.ParOp (Just par_op) op)) td_env bu_env =
  shortCircuitSegOp (const True) lutab pat certs par_op td_env bu_env
    >>= shortCircuitSegOp (const True) lutab pat certs op td_env
shortCircuitMCMem lutab pat certs (Inner (MC.ParOp Nothing op)) td_env bu_env =
  shortCircuitSegOp (const True) lutab pat certs op td_env bu_env

dropLastSegSpace :: SegSpace -> SegSpace
dropLastSegSpace space = space {unSegSpace = init $ unSegSpace space}

isSegThread :: GPU.SegLevel -> Bool
isSegThread GPU.SegThread {} = True
isSegThread _ = False

-- | Computes the slice written at the end of a thread in a 'SegOp'.
threadSlice :: SegSpace -> KernelResult -> Maybe (Slice (TPrimExp Int64 VName))
threadSlice space Returns {} =
  Just $
    Slice $
      map (DimFix . TPrimExp . flip LeafExp (IntType Int64) . fst) $
        unSegSpace space
threadSlice space (RegTileReturns _ dims _) =
  Just
    $ Slice
    $ zipWith
      ( \(_, block_tile_size0, reg_tile_size0) (x0, _) ->
          let x = pe64 $ Var x0
              block_tile_size = pe64 block_tile_size0
              reg_tile_size = pe64 reg_tile_size0
           in DimSlice (x * block_tile_size * reg_tile_size) (block_tile_size * reg_tile_size) 1
      )
      dims
    $ unSegSpace space
threadSlice _ _ = Nothing

bodyToKernelBody :: Body (Aliases GPUMem) -> KernelBody (Aliases GPUMem)
bodyToKernelBody (Body dec stms res) =
  KernelBody dec stms $ map (\(SubExpRes cert subexps) -> Returns ResultNoSimplify cert subexps) res

-- | A helper for all the different kinds of 'SegOp'.
--
-- Consists of four parts:
--
-- 1. Create coalescing relations between the pattern elements and the kernel
-- body results using 'makeSegMapCoals'.
--
-- 2. Process the statements of the 'KernelBody'.
--
-- 3. Check the overlap between the different threads.
--
-- 4. Mark active coalescings as finished, since a 'SegOp' is an array creation
-- point.
shortCircuitSegOpHelper ::
  (Coalesceable rep inner) =>
  -- | The number of returns for which we should drop the last seg space
  Int ->
  -- | Whether we should look at a segop with this lvl.
  (lvl -> Bool) ->
  lvl ->
  LUTabFun ->
  Pat (VarAliases, LetDecMem) ->
  Certs ->
  SegSpace ->
  KernelBody (Aliases rep) ->
  TopdownEnv rep ->
  BotUpEnv ->
  ShortCircuitM rep BotUpEnv
shortCircuitSegOpHelper num_reds lvlOK lvl lutab pat@(Pat ps0) pat_certs space0 kernel_body td_env bu_env = do
  -- We need to drop the last element of the 'SegSpace' for pattern elements
  -- that correspond to reductions.
  let ps_space_and_res =
        zip3 ps0 (replicate num_reds (dropLastSegSpace space0) <> repeat space0) $
          kernelBodyResult kernel_body
  -- Create coalescing relations between pattern elements and kernel body
  -- results
  let (actv0, inhibit0) =
        filterSafetyCond2and5
          (activeCoals bu_env)
          (inhibit bu_env)
          (scals bu_env)
          td_env
          (patElems pat)
      (actv_return, inhibit_return) =
        if num_reds > 0
          then (actv0, inhibit0)
          else foldl (makeSegMapCoals lvlOK lvl td_env kernel_body pat_certs) (actv0, inhibit0) ps_space_and_res

  -- Start from empty references, we'll update with aggregates later.
  let actv0' = M.map (\etry -> etry {memrefs = mempty}) $ actv0 <> actv_return
  -- Process kernel body statements
  bu_env' <-
    mkCoalsTabStms lutab (kernelBodyStms kernel_body) td_env $
      bu_env {activeCoals = actv0', inhibit = inhibit_return}

  let actv_coals_after =
        M.mapWithKey
          ( \k etry ->
              etry
                { memrefs = memrefs etry <> maybe mempty memrefs (M.lookup k $ actv0 <> actv_return)
                }
          )
          $ activeCoals bu_env'

  -- Check partial overlap.
  let checkPartialOverlap bu_env_f (k, entry) = do
        let sliceThreadAccess (p, space, res) =
              case M.lookup (patElemName p) $ vartab entry of
                Just (Coalesced _ (MemBlock _ _ _ ixf) _) ->
                  maybe
                    Undeterminable
                    ( ixfunToAccessSummary
                        . LMAD.slice ixf
                        . fullSlice (LMAD.shape ixf)
                    )
                    $ threadSlice space res
                Nothing -> mempty
            thread_writes = foldMap sliceThreadAccess ps_space_and_res
            source_writes = srcwrts (memrefs entry) <> thread_writes
        destination_uses <-
          case dstrefs (memrefs entry)
            `accessSubtract` dstrefs (maybe mempty memrefs $ M.lookup k $ activeCoals bu_env) of
            Set s ->
              concatMapM
                (aggSummaryMapPartial (scalarTable td_env) $ unSegSpace space0)
                (S.toList s)
            Undeterminable -> pure Undeterminable
        -- Do not allow short-circuiting from a segop-shared memory
        -- block (not in the topdown scope) to an outer memory block.
        if dstmem entry `M.member` scope td_env
          && noMemOverlap td_env destination_uses source_writes
          then pure bu_env_f
          else do
            let (ac, inh) = markFailedCoal (activeCoals bu_env_f, inhibit bu_env_f) k
            pure $ bu_env_f {activeCoals = ac, inhibit = inh}

  bu_env'' <-
    foldM
      checkPartialOverlap
      (bu_env' {activeCoals = actv_coals_after})
      $ M.toList actv_coals_after

  let updateMemRefs entry = do
        wrts <- aggSummaryMapTotal (scalarTable td_env) (unSegSpace space0) $ srcwrts $ memrefs entry
        uses <- aggSummaryMapTotal (scalarTable td_env) (unSegSpace space0) $ dstrefs $ memrefs entry

        -- Add destination uses from the pattern
        let uses' =
              foldMap
                ( \case
                    PatElem _ (_, MemArray _ _ _ (ArrayIn p_mem p_ixf))
                      | p_mem `nameIn` alsmem entry ->
                          ixfunToAccessSummary p_ixf
                    _ -> mempty
                )
                ps0

        pure $ entry {memrefs = MemRefs (uses <> uses') wrts}

  actv <- mapM updateMemRefs $ activeCoals bu_env''
  let bu_env''' = bu_env'' {activeCoals = actv}

  -- Process pattern and return values
  let mergee_writes =
        mapMaybe
          ( \(p, _, _) ->
              fmap (p,) $
                getDirAliasedIxfn' td_env (activeCoals bu_env''') $
                  patElemName p
          )
          ps_space_and_res

  -- Now, for each mergee write, we need to check that it doesn't overlap with any previous uses of the destination.
  let checkMergeeOverlap bu_env_f (p, (m_b, _, ixf)) =
        let as = ixfunToAccessSummary ixf
         in -- Should be @bu_env@ here, because we need to check overlap
            -- against previous uses.
            case M.lookup m_b $ activeCoals bu_env of
              Just coal_entry -> do
                let mrefs =
                      memrefs coal_entry
                    res = noMemOverlap td_env as $ dstrefs mrefs
                    fail_res =
                      let (ac, inh) = markFailedCoal (activeCoals bu_env_f, inhibit bu_env_f) m_b
                       in bu_env_f {activeCoals = ac, inhibit = inh}

                if res
                  then case M.lookup (patElemName p) $ vartab coal_entry of
                    Nothing -> pure bu_env_f
                    Just (Coalesced knd mbd@(MemBlock _ _ _ ixfn) _) -> pure $
                      case freeVarSubstitutions (scope td_env) (scalarTable td_env) ixfn of
                        Just fv_subst ->
                          let entry =
                                coal_entry
                                  { vartab =
                                      M.insert
                                        (patElemName p)
                                        (Coalesced knd mbd fv_subst)
                                        (vartab coal_entry)
                                  }
                              (ac, suc) =
                                markSuccessCoal (activeCoals bu_env_f, successCoals bu_env_f) m_b entry
                           in bu_env_f {activeCoals = ac, successCoals = suc}
                        Nothing ->
                          fail_res
                  else pure fail_res
              _ -> pure bu_env_f

  foldM checkMergeeOverlap bu_env''' mergee_writes

-- | Given a pattern element and the corresponding kernel result, try to put the
-- kernel result directly in the memory block of pattern element
makeSegMapCoals ::
  (Coalesceable rep inner) =>
  (lvl -> Bool) ->
  lvl ->
  TopdownEnv rep ->
  KernelBody (Aliases rep) ->
  Certs ->
  (CoalsTab, InhibitTab) ->
  (PatElem (VarAliases, LetDecMem), SegSpace, KernelResult) ->
  (CoalsTab, InhibitTab)
makeSegMapCoals lvlOK lvl td_env kernel_body pat_certs (active, inhb) (PatElem pat_name (_, MemArray _ _ _ (ArrayIn pat_mem pat_ixf)), space, Returns _ _ (Var return_name))
  | Just (MemBlock tp return_shp return_mem _) <-
      getScopeMemInfo return_name $ scope td_env <> scopeOf (kernelBodyStms kernel_body),
    lvlOK lvl,
    MemMem pat_space <- runReader (lookupMemInfo pat_mem) $ removeScopeAliases $ scope td_env,
    MemMem return_space <-
      scope td_env <> scopeOf (kernelBodyStms kernel_body) <> scopeOfSegSpace space
        & removeScopeAliases
        & runReader (lookupMemInfo return_mem),
    pat_space == return_space =
      case M.lookup pat_mem active of
        Nothing ->
          -- We are not in a transitive case
          case ( maybe False (pat_mem `nameIn`) (M.lookup return_mem inhb),
                 Coalesced
                   InPlaceCoal
                   (MemBlock tp return_shp pat_mem $ resultSlice pat_ixf)
                   mempty
                   & M.singleton return_name
                   & flip (addInvAliasesVarTab td_env) return_name
               ) of
            (False, Just vtab) ->
              ( active
                  <> M.singleton
                    return_mem
                    (CoalsEntry pat_mem pat_ixf (oneName pat_mem) vtab mempty mempty pat_certs),
                inhb
              )
            _ -> (active, inhb)
        Just trans ->
          case ( maybe False (dstmem trans `nameIn`) $ M.lookup return_mem inhb,
                 let Coalesced _ (MemBlock _ _ trans_mem trans_ixf) _ =
                       fromMaybe (error "Impossible") $ M.lookup pat_name $ vartab trans
                  in Coalesced
                       TransitiveCoal
                       (MemBlock tp return_shp trans_mem $ resultSlice trans_ixf)
                       mempty
                       & M.singleton return_name
                       & flip (addInvAliasesVarTab td_env) return_name
               ) of
            (False, Just vtab) ->
              let opts =
                    if dstmem trans == pat_mem
                      then mempty
                      else M.insert pat_name pat_mem $ optdeps trans
               in ( M.insert
                      return_mem
                      ( CoalsEntry
                          (dstmem trans)
                          (dstind trans)
                          (oneName pat_mem <> alsmem trans)
                          vtab
                          opts
                          mempty
                          (certs trans <> pat_certs)
                      )
                      active,
                    inhb
                  )
            _ -> (active, inhb)
  where
    thread_slice =
      unSegSpace space
        & map (DimFix . TPrimExp . flip LeafExp (IntType Int64) . fst)
        & Slice
    resultSlice ixf = LMAD.slice ixf $ fullSlice (LMAD.shape ixf) thread_slice
makeSegMapCoals _ _ td_env _ _ x (_, _, WriteReturns _ return_name _) =
  case getScopeMemInfo return_name $ scope td_env of
    Just (MemBlock _ _ return_mem _) -> markFailedCoal x return_mem
    Nothing -> error "Should not happen?"
makeSegMapCoals _ _ td_env _ _ x (_, _, result) =
  freeIn result
    & namesToList
    & mapMaybe (flip getScopeMemInfo $ scope td_env)
    & foldr (flip markFailedCoal . memName) x

fullSlice :: [TPrimExp Int64 VName] -> Slice (TPrimExp Int64 VName) -> Slice (TPrimExp Int64 VName)
fullSlice shp (Slice slc) =
  Slice $ slc ++ map (\d -> DimSlice 0 d 1) (drop (length slc) shp)

fixPointCoalesce ::
  (Coalesceable rep inner) =>
  LUTabFun ->
  [Param FParamMem] ->
  Body (Aliases rep) ->
  TopdownEnv rep ->
  ShortCircuitM rep CoalsTab
fixPointCoalesce lutab fpar bdy topenv = do
  buenv <- mkCoalsTabStms lutab (bodyStms bdy) topenv (emptyBotUpEnv {inhibit = inhibited topenv})
  let succ_tab = successCoals buenv
      actv_tab = activeCoals buenv
      inhb_tab = inhibit buenv
      -- Allow short-circuiting function parameters that are unique and have
      -- matching index functions, otherwise mark as failed
      handleFunctionParams (a, i, s) (_, u, MemBlock _ _ m ixf) =
        case (u, M.lookup m a) of
          (Unique, Just entry)
            | dstind entry == ixf,
              Set dst_uses <- dstrefs (memrefs entry),
              dst_uses == mempty ->
                let (a', s') = markSuccessCoal (a, s) m entry
                 in (a', i, s')
          _ ->
            let (a', i') = markFailedCoal (a, i) m
             in (a', i', s)
      (actv_tab', inhb_tab', succ_tab') =
        foldl
          handleFunctionParams
          (actv_tab, inhb_tab, succ_tab)
          $ getArrMemAssocFParam fpar

      (succ_tab'', failed_optdeps) = fixPointFilterDeps succ_tab' M.empty
      inhb_tab'' = M.unionWith (<>) failed_optdeps inhb_tab'
  if not $ M.null actv_tab'
    then error ("COALESCING ROOT: BROKEN INV, active not empty: " ++ show (M.keys actv_tab'))
    else
      if M.null $ inhb_tab'' `M.difference` inhibited topenv
        then pure succ_tab''
        else fixPointCoalesce lutab fpar bdy (topenv {inhibited = inhb_tab''})
  where
    fixPointFilterDeps :: CoalsTab -> InhibitTab -> (CoalsTab, InhibitTab)
    fixPointFilterDeps coaltab inhbtab =
      let (coaltab', inhbtab') = foldl filterDeps (coaltab, inhbtab) (M.keys coaltab)
       in if length (M.keys coaltab) == length (M.keys coaltab')
            then (coaltab', inhbtab')
            else fixPointFilterDeps coaltab' inhbtab'

    filterDeps (coal, inhb) mb
      | not (M.member mb coal) = (coal, inhb)
    filterDeps (coal, inhb) mb
      | Just coal_etry <- M.lookup mb coal =
          let failed = M.filterWithKey (failedOptDep coal) (optdeps coal_etry)
           in if M.null failed
                then (coal, inhb) -- all ok
                else -- optimistic dependencies failed for the current
                -- memblock; extend inhibited mem-block mergings.
                  markFailedCoal (coal, inhb) mb
    filterDeps _ _ = error "In ArrayCoalescing.hs, fun filterDeps, impossible case reached!"
    failedOptDep coal _ mr
      | not (mr `M.member` coal) = True
    failedOptDep coal r mr
      | Just coal_etry <- M.lookup mr coal = not $ r `M.member` vartab coal_etry
    failedOptDep _ _ _ = error "In ArrayCoalescing.hs, fun failedOptDep, impossible case reached!"

-- | Perform short-circuiting on 'Stms'.
mkCoalsTabStms ::
  (Coalesceable rep inner) =>
  LUTabFun ->
  Stms (Aliases rep) ->
  TopdownEnv rep ->
  BotUpEnv ->
  ShortCircuitM rep BotUpEnv
mkCoalsTabStms lutab stms0 = traverseStms stms0
  where
    non_negs_in_pats = foldMap (nonNegativesInPat . stmPat) stms0
    traverseStms Empty _ bu_env = pure bu_env
    traverseStms (stm :<| stms) td_env bu_env = do
      -- Compute @td_env@ top down
      let td_env' = updateTopdownEnv td_env stm
      -- Compute @bu_env@ bottom up
      bu_env' <- traverseStms stms td_env' bu_env
      mkCoalsTabStm lutab stm (td_env' {nonNegatives = nonNegatives td_env' <> non_negs_in_pats}) bu_env'

-- | Array (register) coalescing can have one of three shapes:
--      a) @let y    = copy(b^{lu})@
--      b) @let y    = concat(a, b^{lu})@
--      c) @let y[i] = b^{lu}@
--   The intent is to use the memory block of the left-hand side
--     for the right-hand side variable, meaning to store @b@ in
--     @m_y@ (rather than @m_b@).
--   The following five safety conditions are necessary:
--      1. the right-hand side is lastly-used in the current statement
--      2. the allocation of @m_y@ dominates the creation of @b@
--         ^ relax it by hoisting the allocation of @m_y@
--      3. there is no use of the left-hand side memory block @m_y@
--           during the liveness of @b@, i.e., in between its last use
--           and its creation.
--         ^ relax it by pointwise/interval-based checking
--      4. @b@ is a newly created array, i.e., does not aliases anything
--         ^ relax it to support exitential memory blocks for if-then-else
--      5. the new index function of @b@ corresponding to memory block @m_y@
--           can be translated at the definition of @b@, and the
--           same for all variables aliasing @b@.
--   Observation: during the live range of @b@, @m_b@ can only be used by
--                variables aliased with @b@, because @b@ is newly created.
--                relax it: in case @m_b@ is existential due to an if-then-else
--                          then the checks should be extended to the actual
--                          array-creation points.
mkCoalsTabStm ::
  (Coalesceable rep inner) =>
  LUTabFun ->
  Stm (Aliases rep) ->
  TopdownEnv rep ->
  BotUpEnv ->
  ShortCircuitM rep BotUpEnv
mkCoalsTabStm _ (Let (Pat [pe]) _ e) td_env bu_env
  | Just primexp <- primExpFromExp (vnameToPrimExp (scope td_env) (scals bu_env)) e =
      pure $ bu_env {scals = M.insert (patElemName pe) primexp (scals bu_env)}
mkCoalsTabStm lutab (Let patt _ (Match _ cases defbody _)) td_env bu_env = do
  let pat_val_elms = patElems patt
      -- ToDo: 1. we need to record existential memory blocks in alias table on the top-down pass.
      --       2. need to extend the scope table

      --  i) Filter @activeCoals@ by the 2ND AND 5th safety conditions:
      (activeCoals0, inhibit0) =
        filterSafetyCond2and5
          (activeCoals bu_env)
          (inhibit bu_env)
          (scals bu_env)
          td_env
          pat_val_elms

      -- ii) extend @activeCoals@ by transfering the pattern-elements bindings existent
      --     in @activeCoals@ to the body results of the then and else branches, but only
      --     if the current pattern element can be potentially coalesced and also
      --     if the current pattern element satisfies safety conditions 2 & 5.
      res_mem_def = findMemBodyResult activeCoals0 (scope td_env) pat_val_elms defbody
      res_mem_cases = map (findMemBodyResult activeCoals0 (scope td_env) pat_val_elms . caseBody) cases

      subs_def = mkSubsTab patt $ map resSubExp $ bodyResult defbody
      subs_cases = map (mkSubsTab patt . map resSubExp . bodyResult . caseBody) cases

      actv_def_i = foldl (transferCoalsToBody subs_def) activeCoals0 res_mem_def
      actv_cases_i = zipWith (\subs res -> foldl (transferCoalsToBody subs) activeCoals0 res) subs_cases res_mem_cases

      -- eliminate the original pattern binding of the if statement,
      -- @let x = if y[0,0] > 0 then map (+y[0,0]) a else map (+1) b@
      -- @let y[0] = x@
      -- should succeed because @m_y@ is used before @x@ is created.
      aux ac (MemBodyResult m_b _ _ m_r) = if m_b == m_r then ac else M.delete m_b ac
      actv_def = foldl aux actv_def_i res_mem_def
      actv_cases = zipWith (foldl aux) actv_cases_i res_mem_cases

  -- iii) process the then and else bodies
  res_def <- mkCoalsTabStms lutab (bodyStms defbody) td_env (bu_env {activeCoals = actv_def})
  res_cases <- zipWithM (\c a -> mkCoalsTabStms lutab (bodyStms $ caseBody c) td_env (bu_env {activeCoals = a})) cases actv_cases
  let (actv_def0, succ_def0, inhb_def0) = (activeCoals res_def, successCoals res_def, inhibit res_def)

      -- iv) optimistically mark the pattern succesful:
      ((activeCoals1, inhibit1), successCoals1) =
        foldl
          ( foldfun
              ( (actv_def0, succ_def0)
                  : zip (map activeCoals res_cases) (map successCoals res_cases)
              )
          )
          ((activeCoals0, inhibit0), successCoals bu_env)
          (L.transpose $ res_mem_def : res_mem_cases)

      --  v) unify coalescing results of all branches by taking the union
      --     of all entries in the current/then/else success tables.

      actv_res = foldr (M.intersectionWith unionCoalsEntry) activeCoals1 $ actv_def0 : map activeCoals res_cases

      succ_res = foldr (M.unionWith unionCoalsEntry) successCoals1 $ succ_def0 : map successCoals res_cases

      -- vi) The step of filtering by 3rd safety condition is not
      --       necessary, because we perform index analysis of the
      --       source/destination uses, and they should have been
      --       filtered during the analysis of the then/else bodies.
      inhibit_res =
        M.unionsWith
          (<>)
          ( inhibit1
              : zipWith
                ( \actv inhb ->
                    let failed = M.difference actv $ M.intersectionWith unionCoalsEntry actv activeCoals0
                     in snd $ foldl markFailedCoal (failed, inhb) (M.keys failed)
                )
                (actv_def0 : map activeCoals res_cases)
                (inhb_def0 : map inhibit res_cases)
          )
  pure
    bu_env
      { activeCoals = actv_res,
        successCoals = succ_res,
        inhibit = inhibit_res
      }
  where
    foldfun _ _ [] =
      error "Imposible Case 1!!!"
    foldfun _ ((act, _), _) mem_body_results
      | Nothing <- M.lookup (patMem $ head mem_body_results) act =
          error "Imposible Case 2!!!"
    foldfun
      acc
      ((act, inhb), succc)
      mem_body_results@(MemBodyResult m_b _ _ _ : _)
        | Just info <- M.lookup m_b act,
          Just _ <- zipWithM (M.lookup . bodyMem) mem_body_results $ map snd acc =
            -- Optimistically promote to successful coalescing and append!
            let info' =
                  info
                    { optdeps =
                        foldr
                          (\mbr -> M.insert (bodyName mbr) (bodyMem mbr))
                          (optdeps info)
                          mem_body_results
                    }
                (act', succc') = markSuccessCoal (act, succc) m_b info'
             in ((act', inhb), succc')
    foldfun
      acc
      ((act, inhb), succc)
      mem_body_results@(MemBodyResult m_b _ _ _ : _)
        | Just info <- M.lookup m_b act,
          all ((==) m_b . bodyMem) mem_body_results,
          Just info' <- zipWithM (M.lookup . bodyMem) mem_body_results $ map fst acc =
            -- Treating special case resembling:
            -- @let x0 = map (+1) a                                  @
            -- @let x3 = if cond then let x1 = x0 with [0] <- 2 in x1@
            -- @                 else let x2 = x0 with [1] <- 3 in x2@
            -- @let z[1] = x3                                        @
            -- In this case the result active table should be the union
            -- of the @m_x@ entries of the then and else active tables.
            let info'' =
                  foldl unionCoalsEntry info info'
                act' = M.insert m_b info'' act
             in ((act', inhb), succc)
    foldfun _ ((act, inhb), succc) (mbr : _) =
      -- one of the branches has failed coalescing,
      -- hence remove the coalescing of the result.

      (markFailedCoal (act, inhb) (patMem mbr), succc)
mkCoalsTabStm lutab (Let pat _ (Loop arginis lform body)) td_env bu_env = do
  let pat_val_elms = patElems pat

      --  i) Filter @activeCoals@ by the 2nd, 3rd AND 5th safety conditions. In
      --  other words, for each active coalescing target, the creation of the
      --  array we're trying to merge should happen before the allocation of the
      --  merge target and the index function should be translateable.
      (actv0, inhibit0) =
        filterSafetyCond2and5
          (activeCoals bu_env)
          (inhibit bu_env)
          (scals bu_env)
          td_env
          pat_val_elms
      -- ii) Extend @activeCoals@ by transfering the pattern-elements bindings
      --     existent in @activeCoals@ to the loop-body results, but only if:
      --       (a) the pattern element is a candidate for coalescing,        &&
      --       (b) the pattern element satisfies safety conditions 2 & 5,
      --           (conditions (a) and (b) have already been checked above), &&
      --       (c) the memory block of the corresponding body result is
      --           allocated outside the loop, i.e., non-existential,        &&
      --       (d) the init name is lastly-used in the initialization
      --           of the loop variant.
      --     Otherwise fail and remove from active-coalescing table!
      bdy_ress = bodyResult body
      (patmems, argmems, inimems, resmems) =
        L.unzip4 $
          mapMaybe (mapmbFun actv0) (zip3 pat_val_elms arginis $ map resSubExp bdy_ress) -- td_env'

      -- remove the other pattern elements from the active coalescing table:
      coal_pat_names = namesFromList $ map fst patmems
      (actv1, inhibit1) =
        foldl
          ( \(act, inhb) (b, MemBlock _ _ m_b _) ->
              if b `nameIn` coal_pat_names
                then (act, inhb) -- ok
                else markFailedCoal (act, inhb) m_b -- remove from active
          )
          (actv0, inhibit0)
          (getArrMemAssoc pat)

      -- iii) Process the loop's body.
      --      If the memory blocks of the loop result and loop variant param differ
      --      then make the original memory block of the loop result conflict with
      --      the original memory block of the loop parameter. This is done in
      --      order to prevent the coalescing of @a1@, @a0@, @x@ and @db@ in the
      --      same memory block of @y@ in the example below:
      --      @loop(a1 = a0) = for i < n do @
      --      @    let x = map (stencil a1) (iota n)@
      --      @    let db = copy x          @
      --      @    in db                    @
      --      @let y[0] = a1                @
      --      Meaning the coalescing of @x@ in @let db = copy x@ should fail because
      --      @a1@ appears in the definition of @let x = map (stencil a1) (iota n)@.
      res_mem_bdy = zipWith (\(b, m_b) (r, m_r) -> MemBodyResult m_b b r m_r) patmems resmems
      res_mem_arg = zipWith (\(b, m_b) (r, m_r) -> MemBodyResult m_b b r m_r) patmems argmems
      res_mem_ini = zipWith (\(b, m_b) (r, m_r) -> MemBodyResult m_b b r m_r) patmems inimems

      actv2 =
        let subs_res = mkSubsTab pat $ map resSubExp $ bodyResult body
            actv11 = foldl (transferCoalsToBody subs_res) actv1 res_mem_bdy
            subs_arg = mkSubsTab pat $ map (Var . paramName . fst) arginis
            actv12 = foldl (transferCoalsToBody subs_arg) actv11 res_mem_arg
            subs_ini = mkSubsTab pat $ map snd arginis
         in foldl (transferCoalsToBody subs_ini) actv12 res_mem_ini

      -- The code below adds an aliasing relation to the loop-arg memory
      --   so that to prevent, e.g., the coalescing of an iterative stencil
      --   (you need a buffer for the result and a separate one for the stencil).
      -- @ let b =               @
      -- @    loop (a) for i<N do@
      -- @        stencil a      @
      -- @  ...                  @
      -- @  y[slc_y] = b         @
      -- This should fail coalescing because we are aliasing @m_a@ with
      --   the memory block of the result.
      insertMemAliases tab (MemBodyResult _ _ _ m_r, MemBodyResult _ _ _ m_a) =
        if m_r == m_a
          then tab
          else case M.lookup m_r tab of
            Nothing -> tab
            Just etry ->
              M.insert m_r (etry {alsmem = alsmem etry <> oneName m_a}) tab
      actv3 = foldl insertMemAliases actv2 (zip res_mem_bdy res_mem_arg)
      -- analysing the loop body starts from a null memory-reference set;
      --  the results of the loop body iteration are aggregated later
      actv4 = M.map (\etry -> etry {memrefs = mempty}) actv3
  res_env_body <-
    mkCoalsTabStms
      lutab
      (bodyStms body)
      td_env'
      ( bu_env
          { activeCoals = actv4,
            inhibit = inhibit1
          }
      )
  let scals_loop = scals res_env_body
      (res_actv0, res_succ0, res_inhb0) = (activeCoals res_env_body, successCoals res_env_body, inhibit res_env_body)
      -- iv) Aggregate memory references across loop and filter unsound coalescing
      -- a) Filter the active-table by the FIRST SOUNDNESS condition, namely:
      --     W_i does not overlap with Union_{j=i+1..n} U_j,
      --     where W_i corresponds to the Write set of src mem-block m_b,
      --     and U_j correspond to the uses of the destination
      --     mem-block m_y, in which m_b is coalesced into.
      --     W_i and U_j correspond to the accesses within the loop body.
      mb_loop_idx = mbLoopIndexRange lform
  res_actv1 <- filterMapM1 (loopSoundness1Entry scals_loop mb_loop_idx) res_actv0

  -- b) Update the memory-reference summaries across loop:
  --   W = Union_{i=0..n-1} W_i Union W_{before-loop}
  --   U = Union_{i=0..n-1} U_i Union U_{before-loop}
  res_actv2 <- mapM (aggAcrossLoopEntry (scope td_env' <> scopeOf (bodyStms body)) scals_loop mb_loop_idx) res_actv1

  -- c) check soundness of the successful promotions for:
  --      - the entries that have been promoted to success during the loop-body pass
  --      - for all the entries of active table
  --    Filter the entries by the SECOND SOUNDNESS CONDITION, namely:
  --      Union_{i=1..n-1} W_i does not overlap the before-the-loop uses
  --        of the destination memory block.
  let res_actv3 = M.filterWithKey (loopSoundness2Entry actv3) res_actv2

  let tmp_succ =
        M.filterWithKey (okLookup actv3) $
          M.difference res_succ0 (successCoals bu_env)
      ver_succ = M.filterWithKey (loopSoundness2Entry actv3) tmp_succ
  let suc_fail = M.difference tmp_succ ver_succ
      (res_succ, res_inhb1) = foldl markFailedCoal (res_succ0, res_inhb0) $ M.keys suc_fail
      --
      act_fail = M.difference res_actv0 res_actv3
      (_, res_inhb) = foldl markFailedCoal (res_actv0, res_inhb1) $ M.keys act_fail
      res_actv =
        M.mapWithKey (addBeforeLoop actv3) res_actv3

      -- v) optimistically mark the pattern succesful if there is any chance to succeed
      ((fin_actv1, fin_inhb1), fin_succ1) =
        foldl foldFunOptimPromotion ((res_actv, res_inhb), res_succ) $
          L.zip4 patmems argmems resmems inimems
      (fin_actv2, fin_inhb2) =
        M.foldlWithKey
          ( \acc k _ ->
              if k `nameIn` namesFromList (map (paramName . fst) arginis)
                then markFailedCoal acc k
                else acc
          )
          (fin_actv1, fin_inhb1)
          fin_actv1
  pure bu_env {activeCoals = fin_actv2, successCoals = fin_succ1, inhibit = fin_inhb2}
  where
    allocs_bdy = foldl getAllocs (alloc td_env') $ bodyStms body
    td_env_allocs = td_env' {alloc = allocs_bdy, scope = scope td_env' <> scopeOf (bodyStms body)}
    td_env' = updateTopdownEnvLoop td_env arginis lform
    getAllocs tab (Let (Pat [pe]) _ (Op (Alloc _ sp))) =
      M.insert (patElemName pe) sp tab
    getAllocs tab _ = tab
    okLookup tab m _
      | Just _ <- M.lookup m tab = True
    okLookup _ _ _ = False
    --
    mapmbFun actv0 (patel, (arg, ini), bdyres)
      | b <- patElemName patel,
        (_, MemArray _ _ _ (ArrayIn m_b _)) <- patElemDec patel,
        a <- paramName arg,
        -- Not safe to short-circuit if the index function of this
        -- parameter is variant to the loop.
        not $ any ((`nameIn` freeIn (paramDec arg)) . paramName . fst) arginis,
        Var a0 <- ini,
        Var r <- bdyres,
        Just coal_etry <- M.lookup m_b actv0,
        Just _ <- M.lookup b (vartab coal_etry),
        Just (MemBlock _ _ m_a _) <- getScopeMemInfo a (scope td_env_allocs),
        Just (MemBlock _ _ m_a0 _) <- getScopeMemInfo a0 (scope td_env_allocs),
        Just (MemBlock _ _ m_r _) <- getScopeMemInfo r (scope td_env_allocs),
        Just nms <- M.lookup a lutab,
        a0 `nameIn` nms,
        m_r `elem` M.keys (alloc td_env_allocs) =
          Just ((b, m_b), (a, m_a), (a0, m_a0), (r, m_r))
    mapmbFun _ (_patel, (_arg, _ini), _bdyres) = Nothing
    foldFunOptimPromotion ::
      ((CoalsTab, InhibitTab), CoalsTab) ->
      ((VName, VName), (VName, VName), (VName, VName), (VName, VName)) ->
      ((CoalsTab, InhibitTab), CoalsTab)
    foldFunOptimPromotion ((act, inhb), succc) ((b, m_b), (a, m_a), (_r, m_r), (b_i, m_i))
      | m_r == m_i,
        Just info <- M.lookup m_i act,
        Just vtab_i <- addInvAliasesVarTab td_env (vartab info) b_i =
          Exc.assert
            (m_r == m_b && m_a == m_b)
            ((M.insert m_b (info {vartab = vtab_i}) act, inhb), succc)
      | m_r == m_i =
          Exc.assert
            (m_r == m_b && m_a == m_b)
            (markFailedCoal (act, inhb) m_b, succc)
      | Just info_b0 <- M.lookup m_b act,
        Just info_a0 <- M.lookup m_a act,
        Just info_i <- M.lookup m_i act,
        M.member m_r succc,
        Just vtab_i <- addInvAliasesVarTab td_env (vartab info_i) b_i,
        [Just info_b, Just info_a] <- map translateIxFnInScope [(b, info_b0), (a, info_a0)] =
          let info_b' = info_b {optdeps = M.insert b_i m_i $ optdeps info_b}
              info_a' = info_a {optdeps = M.insert b_i m_i $ optdeps info_a}
              info_i' =
                info_i
                  { optdeps = M.insert b m_b $ optdeps info_i,
                    memrefs = mempty,
                    vartab = vtab_i
                  }
              act' = M.insert m_i info_i' act
              (act1, succc1) =
                foldl
                  (\acc (m, info) -> markSuccessCoal acc m info)
                  (act', succc)
                  [(m_b, info_b'), (m_a, info_a')]
           in -- ToDo: make sure that ixfun translates and update substitutions (?)
              ((act1, inhb), succc1)
    foldFunOptimPromotion ((act, inhb), succc) ((_, m_b), (_a, m_a), (_r, m_r), (_b_i, m_i)) =
      Exc.assert
        (m_r /= m_i)
        (foldl markFailedCoal (act, inhb) [m_b, m_a, m_r, m_i], succc)

    translateIxFnInScope (x, info)
      | Just (Coalesced knd mbd@(MemBlock _ _ _ ixfn) _subs0) <- M.lookup x (vartab info),
        isInScope td_env (dstmem info) =
          let scope_tab =
                scope td_env
                  <> scopeOfFParams (map fst arginis)
           in case freeVarSubstitutions scope_tab (scals bu_env) ixfn of
                Just fv_subst ->
                  Just $ info {vartab = M.insert x (Coalesced knd mbd fv_subst) (vartab info)}
                Nothing -> Nothing
    translateIxFnInScope _ = Nothing
    se0 = intConst Int64 0
    mbLoopIndexRange ::
      LoopForm ->
      Maybe (VName, (TPrimExp Int64 VName, TPrimExp Int64 VName))
    mbLoopIndexRange (WhileLoop _) = Nothing
    mbLoopIndexRange (ForLoop inm _inttp seN) = Just (inm, (pe64 se0, pe64 seN))
    addBeforeLoop actv_bef m_b etry =
      case M.lookup m_b actv_bef of
        Nothing -> etry
        Just etry0 ->
          etry {memrefs = memrefs etry0 <> memrefs etry}
    aggAcrossLoopEntry scope_loop scal_tab idx etry = do
      wrts <-
        aggSummaryLoopTotal (scope td_env) scope_loop scal_tab idx $
          (srcwrts . memrefs) etry
      uses <-
        aggSummaryLoopTotal (scope td_env) scope_loop scal_tab idx $
          (dstrefs . memrefs) etry
      pure $ etry {memrefs = MemRefs uses wrts}
    loopSoundness1Entry scal_tab idx etry = do
      let wrt_i = (srcwrts . memrefs) etry
      use_p <-
        aggSummaryLoopPartial (scal_tab <> scalarTable td_env) idx $
          dstrefs $
            memrefs etry
      pure $ noMemOverlap td_env' wrt_i use_p
    loopSoundness2Entry :: CoalsTab -> VName -> CoalsEntry -> Bool
    loopSoundness2Entry old_actv m_b etry =
      case M.lookup m_b old_actv of
        Nothing -> True
        Just etry0 ->
          let uses_before = (dstrefs . memrefs) etry0
              write_loop = (srcwrts . memrefs) etry
           in noMemOverlap td_env write_loop uses_before

-- The case of in-place update:
--   @let x' = x with slice <- elm@
mkCoalsTabStm lutab stm@(Let pat@(Pat [x']) _ (BasicOp (Update safety x _ _elm))) td_env bu_env
  | [(_, MemBlock _ _ m_x _)] <- getArrMemAssoc pat = do
      -- (a) filter by the 3rd safety for @elm@ and @x'@
      let (actv, inhbt) = recordMemRefUses td_env bu_env stm
          -- (b) if @x'@ is in active coalesced table, then add an entry for @x@ as well
          (actv', inhbt') =
            case M.lookup m_x actv of
              Nothing -> (actv, inhbt)
              Just info ->
                case M.lookup (patElemName x') (vartab info) of
                  Nothing -> markFailedCoal (actv, inhbt) m_x
                  Just (Coalesced k mblk@(MemBlock _ _ _ x_indfun) _) ->
                    case freeVarSubstitutions (scope td_env) (scals bu_env) x_indfun of
                      Just fv_subs
                        | isInScope td_env (dstmem info) ->
                            let coal_etry_x = Coalesced k mblk fv_subs
                                info' =
                                  info
                                    { vartab =
                                        M.insert x coal_etry_x $
                                          M.insert (patElemName x') coal_etry_x (vartab info)
                                    }
                             in (M.insert m_x info' actv, inhbt)
                      _ ->
                        markFailedCoal (actv, inhbt) m_x

      -- (c) this stm is also a potential source for coalescing, so process it
      actv'' <-
        if safety == Unsafe
          then mkCoalsHelper3PatternMatch stm lutab td_env {inhibited = inhbt'} bu_env {activeCoals = actv'}
          else pure actv'
      pure $ bu_env {activeCoals = actv'', inhibit = inhbt'}

-- The case of flat in-place update:
--   @let x' = x with flat-slice <- elm@
mkCoalsTabStm lutab stm@(Let pat@(Pat [x']) _ (BasicOp (FlatUpdate x _ _elm))) td_env bu_env
  | [(_, MemBlock _ _ m_x _)] <- getArrMemAssoc pat = do
      -- (a) filter by the 3rd safety for @elm@ and @x'@
      let (actv, inhbt) = recordMemRefUses td_env bu_env stm
          -- (b) if @x'@ is in active coalesced table, then add an entry for @x@ as well
          (actv', inhbt') =
            case M.lookup m_x actv of
              Nothing -> (actv, inhbt)
              Just info ->
                case M.lookup (patElemName x') (vartab info) of
                  -- this case should not happen, but if it can that
                  -- just fail conservatively
                  Nothing -> markFailedCoal (actv, inhbt) m_x
                  Just (Coalesced k mblk@(MemBlock _ _ _ x_indfun) _) ->
                    case freeVarSubstitutions (scope td_env) (scals bu_env) x_indfun of
                      Just fv_subs
                        | isInScope td_env (dstmem info) ->
                            let coal_etry_x = Coalesced k mblk fv_subs
                                info' =
                                  info
                                    { vartab =
                                        M.insert x coal_etry_x $
                                          M.insert (patElemName x') coal_etry_x (vartab info)
                                    }
                             in (M.insert m_x info' actv, inhbt)
                      _ ->
                        markFailedCoal (actv, inhbt) m_x

      -- (c) this stm is also a potential source for coalescing, so process it
      actv'' <- mkCoalsHelper3PatternMatch stm lutab td_env {inhibited = inhbt'} bu_env {activeCoals = actv'}
      pure $ bu_env {activeCoals = actv'', inhibit = inhbt'}
--
mkCoalsTabStm _ (Let pat _ (BasicOp Update {})) _ _ =
  error $ "In ArrayCoalescing.hs, fun mkCoalsTabStm, illegal pattern for in-place update: " ++ show pat
-- default handling
mkCoalsTabStm lutab stm@(Let pat aux (Op op)) td_env bu_env = do
  -- Process body
  on_op <- asks onOp
  bu_env' <- on_op lutab pat (stmAuxCerts aux) op td_env bu_env
  activeCoals' <- mkCoalsHelper3PatternMatch stm lutab td_env bu_env'
  pure $ bu_env' {activeCoals = activeCoals'}
mkCoalsTabStm lutab stm@(Let pat _ e) td_env bu_env = do
  --   i) Filter @activeCoals@ by the 3rd safety condition:
  --      this is now relaxed by use of LMAD eqs:
  --      the memory referenced in stm are added to memrefs::dstrefs
  --      in corresponding coal-tab entries.
  let (activeCoals', inhibit') = recordMemRefUses td_env bu_env stm

      --  ii) promote any of the entries in @activeCoals@ to @successCoals@ as long as
      --        - this statement defined a variable consumed in a coalesced statement
      --        - and safety conditions 2, 4, and 5 are satisfied.
      --      AND extend @activeCoals@ table for any definition of a variable that
      --      aliases a coalesced variable.
      safe_4 = createsNewArrOK e
      ((activeCoals'', inhibit''), successCoals') =
        foldl (foldfun safe_4) ((activeCoals', inhibit'), successCoals bu_env) (getArrMemAssoc pat)

  -- iii) record a potentially coalesced statement in @activeCoals@
  activeCoals''' <- mkCoalsHelper3PatternMatch stm lutab td_env bu_env {successCoals = successCoals', activeCoals = activeCoals''}
  pure bu_env {activeCoals = activeCoals''', inhibit = inhibit'', successCoals = successCoals'}
  where
    foldfun safe_4 ((a_acc, inhb), s_acc) (b, MemBlock tp shp mb _b_indfun) =
      case M.lookup mb a_acc of
        Nothing -> ((a_acc, inhb), s_acc)
        Just info@(CoalsEntry x_mem _ _ vtab _ _ certs) ->
          let failed = markFailedCoal (a_acc, inhb) mb
           in case M.lookup b vtab of
                Nothing ->
                  -- we hit the definition of some variable @b@ aliased with
                  --    the coalesced variable @x@, hence extend @activeCoals@, e.g.,
                  --       @let x = map f arr  @
                  --       @let b = alias x  @ <- current statement
                  --       @ ... use of b ...  @
                  --       @let c = alias b    @ <- currently fails
                  --       @let y[i] = x       @
                  -- where @alias@ can be @transpose@, @slice@, @reshape@.
                  -- We use getTransitiveAlias helper function to track the aliasing
                  --    through the td_env, and to find the updated ixfun of @b@:
                  case getDirAliasedIxfn td_env a_acc b of
                    Nothing -> (failed, s_acc)
                    Just (_, _, b_indfun') ->
                      case ( freeVarSubstitutions (scope td_env) (scals bu_env) b_indfun',
                             freeVarSubstitutions (scope td_env) (scals bu_env) certs
                           ) of
                        (Just fv_subst, Just fv_subst') ->
                          let mem_info = Coalesced TransitiveCoal (MemBlock tp shp x_mem b_indfun') (fv_subst <> fv_subst')
                              info' = info {vartab = M.insert b mem_info vtab}
                           in ((M.insert mb info' a_acc, inhb), s_acc)
                        _ -> (failed, s_acc)
                Just (Coalesced k mblk@(MemBlock _ _ _ new_indfun) _) ->
                  -- we are at the definition of the coalesced variable @b@
                  -- if 2,4,5 hold promote it to successful coalesced table,
                  -- or if e = transpose, etc. then postpone decision for later on
                  let safe_2 = isInScope td_env x_mem
                   in case ( freeVarSubstitutions (scope td_env) (scals bu_env) new_indfun,
                             freeVarSubstitutions (scope td_env) (scals bu_env) certs
                           ) of
                        (Just fv_subst, Just fv_subst')
                          | safe_2 ->
                              let mem_info = Coalesced k mblk (fv_subst <> fv_subst')
                                  info' = info {vartab = M.insert b mem_info vtab}
                               in if safe_4
                                    then -- array creation point, successful coalescing verified!

                                      let (a_acc', s_acc') = markSuccessCoal (a_acc, s_acc) mb info'
                                       in ((a_acc', inhb), s_acc')
                                    else -- this is an invertible alias case of the kind
                                    -- @ let b    = alias a @
                                    -- @ let x[i] = b @
                                    -- do not promote, but update the index function

                                      ((M.insert mb info' a_acc, inhb), s_acc)
                        _ -> (failed, s_acc) -- fail!

ixfunToAccessSummary :: LMAD.LMAD (TPrimExp Int64 VName) -> AccessSummary
ixfunToAccessSummary = Set . S.singleton

-- | Check safety conditions 2 and 5 and update new substitutions:
-- called on the pat-elements of loop and if-then-else expressions.
--
-- The safety conditions are: The allocation of merge target should dominate the
-- creation of the array we're trying to merge and the new index function of the
-- array can be translated at the definition site of b. The latter requires that
-- any variables used in the index function of the target array are available at
-- the definition site of b.
filterSafetyCond2and5 ::
  (HasMemBlock (Aliases rep)) =>
  CoalsTab ->
  InhibitTab ->
  ScalarTab ->
  TopdownEnv rep ->
  [PatElem (VarAliases, LetDecMem)] ->
  (CoalsTab, InhibitTab)
filterSafetyCond2and5 act_coal inhb_coal scals_env td_env pes =
  foldl helper (act_coal, inhb_coal) pes
  where
    helper (acc, inhb) patel = do
      -- For each pattern element in the input list
      case (patElemName patel, patElemDec patel) of
        (b, (_, MemArray tp0 shp0 _ (ArrayIn m_b _idxfn_b))) ->
          -- If it is an array in memory block m_b
          case M.lookup m_b acc of
            Nothing -> (acc, inhb)
            Just info@(CoalsEntry x_mem _ _ vtab _ _ certs) ->
              -- And m_b we're trying to coalesce m_b
              let failed = markFailedCoal (acc, inhb) m_b
               in -- It is not safe to short circuit if some other pattern
                  -- element is aliased to this one, as this indicates the
                  -- two pattern elements reference the same physical
                  -- value somehow.
                  if any ((`nameIn` aliasesOf patel) . patElemName) pes
                    then failed
                    else case M.lookup b vtab of
                      Nothing ->
                        case getDirAliasedIxfn td_env acc b of
                          Nothing -> failed
                          Just (_, _, b_indfun') ->
                            -- And we have the index function of b
                            case ( freeVarSubstitutions (scope td_env) scals_env b_indfun',
                                   freeVarSubstitutions (scope td_env) scals_env certs
                                 ) of
                              (Just fv_subst, Just fv_subst') ->
                                let mem_info = Coalesced TransitiveCoal (MemBlock tp0 shp0 x_mem b_indfun') (fv_subst <> fv_subst')
                                    info' = info {vartab = M.insert b mem_info vtab}
                                 in (M.insert m_b info' acc, inhb)
                              _ -> failed
                      Just (Coalesced k (MemBlock pt shp _ new_indfun) _) ->
                        let safe_2 = isInScope td_env x_mem
                         in case ( freeVarSubstitutions (scope td_env) scals_env new_indfun,
                                   freeVarSubstitutions (scope td_env) scals_env certs
                                 ) of
                              (Just fv_subst, Just fv_subst')
                                | safe_2 ->
                                    let mem_info = Coalesced k (MemBlock pt shp x_mem new_indfun) (fv_subst <> fv_subst')
                                        info' = info {vartab = M.insert b mem_info vtab}
                                     in (M.insert m_b info' acc, inhb)
                              _ -> failed
        _ -> (acc, inhb)

-- |   Pattern matches a potentially coalesced statement and
--     records a new association in @activeCoals@
mkCoalsHelper3PatternMatch ::
  (Coalesceable rep inner) =>
  Stm (Aliases rep) ->
  LUTabFun ->
  TopdownEnv rep ->
  BotUpEnv ->
  ShortCircuitM rep CoalsTab
mkCoalsHelper3PatternMatch stm lutab td_env bu_env = do
  clst <- genCoalStmtInfo lutab td_env (scope td_env) stm
  case clst of
    Nothing -> pure activeCoals_tab
    Just clst' -> pure $ foldl processNewCoalesce activeCoals_tab clst'
  where
    successCoals_tab = successCoals bu_env
    activeCoals_tab = activeCoals bu_env
    processNewCoalesce acc (knd, alias_fn, x, m_x, ind_x, b, m_b, _, tp_b, shp_b, certs) =
      -- test whether we are in a transitive coalesced case, i.e.,
      --      @let b = scratch ...@
      --      @.....@
      --      @let x[j] = b@
      --      @let y[i] = x@
      -- and compose the index function of @x@ with that of @y@,
      -- and update aliasing of the @m_b@ entry to also contain @m_y@
      -- on top of @m_x@, i.e., transitively, any use of @m_y@ should
      -- be checked for the lifetime of @b@.
      let proper_coals_tab = case knd of
            InPlaceCoal -> activeCoals_tab
            _ -> successCoals_tab
          (m_yx, ind_yx, mem_yx_al, x_deps, certs') =
            case M.lookup m_x proper_coals_tab of
              Nothing ->
                (m_x, alias_fn ind_x, oneName m_x, M.empty, mempty)
              Just (CoalsEntry m_y ind_y y_al vtab x_deps0 _ certs'') ->
                let ind = case M.lookup x vtab of
                      Just (Coalesced _ (MemBlock _ _ _ ixf) _) ->
                        ixf
                      Nothing ->
                        ind_y
                 in (m_y, alias_fn ind, oneName m_x <> y_al, x_deps0, certs <> certs'')
          m_b_aliased_m_yx = areAnyAliased td_env m_b [m_yx] -- m_b \= m_yx
       in if not m_b_aliased_m_yx && isInScope td_env m_yx -- nameIn m_yx (alloc td_env)
      -- Finally update the @activeCoals@ table with a fresh
      --   binding for @m_b@; if such one exists then overwrite.
      -- Also, add all variables from the alias chain of @b@ to
      --   @vartab@, for example, in the case of a sequence:
      --   @ b0 = if cond then ... else ... @
      --   @ b1 = alias0 b0 @
      --   @ b  = alias1 b1 @
      --   @ x[j] = b @
      -- Then @b1@ and @b0@ should also be added to @vartab@ if
      --   @alias1@ and @alias0@ are invertible, otherwise fail early!
            then
              let mem_info = Coalesced knd (MemBlock tp_b shp_b m_yx ind_yx) M.empty
                  opts' =
                    if m_yx == m_x
                      then M.empty
                      else M.insert x m_x x_deps
                  vtab = M.singleton b mem_info
                  mvtab = addInvAliasesVarTab td_env vtab b

                  is_inhibited = case M.lookup m_b $ inhibited td_env of
                    Just nms -> m_yx `nameIn` nms
                    Nothing -> False
               in case (is_inhibited, mvtab) of
                    (True, _) -> acc -- fail due to inhibited
                    (_, Nothing) -> acc -- fail early due to non-invertible aliasing
                    (_, Just vtab') ->
                      -- successfully adding a new coalesced entry
                      let coal_etry =
                            CoalsEntry
                              m_yx
                              ind_yx
                              mem_yx_al
                              vtab'
                              opts'
                              mempty
                              (certs <> certs')
                       in M.insert m_b coal_etry acc
            else acc

-- | Information about a particular short-circuit point
type SSPointInfo =
  ( CoalescedKind,
    LMAD -> LMAD,
    VName,
    VName,
    LMAD,
    VName,
    VName,
    LMAD,
    PrimType,
    Shape,
    Certs
  )

-- | Given an op, return a list of potential short-circuit points
type GenSSPoint rep op =
  LUTabFun ->
  TopdownEnv rep ->
  ScopeTab rep ->
  Pat (VarAliases, LetDecMem) ->
  Certs ->
  op ->
  Maybe [SSPointInfo]

genSSPointInfoSeqMem ::
  GenSSPoint SeqMem (Op (Aliases SeqMem))
genSSPointInfoSeqMem _ _ _ _ _ _ =
  Nothing

-- | For 'SegOp', we currently only handle 'SegMap', and only under the following
-- circumstances:
--
--  1. The 'SegMap' has only one return/pattern value, which is a 'Returns'.
--
--  2. The 'KernelBody' contains an 'Index' statement that is indexing an array using
--  only the values from the 'SegSpace'.
--
--  3. The array being indexed is last-used in that statement, is free in the
--  'SegMap', is unique or has been recently allocated (specifically, it should
--  not be a non-unique argument to the enclosing function), has elements with
--  the same bit-size as the pattern elements, and has the exact same 'LMAD' as
--  the pattern of the 'SegMap' statement.
--
-- There can be multiple candidate arrays, but the current implementation will
-- always just try the first one.
--
-- The first restriction could be relaxed by trying to match up arrays in the
-- 'KernelBody' with patterns of the 'SegMap', but the current implementation
-- should be enough to handle many common cases.
--
-- The result of the 'SegMap' is treated as the destination, while the candidate
-- array from inside the body is treated as the source.
genSSPointInfoSegOp ::
  (Coalesceable rep inner) => GenSSPoint rep (SegOp lvl (Aliases rep))
genSSPointInfoSegOp
  lutab
  td_env
  scopetab
  (Pat [PatElem dst (_, MemArray dst_pt _ _ (ArrayIn dst_mem dst_ixf))])
  certs
  (SegMap _ space _ kernel_body@KernelBody {kernelBodyResult = [Returns {}]})
    | (src, MemBlock src_pt shp src_mem src_ixf) : _ <-
        mapMaybe getPotentialMapShortCircuit $
          stmsToList $
            kernelBodyStms kernel_body =
        Just [(MapCoal, id, dst, dst_mem, dst_ixf, src, src_mem, src_ixf, src_pt, shp, certs)]
    where
      iterators = map fst $ unSegSpace space
      frees = freeIn kernel_body

      getPotentialMapShortCircuit (Let (Pat [PatElem x _]) _ (BasicOp (Index src slc)))
        | Just inds <- sliceIndices slc,
          L.sort inds == L.sort (map Var iterators),
          Just last_uses <- M.lookup x lutab,
          src `nameIn` last_uses,
          Just memblock@(MemBlock src_pt _ src_mem src_ixf) <-
            getScopeMemInfo src scopetab,
          src_mem `nameIn` last_uses,
          -- The 'alloc' table contains allocated memory blocks, including
          -- unique memory blocks from the enclosing function. It does _not_
          -- include non-unique memory blocks from the enclosing function.
          src_mem `M.member` alloc td_env,
          src `nameIn` frees,
          src_ixf == dst_ixf,
          primBitSize src_pt == primBitSize dst_pt =
            Just (src, memblock)
      getPotentialMapShortCircuit _ = Nothing
genSSPointInfoSegOp _ _ _ _ _ _ =
  Nothing

genSSPointInfoMemOp ::
  GenSSPoint rep (inner (Aliases rep)) ->
  GenSSPoint rep (MemOp inner (Aliases rep))
genSSPointInfoMemOp onOp lutab td_end scopetab pat certs (Inner op) =
  onOp lutab td_end scopetab pat certs op
genSSPointInfoMemOp _ _ _ _ _ _ _ = Nothing

genSSPointInfoGPUMem ::
  GenSSPoint GPUMem (Op (Aliases GPUMem))
genSSPointInfoGPUMem = genSSPointInfoMemOp f
  where
    f lutab td_env scopetab pat certs (GPU.SegOp op) =
      genSSPointInfoSegOp lutab td_env scopetab pat certs op
    f _ _ _ _ _ _ = Nothing

genSSPointInfoMCMem ::
  GenSSPoint MCMem (Op (Aliases MCMem))
genSSPointInfoMCMem = genSSPointInfoMemOp f
  where
    f lutab td_env scopetab pat certs (MC.ParOp Nothing op) =
      genSSPointInfoSegOp lutab td_env scopetab pat certs op
    f _ _ _ _ _ _ = Nothing

genCoalStmtInfo ::
  (Coalesceable rep inner) =>
  LUTabFun ->
  TopdownEnv rep ->
  ScopeTab rep ->
  Stm (Aliases rep) ->
  ShortCircuitM rep (Maybe [SSPointInfo])
-- CASE a) @let x <- copy(b^{lu})@
genCoalStmtInfo lutab td_env scopetab (Let pat aux (BasicOp (Replicate (Shape []) (Var b))))
  | Pat [PatElem x (_, MemArray _ _ _ (ArrayIn m_x ind_x))] <- pat,
    Just last_uses <- M.lookup x lutab,
    Just (MemBlock tpb shpb m_b ind_b) <- getScopeMemInfo b scopetab,
    sameSpace td_env m_x m_b,
    b `nameIn` last_uses =
      pure $ Just [(CopyCoal, id, x, m_x, ind_x, b, m_b, ind_b, tpb, shpb, stmAuxCerts aux)]
-- CASE c) @let x[i] = b^{lu}@
genCoalStmtInfo lutab td_env scopetab (Let pat aux (BasicOp (Update _ x slice_x (Var b))))
  | Pat [PatElem x' (_, MemArray _ _ _ (ArrayIn m_x ind_x))] <- pat,
    Just last_uses <- M.lookup x' lutab,
    Just (MemBlock tpb shpb m_b ind_b) <- getScopeMemInfo b scopetab,
    sameSpace td_env m_x m_b,
    b `nameIn` last_uses =
      pure $ Just [(InPlaceCoal, (`updateIndFunSlice` slice_x), x, m_x, ind_x, b, m_b, ind_b, tpb, shpb, stmAuxCerts aux)]
  where
    updateIndFunSlice :: LMAD -> Slice SubExp -> LMAD
    updateIndFunSlice ind_fun slc_x =
      let slc_x' = map (fmap pe64) $ unSlice slc_x
       in LMAD.slice ind_fun $ Slice slc_x'
genCoalStmtInfo lutab td_env scopetab (Let pat aux (BasicOp (FlatUpdate x slice_x b)))
  | Pat [PatElem x' (_, MemArray _ _ _ (ArrayIn m_x ind_x))] <- pat,
    Just last_uses <- M.lookup x' lutab,
    Just (MemBlock tpb shpb m_b ind_b) <- getScopeMemInfo b scopetab,
    sameSpace td_env m_x m_b,
    b `nameIn` last_uses =
      pure $ Just [(InPlaceCoal, (`updateIndFunSlice` slice_x), x, m_x, ind_x, b, m_b, ind_b, tpb, shpb, stmAuxCerts aux)]
  where
    updateIndFunSlice :: LMAD -> FlatSlice SubExp -> LMAD
    updateIndFunSlice ind_fun (FlatSlice offset dims) =
      LMAD.flatSlice ind_fun $ FlatSlice (pe64 offset) $ map (fmap pe64) dims

-- CASE b) @let x = concat(a, b^{lu})@
genCoalStmtInfo lutab td_env scopetab (Let pat aux (BasicOp (Concat concat_dim (b0 :| bs) _)))
  | Pat [PatElem x (_, MemArray _ _ _ (ArrayIn m_x ind_x))] <- pat,
    Just last_uses <- M.lookup x lutab =
      pure $
        let (res, _, _) = foldl (markConcatParts last_uses x m_x ind_x) ([], zero, True) (b0 : bs)
         in if null res then Nothing else Just res
  where
    zero = pe64 $ intConst Int64 0
    markConcatParts _ _ _ _ acc@(_, _, False) _ = acc
    markConcatParts last_uses x m_x ind_x (acc, offs, True) b
      | Just (MemBlock tpb shpb@(Shape dims@(_ : _)) m_b ind_b) <- getScopeMemInfo b scopetab,
        Just d <- maybeNth concat_dim dims,
        offs' <- offs + pe64 d =
          if b `nameIn` last_uses && sameSpace td_env m_x m_b
            then
              let slc =
                    Slice $
                      map (unitSlice zero . pe64) (take concat_dim dims)
                        <> [unitSlice offs (pe64 d)]
                        <> map (unitSlice zero . pe64) (drop (concat_dim + 1) dims)
               in ( acc ++ [(ConcatCoal, (`LMAD.slice` slc), x, m_x, ind_x, b, m_b, ind_b, tpb, shpb, stmAuxCerts aux)],
                    offs',
                    True
                  )
            else (acc, offs', True)
      | otherwise = (acc, offs, False)
-- case d) short-circuit points from ops. For instance, the result of a segmap
-- can be considered a short-circuit point.
genCoalStmtInfo lutab td_env scopetab (Let pat aux (Op op)) = do
  ss_op <- asks ssPointFromOp
  pure $ ss_op lutab td_env scopetab pat (stmAuxCerts aux) op
-- CASE other than a), b), c), or d) not supported
genCoalStmtInfo _ _ _ _ = pure Nothing

sameSpace :: (Coalesceable rep inner) => TopdownEnv rep -> VName -> VName -> Bool
sameSpace td_env m_x m_b
  | Just (MemMem pat_space) <- nameInfoToMemInfo <$> M.lookup m_x scope',
    Just (MemMem return_space) <- nameInfoToMemInfo <$> M.lookup m_b scope' =
      pat_space == return_space
  | otherwise = False
  where
    scope' = removeScopeAliases $ scope td_env

data MemBodyResult = MemBodyResult
  { patMem :: VName,
    _patName :: VName,
    bodyName :: VName,
    bodyMem :: VName
  }

-- | Results in pairs of pattern-blockresult pairs of (var name, mem block)
--   for those if-patterns that are candidates for coalescing.
findMemBodyResult ::
  (HasMemBlock (Aliases rep)) =>
  CoalsTab ->
  ScopeTab rep ->
  [PatElem (VarAliases, LetDecMem)] ->
  Body (Aliases rep) ->
  [MemBodyResult]
findMemBodyResult activeCoals_tab scope_env patelms bdy =
  mapMaybe
    findMemBodyResult'
    (zip patelms $ map resSubExp $ bodyResult bdy)
  where
    scope_env' = scope_env <> scopeOf (bodyStms bdy)
    findMemBodyResult' (patel, se_r) =
      case (patElemName patel, patElemDec patel, se_r) of
        (b, (_, MemArray _ _ _ (ArrayIn m_b _)), Var r) ->
          case getScopeMemInfo r scope_env' of
            Nothing -> Nothing
            Just (MemBlock _ _ m_r _) ->
              case M.lookup m_b activeCoals_tab of
                Nothing -> Nothing
                Just coal_etry ->
                  case M.lookup b (vartab coal_etry) of
                    Nothing -> Nothing
                    Just _ -> Just $ MemBodyResult m_b b r m_r
        _ -> Nothing

-- | transfers coalescing from if-pattern to then|else body result
--   in the active coalesced table. The transfer involves, among
--   others, inserting @(r,m_r)@ in the optimistically-dependency
--   set of @m_b@'s entry and inserting @(b,m_b)@ in the opt-deps
--   set of @m_r@'s entry. Meaning, ultimately, @m_b@ can be merged
--   if @m_r@ can be merged (and vice-versa). This is checked by a
--   fix point iteration at the function-definition level.
transferCoalsToBody ::
  M.Map VName (TPrimExp Int64 VName) -> -- (PrimExp VName)
  CoalsTab ->
  MemBodyResult ->
  CoalsTab
transferCoalsToBody exist_subs activeCoals_tab (MemBodyResult m_b b r m_r)
  | -- the @Nothing@ pattern for the two lookups cannot happen
    -- because they were already cheked in @findMemBodyResult@
    Just etry <- M.lookup m_b activeCoals_tab,
    Just (Coalesced knd (MemBlock btp shp _ ind_b) subst_b) <- M.lookup b $ vartab etry =
      -- by definition of if-stmt, r and b have the same basic type, shape and
      -- index function, hence, for example, do not need to rebase
      -- We will check whether it is translatable at the definition point of r.
      let ind_r = LMAD.substitute exist_subs ind_b
          subst_r = M.union exist_subs subst_b
          mem_info = Coalesced knd (MemBlock btp shp (dstmem etry) ind_r) subst_r
       in if m_r == m_b -- already unified, just add binding for @r@
            then
              let etry' =
                    etry
                      { optdeps = M.insert b m_b (optdeps etry),
                        vartab = M.insert r mem_info (vartab etry)
                      }
               in M.insert m_r etry' activeCoals_tab
            else -- make them both optimistically depend on each other

              let opts_x_new = M.insert r m_r (optdeps etry)
                  -- Here we should translate the @ind_b@ field of @mem_info@
                  -- across the existential introduced by the if-then-else
                  coal_etry =
                    etry
                      { vartab = M.singleton r mem_info,
                        optdeps = M.insert b m_b (optdeps etry)
                      }
               in M.insert m_b (etry {optdeps = opts_x_new}) $
                    M.insert m_r coal_etry activeCoals_tab
  | otherwise = error "Impossible"

mkSubsTab ::
  Pat (aliases, LetDecMem) ->
  [SubExp] ->
  M.Map VName (TPrimExp Int64 VName)
mkSubsTab pat res =
  let pat_elms = patElems pat
   in M.fromList $ mapMaybe mki64subst $ zip pat_elms res
  where
    mki64subst (a, Var v)
      | (_, MemPrim (IntType Int64)) <- patElemDec a = Just (patElemName a, le64 v)
    mki64subst (a, se@(Constant (IntValue (Int64Value _)))) = Just (patElemName a, pe64 se)
    mki64subst _ = Nothing

computeScalarTable ::
  (Coalesceable rep inner) =>
  ScopeTab rep ->
  Stm (Aliases rep) ->
  ScalarTableM rep (M.Map VName (PrimExp VName))
computeScalarTable scope_table (Let (Pat [pe]) _ e)
  | Just primexp <- primExpFromExp (vnameToPrimExp scope_table mempty) e =
      pure $ M.singleton (patElemName pe) primexp
computeScalarTable scope_table (Let _ _ (Loop loop_inits loop_form body)) =
  concatMapM
    ( computeScalarTable $
        scope_table
          <> scopeOfFParams (map fst loop_inits)
          <> scopeOfLoopForm loop_form
          <> scopeOf (bodyStms body)
    )
    (stmsToList $ bodyStms body)
computeScalarTable scope_table (Let _ _ (Match _ cases body _)) = do
  body_tab <- concatMapM (computeScalarTable $ scope_table <> scopeOf (bodyStms body)) (stmsToList $ bodyStms body)
  cases_tab <-
    concatMapM
      ( \(Case _ b) ->
          concatMapM
            (computeScalarTable $ scope_table <> scopeOf (bodyStms b))
            ( stmsToList $
                bodyStms body
            )
      )
      cases
  pure $ body_tab <> cases_tab
computeScalarTable scope_table (Let _ _ (Op op)) = do
  on_op <- asks scalarTableOnOp
  on_op scope_table op
computeScalarTable _ _ = pure mempty

computeScalarTableMemOp ::
  ComputeScalarTable rep (inner (Aliases rep)) -> ComputeScalarTable rep (MemOp inner (Aliases rep))
computeScalarTableMemOp _ _ (Alloc _ _) = pure mempty
computeScalarTableMemOp onInner scope_table (Inner op) = onInner scope_table op

computeScalarTableSegOp ::
  (Coalesceable rep inner) =>
  ComputeScalarTable rep (GPU.SegOp lvl (Aliases rep))
computeScalarTableSegOp scope_table segop = do
  concatMapM
    ( computeScalarTable $
        scope_table
          <> scopeOf (kernelBodyStms $ segBody segop)
          <> scopeOfSegSpace (segSpace segop)
    )
    (stmsToList $ kernelBodyStms $ segBody segop)

computeScalarTableGPUMem ::
  ComputeScalarTable GPUMem (GPU.HostOp NoOp (Aliases GPUMem))
computeScalarTableGPUMem scope_table (GPU.SegOp segop) =
  computeScalarTableSegOp scope_table segop
computeScalarTableGPUMem _ (GPU.SizeOp _) = pure mempty
computeScalarTableGPUMem _ (GPU.OtherOp NoOp) = pure mempty
computeScalarTableGPUMem scope_table (GPU.GPUBody _ body) =
  concatMapM
    (computeScalarTable $ scope_table <> scopeOf (bodyStms body))
    (stmsToList $ bodyStms body)

computeScalarTableMCMem ::
  ComputeScalarTable MCMem (MC.MCOp NoOp (Aliases MCMem))
computeScalarTableMCMem _ (MC.OtherOp NoOp) = pure mempty
computeScalarTableMCMem scope_table (MC.ParOp par_op segop) =
  (<>)
    <$> maybe (pure mempty) (computeScalarTableSegOp scope_table) par_op
    <*> computeScalarTableSegOp scope_table segop

filterMapM1 :: (Eq k, Monad m) => (v -> m Bool) -> M.Map k v -> m (M.Map k v)
filterMapM1 f m = fmap M.fromAscList $ filterM (f . snd) $ M.toAscList m
