-- Copyright 2022 Google LLC
--
-- Use of this source code is governed by a BSD-style
-- license that can be found in the LICENSE file or at
-- https://developers.google.com/open-source/licenses/bsd

{-# LANGUAGE StrictData #-}
{-# LANGUAGE UndecidableInstances #-}

module Occurrence where

import Prelude hiding (abs, iterate, max, sum)
import Prelude qualified

import Control.Monad.State.Strict
import Data.Hashable
import Data.IntMap.Strict qualified as M
import Data.List (foldl')
import Data.Store (Store (..))
import GHC.Generics (Generic (..))

import PPrint
import IRVariants
import Name

-- === Use Analysis ===

-- We want to know how a binding is used in a Dex program, so as to
-- know whether inlining it is guaranteed to conserve work and code
-- size.

-- Conservation of code size is easy.  Coservation of work, however,
-- is more subtle.  Inlining the computation of x into
--   for i. x
-- does not conserve work (unless the index set has size 0 or 1), but
-- inlining x into
--   for i. x.i
-- does conserve work -- provided we have access to the loop body
-- that constructs it so we can beta-reduce against the indexing
-- operation.

-- The main difficulty is that Dex allows indexing by expressions
-- other than referencing a loop binder, so to solve subtle cases we
-- need to translate iteration-space expressions into data-space.
-- For example, inlining x into
--   for i. x.(Left i)
--   for j. x.(Right j)
-- conserves work, because the two injections have disjoint ranges.

-- We proceed with a four-phase translation.
--
-- - Phase 1 translates Dex source into the Access structure.  An
--   Access mirrors the control pattern of the Dex source, abstracting
--   away most Dex syntax but still referring to iteration space.
--   Access retains nearly complete usage information.
--
-- - Phase 2 transposes Access into Use by turning iteration space
--   into data space.  (Use is like Access, except organized in data
--   space.)  Use retains about the same amount of usage information
--   as Access.
--
-- - Phase 3 collapses Use into CollapsedUse.  This process normalizes
--   the subsets of the data space that are actually accessed, such as
--   realizing that the subsets of all `Left` values and all `Right`
--   values cannot intersect.  This is where most of the
--   overapproximations are; CollapsedUse forgets much of the
--   detailed usage information Use retained.
--
-- - Phase 4 actually extracts the inlining-relevant information from
--   CollapsedUse.

-- === Algebra ===

-- The algebraic structure we want in order to represent uses is
-- max-plus: max corresponds to `if` or `case` in the Dex program, and
-- plus corresponds to multiple sequential statements.
--
-- This is almost a semiring, in that
-- - max and plus are both monoids
-- - max (and plus) are commutative
-- - plus distributes over max
-- - however, in our case, all actual uses that will occur are
--   "positive", which means we can use `zero` (the identity for plus)
--   as the identity for max as well, except that it then doesn't obey
--   the annihilation law.

class MaxPlus a where
  zero :: a
  max :: a -> a -> a
  plus :: a -> a -> a

instance (MaxPlus (e n)) => MaxPlus (NameMapE c e n) where
  zero = mempty
  max  = unionWithNameMapE max
  plus = unionWithNameMapE plus

-- === Access ===

-- We represent an upper bound on the pattern by which (the object
-- bound to) a name is accessed.  The goal is to determine whether
-- inlining the name could duplicate work, so this is an upper bound
-- (as opposed to an estimate or a lower bound).

-- Each element (or subarray) of the object may be accessed a known or
-- unbounded number of times.  We also highlight "accessed once" and
-- "not accessed" as patterns, because that's what will mainly
-- determine inlining decisions.

-- We also use the Count type to count how many times a binding is
-- referenced statically.

data Count =
    Bounded Int
  | Unbounded
  deriving (Eq, Ord, Show, Generic)

instance Hashable Count
instance Store Count

pattern Zero :: Count
pattern Zero = Bounded 0

pattern One :: Count
pattern One = Bounded 1

pattern Two :: Count
pattern Two = Bounded 2

instance MaxPlus Count where
  zero = Zero
  max Unbounded _ = Unbounded
  max _ Unbounded = Unbounded
  max (Bounded i1) (Bounded i2) = Bounded $ i1 `Prelude.max` i2
  plus Unbounded _ = Unbounded
  plus _ Unbounded = Unbounded
  plus (Bounded i1) (Bounded i2) = Bounded $ i1 + i2

instance Bounded Count where
  maxBound = Unbounded
  minBound = zero

-- An access bound for a compound object determines some paths to some
-- leaves or subtrees of the object, and how many times that leaf or
-- subtree is accessed (at most).  Note that the access bound may
-- depend on names in the scope `n`, such as indices of enclosing
-- loops.

-- The Access represents accesses in "iteration space", i.e., the
-- binders in the `ForEach` constructors correspond to loops in the
-- user program.
data Access (n::S) where
  -- This chain of indices of this object is accessed this many times.
  -- Corresponds to indexing.
  Location :: [IxExpr n] -> Count -> Access n
  -- The object may be accessed in any one of these ways (as, e.g., in
  -- different arms of a downstream `case` expression).
  -- Corresponds to case expressions.
  Any :: [Access n] -> Access n
  -- The object may be accessed in some or all of these ways (e.g.,
  -- from multiple references).
  -- Corresponds to sequential statements (e.g., x.i; x.j)
  All :: [Access n] -> Access n
  -- The object may be accessed in any of the ways given by the second
  -- argument, as the name in the `Binder` argument varies over its
  -- type.
  -- Corresponds to `for`.
  ForEach :: forall n l. (NameBinder (AtomNameC SimpIR) n l) -> Access l -> Access n

deriving instance Show (Access n)

accessOnce :: Access n
accessOnce = Location [] One

accessMuch :: Access n -> Access n
accessMuch = \case
  Location ixs _ -> Location ixs Unbounded
  Any as -> Any $ map accessMuch as
  All as -> All $ map accessMuch as
  ForEach b a -> ForEach b $ accessMuch a

-- Prepend the given indexing expression onto an existing Access
location :: IxExpr n -> Access n -> Access n
location ix (Location ixs a) = Location (ix:ixs) a
location _ _ = error "Cannot nest general Access in Location"

-- These are the indexing expressions we model.  An IxExpr n is an abstraction
-- of a function from names in the n scope, to either a single index location or
-- to a set of index locations.  We only track the operations we have a hope of
-- being able to analyze usefully.
data IxExpr (n::S) =
  -- Indexing by an identifier in scope
    Var (Name (AtomNameC SimpIR) n)
  -- An index of product type
  | Product [IxExpr n]
  -- An index of sum type
  | Inject Int (IxExpr n)
  -- Unknown deterministic function of the given names
  | Deterministic [Name (AtomNameC SimpIR) n]
  -- All indices are accessed, or we are conservatively assuming they are
  -- because we don't know which are accessed.  See note "IxAll semantics" in
  -- "Notes".
  | IxAll
  -- TODO(precision) Could add more constructors to capture other indexing
  -- expressions, such as `from_ordinal _ $ ordinal i + 1`.  If that's
  -- the only indexing expression the binding is safe to inline, and
  -- inlining even preserves locality, not just work.
  deriving (Eq, Ord, Show)

-- We give Access semantics by a notional translation into a piece of
-- code that dynamically counts the number of times every index is
-- actually accessed.  Except, Access intentionally loses information,
-- so the semantics of Access is actually a program that has some
-- control path that accesses at least as much as the original.
--
-- - An Access becomes a function consuming a buffer reference
--   and returning unit.  The buffer has the same type as the
--   object whose accesses we are analyzing, except for integer
--   leaves.
--
-- - For every control path of the source program, there is a control
--   path of the Access program that accumulates at least as many
--   counts into each element of the buffer as the number of times the
--   source program read that element.
--
-- - `Location exprs ct` turns into accumulation
--     \ref. ref!<exprs> += ct
--     Note that this += must be understood as pointwise if
--     the source indexing expression is not fully saturated.
--
-- - `Any opts` turns into case
--     \ref. index = randint $ length opts
--       case index of
--         1 -> compile opt1 ref
--         2 -> compile opt2 ref
--         ...
--
-- - `All parts` turns into sequentialization
--     \ref.
--       compile part1 ref
--       compile part2 ref
--       ...
--
-- - `ForEach b sub` is iteration
--     \ref. for b in (type of b).
--       compile sub ref
--   The type that `b` iterates over can in principle be fetched from
--   the Dex program that led to this `ForEach` constructor.  It only
--   matters for the semantics, so we don't actually store it.

instance MaxPlus (Access n) where
  zero = Location [] Zero
  a1 `max` a2 = Any [a1, a2]
  a1 `plus` a2 = All [a1, a2]

-- === Use ===

-- The Use structure is a "data space" representation of accessing
-- arrays.
--
-- A Use implicitly iterates over as many dimensions of the accessed
-- object as there are nesting levels of At constructors in the Use.
-- We do not use explicit binders for this because they always appear
-- in the same nesting order, so it's simplest to leave them implicit.
-- For example, the case
--   xs = ...
--   for i. for j. xs.i.j
-- is represented by the Use
--   At IxAll (At IxAll One)
-- meaning "we access every element or subarray once, two dimensions
-- deep".
--
-- Note that a Use can still depend on iteration-space indices (from
-- the `iter` scope)---it just doesn't have internal binders that are
-- meant to be interpreted in iteration space.  For example, the Use
-- for `xs` in
--   for i.
--     xs = ...
--     xs.i
-- is
--   At (Var i) One
-- meaning that `xs` is accessed at the location `i`.

data Use (iter::S) =
  -- The object is accessed in its entirety
    Const Count
  -- The object is accessed at the positions given by the first
  -- argument.  The subarray at those positions is accessed according
  -- to the sub-Use.  Note that at this point we permit the `IxAll`
  -- constructor in `IxExpr`, so an `IxExpr` may represent a set of
  -- positions.
  | At (IxExpr iter) (Use iter)
  -- The object is accessed in at most one of the given Use patterns.
  | Max [Use iter]
  -- The object is accessed in all of the given Use patterns.
  | Sum [Use iter]
  deriving (Eq, Show)

-- The semantics of Use are similar to those for Access.  Here is the
-- compilation that turns a Use into a program that dynamically counts
-- how many times each index is used.
--
-- Note that IxExpr permits the `IxAll` constructor, so an `IxExpr`
-- represents a subset of indices rather than a single index.
--
-- - `Const ct` becomes accumulation
--     \ref. ref += ct
--     Also polymorphic in the ref type, applying elementwise
--
-- - `Max alts` becomes case
--     \ref. case randint (length alts) of
--       1 -> compile alt1 ref
--       2 -> compile alt2 ref
--       ...
--
-- - `Sum parts` becomes sequencing
--     \ref.
--       compile alt1 ref
--       compile alt2 ref
--       ...
--
-- - `At expr use` becomes testing and slicing
--     \ref. for b in <next object dim>.
--       if b in expr then (compile use ref!b)

-- This semantics implies a transformation that we perform eagerly:
-- - It's safe to distribute Sum into Max, so that Max ends up on the
--   outside.

-- All the At constructors created by `interp`, below, on any path
-- through a sound Access refer to the same dimensions of the object
-- being indexed, in the same order, with no gaps.  Why?  Because each
-- is introduced at the same point in the sequence of index
-- expressions applied to the object, regardless of the
-- iteration-space control path that got there.
--
-- Another way to see this is that the compilation of At is the only
-- thing that cares about and alters the type of the reference being
-- passed around.  Ergo, all the At constructors must occur in the
-- same order.

-- === Tropical semiring on Use ===

-- We lightly normalize Use data under these operations.  The main
-- invariants are
--
-- - Max constructors appear outside of adjacent Sum constructors
-- - If the `Use` represents no use, it's represented as `Const Zero`
--
-- We also take the opportunity to peephole-simplify the Use
-- structure.

instance MaxPlus (Use iter) where
  zero = Const zero

  max (Const Unbounded) _ = Const Unbounded
  max _ (Const Unbounded) = Const Unbounded
  max (Const b1) (Const b2) = Const $ b1 `max` b2
  max (Max uses1) (Max uses2) = Max $ uses1 ++ uses2
  -- This is trying to eliminate zero-valued elements, because they can cause
  -- (potentially severe) loss of precision downstream.  However, that relies on
  -- upstream construction canonicalizing all semantically-zero Uses to the
  -- `zero` use.
  -- TODO Prove or unit-test that upstream cannot construct semantically-zero
  -- Use objects that are not `== zero`.
  max (Max uses1) use2
    | use2 == zero = Max uses1
    | otherwise = Max $ uses1 ++ [use2]
  max use1 (Max uses2)
    | use1 == zero = Max uses2
    | otherwise = Max $ use1:uses2
  max use1 use2
    | use1 == zero = use2
    | use2 == zero = use1
    | otherwise = Max [use1, use2]

  plus (Const Unbounded) _ = Const Unbounded
  plus _ (Const Unbounded) = Const Unbounded
  plus (Const b1) (Const b2) = Const $ b1 `plus` b2
  plus (Max uses1) (Max uses2) = Max uses' where
    uses' = concat $ flip map uses1 (\use1 ->
      flip map uses2 \use2 -> plus use1 use2)
  plus u1@(Max _) use2
    | use2 == zero = u1
    | otherwise = plus u1 (Max [use2])
  plus use1 u2@(Max _)
    | use1 == zero = u2
    | otherwise = plus (Max [use1]) u2
  plus (Sum uses1) (Sum uses2) = Sum $ uses1 ++ uses2
  plus (Sum uses1) use2
    -- TODO(code cleanliness) Is there a way to deduplicate these zero checks?
    | use2 == zero = Sum uses1
    | otherwise = Sum $ uses1 ++ [use2]
  plus use1 (Sum uses2)
    | use1 == zero = Sum uses2
    | otherwise = Sum $ use1:uses2
  plus use1 use2
    | use1 == zero = use2
    | use2 == zero = use1
    | otherwise = Sum [use1, use2]

-- === Interpreting Access into Use ===

-- The main interpreter is pretty straightforward.  The main trick is
-- just semantics: an indexing expression becomes a (nested) `At`,
-- which represents a data-space loop that checks that the data-space
-- variable equals the indexing expression.
interp :: Access iter -> Use iter
interp = \case
  Location [] ct -> Const ct
  Location (expr:exprs) ct -> At expr $ interp $ Location exprs ct
  Any accesses -> foldl'  max zero $ map interp accesses
  All accesses -> foldl' plus zero $ map interp accesses
  ForEach b access -> iterate False b $ interp access

-- The interesting function is resolution of iteration-space loops.
-- We have to construct a Use that can refer to subsets of data space
-- (as opposed to single points, as in Access).
-- The `Bool` argument means "have we already observed that upstream
-- indexing touches different parts of the object?"  Even if the
-- answer is "yes", we do still have to recur, because we need to
-- translate the interation-space binder into data space.

iterate :: Bool -> (NameBinder (AtomNameC SimpIR) iter iter')
          -> Use iter' -> (Use iter)
iterate True _ (Const ct) = Const ct
iterate False _ (Const ct) =
  -- If we haven't proven upstream that indexing is injective,
  -- we have to assume here that it collides with itself.
  -- e.g., for i. xs
  -- Note: If we have a bound on the size of the index set traversed
  -- by `ib`, then we can compute a bound on the number of accesses,
  -- but there's no need to do that until we make inlining decisions
  -- on a more fine-grained basis than "was this used more than once?"
  if ct == zero then zero else Const Unbounded
iterate injective ib (At expr use) = At expr' use' where
  (expr', injective') =
    flip runState injective $ iterateIxExpr ib expr
  use' = iterate injective' ib use
iterate injective ib (Max uses) =
  -- TODO(correctness) This can be an underapproximation if the upstream
  -- structure becomes expressive enough, but it's OK for now.  For example,
  -- consider the access pattern
  --   for i.
  --     if ...
  --       then xs.i
  --       else xs.(i+1)
  -- Here inlining xs will in general duplicate work, because the `if`
  -- could go different ways in different iterations, thus accessing
  -- some elements of `xs` more than once (and others not at all).
  --
  -- However, if we were to distribute the `for` into the `if` as
  -- this code does, we would get
  --   if ...
  --     then for i. xs.i
  --     else for i. xs.(i+1)
  -- Now, `xs` _is_ work-preserving to inline here, because each
  -- branch accesses each element of `xs` only once.  So mistaking
  -- the former for the latter would be a mistake.
  --
  -- However, the current Access structure cannot actually _prove_
  -- that inlining `xs` into the latter expression is safe, because
  -- the addition will register as an unknown function of `i`.  The
  -- distribution bug is therefore masked.
  --
  -- In fact, that currently always happens.  All distinct functions
  -- that produce indices of type n that Access can currently
  -- represent are either `Deterministic` or have disjoint ranges, so this
  -- kind of underapproximation cannot yet lead to a mistaken
  -- assertion that some binding is safe to inline.
  foldl' max zero $ map (iterate injective ib) uses
iterate injective ib (Sum uses) =
  foldl' plus zero $ map (iterate injective ib) uses

iterateIxExpr :: NameBinder (AtomNameC SimpIR) iter iter'
              -> IxExpr iter' -> State Bool (IxExpr iter)
iterateIxExpr ib (Var i) =
  case hoist ib i of
    HoistSuccess i' -> return $ Var i'
    HoistFailure _ ->
      -- If we hit a reference to the iteration binder,
      -- that constitutes a proof that indexing is injective,
      -- because all constructors of `IxExpr` that this could
      -- be nested under represent invertible functions.
      -- If we were already injective, this case is an
      -- over-approximation: the object is actually accessed at just
      -- one value of the index, but since we don't know which one, we
      -- assume it was all of them.
      put True >> return IxAll
iterateIxExpr ib (Product elts) =
  -- Rectangular approximation to something like xs.(j, j)
  Product <$> mapM (iterateIxExpr ib) elts
iterateIxExpr ib (Inject i elt) =
  Inject i <$> iterateIxExpr ib elt
iterateIxExpr ib (Deterministic names) =
  case hoist ib (ListE names) of
    HoistSuccess (ListE names') -> return $ Deterministic names'
    -- If we hit a reference to an unknown function of the binder, we
    -- must assume it can hit any data index (which we
    -- over-approximate as it hitting all of them); and we also do not
    -- gain a proof that the indexing is injective, because different
    -- values of the iteration-space binder could lead to the same
    -- value of the data-space binder.
    HoistFailure _ -> return IxAll
iterateIxExpr _ IxAll = return IxAll

-- === CollapsedUse ===

-- CollapsedUse is like Use, but now it is also independent of binders
-- in scope outside the inlining decision under consideration.
-- Collapsing Use into CollapsedUse perforce over-approximates those
-- binders, potentially severely.

-- A CollapsedUse is either a count of the accesses to the whole
-- object, or a subset of the leading dimension of the object.  In the
-- latter case, the object is not accessed outside that subset.  The
-- Subset structure maintains a map of further subsets, and how later
-- dimensions are accessed in each of them.

-- Note that we care not just about which elements are used, but also
-- the indexing depth of the object.  That is `CConst One` and
-- `CSubset (SAll (CConst One))` both mean "used everywhere exactly
-- once", but the latter also means "do not inline unless you have an
-- explicit `for` to cancel against indexing".
data CollapsedUse =
    CConst Count
  | CSubset (Subset CollapsedUse)
  deriving (Eq, Show)

-- A subset of an index type.  The constructors correspond to possible
-- index types---each one gives known information about the structure
-- of the type, and relevant available subset information.
-- A trie on subtypes of an index type that supports subset
-- intersection as efficiently as is practicable.
data Subset a =
  -- The index could be any type, and all of it is accessed.
    SAll a
  -- The index is a sum.  We maintain a map from constructors to
  -- subsets of their respective types.  The Maybe is a performance
  -- optimization: a `(Just a)` is associated with the whole type,
  -- which we just store here instead of broadcasting it over all the
  -- cases.
  | Inj  (Maybe a) (M.IntMap (Subset a))
  -- The index is a product.  We maintain one rectangle as a subset.
  | Prod [Subset ()] a
  deriving (Eq, Show)

-- === Collapsing Use to CollapsedUse ===

collapse :: Use iter -> CollapsedUse
collapse (Const ct) = (CConst ct)
collapse (At expr use) =
  let use' = collapse use in
  if use' == zero then zero
  else CSubset $ exprToSubset expr use'
collapse (Sum uses) = foldl' plus zero $ map collapse uses
collapse (Max uses) = foldl' max zero $ map collapse uses

exprToSubset :: IxExpr n -> a -> Subset a
-- An iteration-space variable approximates to "all", because at this
-- point it can take any value.  This is a severe over-approximation,
-- but simplifies the structure.
--
-- - In particular, distinct iteration-space variables will collide
--   rarely, but we conservatively assume they always do.  Consider
--   inlining x into
--     for i. x.j.i
--     for i. x.k.i
--   where j and k are free outside the binding of x.  Whenever j and
--   k are distinct, inlining leads to savings; potentially large
--   savings if the j and k spaces are large.  But, they may
--   occasionally be equal, in which case inlining would duplicate
--   work (but on just those elements).  We conservatively miss this
--   inlining opportunity, but perhaps should revisit that.
exprToSubset (Var _) x = SAll x
exprToSubset (Product exprs) x = Prod (map (flip exprToSubset ()) exprs) x
exprToSubset (Inject i expr) x =
  Inj Nothing $ M.singleton i $ exprToSubset expr x
exprToSubset IxAll   x = SAll x
-- A Deterministic that survies to this point should only reference
-- iteration-space variables that are scoped over the entire inlining,
-- like the `Var` constructor.  We overapproximate with `SAll` for the
-- same reason.
exprToSubset (Deterministic _) x = SAll x

instance MaxPlus CollapsedUse where
  zero = CConst zero

  plus :: CollapsedUse -> CollapsedUse -> CollapsedUse
  plus (CConst ct1) (CConst ct2) = CConst $ ct1 `plus` ct2
  plus (CSubset s1) (CSubset s2) = CSubset $ s1 `plus` s2
  plus (CSubset s1) (CConst ct2)
    | ct2 == zero = CSubset s1
    | otherwise = CSubset $ s1 `plus` (SAll (CConst ct2))
  plus (CConst ct1) (CSubset s2)
    | ct1 == zero = CSubset s2
    | otherwise = CSubset $ (SAll (CConst ct1)) `plus` s2

  max :: CollapsedUse -> CollapsedUse -> CollapsedUse
  max (CConst ct1) (CConst ct2) = CConst $ ct1 `max` ct2
  max (CSubset s1) (CSubset s2) = CSubset $ s1 `max` s2
  max (CSubset s1) (CConst ct2)
    | ct2 == zero = CSubset s1
    | otherwise = CSubset $ s1 `max` (SAll (CConst ct2))
  max (CConst ct1) (CSubset s2)
    | ct1 == zero = CSubset s2
    | otherwise = CSubset $ (SAll (CConst ct1)) `max` s2

-- === Subset operations ===

instance MaxPlus () where
  zero = ()
  max _ _ = ()
  plus _ _ = ()

instance MaxPlus a => MaxPlus (Subset a) where
  zero = SAll zero
  max = smerge max
  plus = smerge plus

smerge :: MaxPlus a => (a -> a -> a) -> Subset a -> Subset a -> Subset a
smerge merge (SAll a1) (SAll a2) = SAll $ merge a1 a2
smerge merge (SAll a) i@(Inj _ _) = smerge merge (Inj (Just a) M.empty) i
smerge merge i@(Inj _ _) (SAll a) = smerge merge i (Inj (Just a) M.empty)
smerge merge (Inj a1 elts1) (Inj a2 elts2) = Inj a' elts' where
  a' = case (a1, a2) of
    (Nothing, Just a)  -> Just a
    (Just a, Nothing)  -> Just a
    (Just a1', Just a2') -> Just $ merge a1' a2'
    (Nothing, Nothing) -> Nothing
  elts' = M.unionWith (smerge merge) elts1 elts2
smerge merge (Prod elts1 a1) (Prod elts2 a2) = Prod elts' a' where
  -- We approximate the max or the sum of two rectangles as the
  -- smallest rectangle that contains them.
  -- TODO(precision) This needlessly loses precision if one of the aruments is
  -- actually a zero use; but we hope they are all normalized away upstream.
  elts' = zipWith max elts1 elts2
  -- The payload of that rectangle is the requested merge if the
  -- arguments intersect, but always the max if they do not.
  a' = if all id $ zipWith intersect elts1 elts2 then
         merge a1 a2
       else
         -- always max, regardless of the merge function
         a1 `max` a2
smerge merge (SAll a1) (Prod _ a2) = SAll $ merge a1 a2
smerge merge (Prod _ a1) (SAll a2) = SAll $ merge a1 a2
smerge _ (Inj _ _) (Prod _ _) = error "Impossible: index type both a sum and a product"
smerge _ (Prod _ _) (Inj _ _) = error "Impossible: index type both a sum and a product"

intersect :: Subset a -> Subset a -> Bool
intersect (SAll _) _ = True
intersect _ (SAll _) = True
intersect (Inj (Just _) _) _ = True
intersect _ (Inj (Just _) _) = True
intersect (Inj Nothing elts1) (Inj Nothing elts2) =
  not $ M.null $ M.intersection elts1 elts2
intersect (Prod elts1 _) (Prod elts2 _) =
  all id $ zipWith intersect elts1 elts2
intersect (Inj _ _) (Prod _ _) = error "Impossible: index type both a sum and a product"
intersect (Prod _ _) (Inj _ _) = error "Impossible: index type both a sum and a product"

-- === Overapproximating a CollapsedUse with a final Count ===

type DynUseInfo = (Int, Count)  -- The int is maximum number of indexed dimensions

instance MaxPlus Int where
  zero = 0
  max = Prelude.max
  plus = (+)

instance (MaxPlus a, MaxPlus b) => MaxPlus (a, b) where
  zero = (zero, zero)
  (a1, b1) `max` (a2, b2) = (a1 `max` a2, b1 `max` b2)
  (a1, b1) `plus` (a2, b2) = (a1 `plus` a2, b1 `plus` b2)

class ApproxConst a where
  approxConst :: a -> DynUseInfo

instance ApproxConst CollapsedUse where
  approxConst (CConst ct) = (0, ct)
  approxConst (CSubset subset) = (1, zero) `plus` approxConst subset

instance ApproxConst a => ApproxConst (Subset a) where
  approxConst (SAll use) = approxConst use
  approxConst (Inj (Just all_use) uses) =
    -- TODO(optimization) Short-circuit if the all_use is heavy enough
    approxConst all_use `plus` approxConst uses
  approxConst (Inj Nothing uses) = approxConst uses
  approxConst (Prod _ use) = approxConst use

instance ApproxConst v => ApproxConst (M.IntMap v) where
  approxConst uses = foldl' max zero $ map approxConst $ M.elems uses

-- === Usage info including static usage ===

-- The `Count` is the count of static occurrences, i.e., how many times
-- references to this binding occur in the source code.  We track this
-- separately from the (dynamic) Access so we can control code blow-up during
-- inlining.
data AccessInfo n = AccessInfo Count (Access n)
  deriving (Show)

instance MaxPlus (AccessInfo n) where
  zero = AccessInfo zero zero
  -- Note that, since `max` corresponds to `case`, the static count is
  -- still added
  (AccessInfo s1 d1) `max` (AccessInfo s2 d2) =
    AccessInfo (s1 `plus` s2) $ d1 `max` d2
  (AccessInfo s1 d1) `plus` (AccessInfo s2 d2) =
    AccessInfo (s1 `plus` s2) $ d1 `plus` d2

instance SinkableE AccessInfo where
  sinkingProofE rename (AccessInfo n a) = AccessInfo n $ sinkingProofE rename a

data UsageInfo = UsageInfo Count DynUseInfo  -- The `Count` is the static use count
  deriving (Eq, Generic, Show)

instance MaxPlus UsageInfo where
  zero = UsageInfo zero zero
  -- Note that, since `max` corresponds to `case`, the static count is
  -- still added
  (UsageInfo s1 d1) `max` (UsageInfo s2 d2) =
    UsageInfo (s1 `plus` s2) $ d1 `max` d2
  (UsageInfo s1 d1) `plus` (UsageInfo s2 d2) =
    UsageInfo (s1 `plus` s2) $ d1 `plus` d2

usageInfo :: AccessInfo n -> UsageInfo
usageInfo (AccessInfo s dyn) =
  UsageInfo s $ approxConst $ collapse $ interp dyn
{-# SCC usageInfo #-}

-- === Notes ===

-- Projection (inlining tuple-typed bindings)
--
-- Adding Projection is pretty easy: projections slice the reference,
-- so have a defined order interleaved with At.
-- - If a body is parametrically polymorphic in the reference, it's safe
--   to wrap it with a sum of all the projections
-- - So does this mean that Projections will be places where formal Sums
--   and Maxes have to rest?  That seems annoying.
-- - I think it's the case that a Sum-Max tree over Projections, where
--   everything underneath is parametrically polymorphic in the
--   reference type (i.e., when we are computing the final answer),
--   can just become the relevant sums and maxes of the pieces.  (Sum
--   over projecting to distinct indices turns into max; sum over
--   colliding indices turns into sum; max turns into max whether the
--   indices collide or not.)
--   - max over Sum [Projection-with-distinct-indices] can be approximated as
--     Sum [indexwise max], but I only want to do that once I've accumulated
--     all the maxes I am going to?  Or does that not matter?
--   - So maybe Sum [Projection-with-distinct-indices] is a normal form
--     without needing formal Max nodes?
--   - Come to think of it, is the same thing true of At?

-- Sum and Max
--
-- In general:
-- - max over Sum unions the indices and does an indexwise max and emits Sum
-- - max over Max unions the indices and does an indexwise max and emits Max
-- - sum over Sum unions the indices and does an indexwise sum and emits Sum
-- - sum over Max unions the indices and does an indexwise sum and emits Max
-- - So if I'm never going to turn formal Sum or Max nodes back into sum or max
--   operations, there is no difference between them.
-- - At the end, when approximating with Const, I max over distinct indices
--   regardless of whether it's a formal Max or a formal Sum node.
-- - However, the distinction between Sum and Max matters if I'm trying to
--   compute whether a given object is accessed exactly once everywhere,
--   so as to try to inline an effectful operation.  (That said, inlining
--   effects also requires paying attention to the order of the accesses.)

-- IxAll semantics
--
-- IxAll can actually also mean "accessed at one unknown index".  Logically this
-- is fine because the Access is an overapproximation.  It can arise when, e.g.,
-- indexing by the content of a reference, since the analysis doesn't trace what
-- values might occupy that reference.
--
-- Some examples may help clarify.
--
-- Example 1
--   xs = for ...
--   with_state ... \ref.
--     xs.(get ref)
--
-- - The indexing is represented by `Location IxAll One`
-- - It doesn't change when it goes above the `ref` binder
-- - `xs` will look inlinable, which is fine because it's only accessed once.
--   In fact, it's probably very beneficial to inline `xs` because it's only
--   accessed at one location, but the analysis doesn't notice that.
--
-- Example 2
--   xs = for ...
--   with_state ... \ref.
--     for i.
--       xs.(get ref)
--
-- - The indexing is represented by `Location IxAll One`
-- - Going above the `i` binder turns this into `Location IxAll Unbounded`
--   because there is no proof that the accesses are distinct across loop
--   iterations (and indeed they are not if the reference is not modified)
-- - `xs` will not look inlinable, which is sound and not even too conservative
--   because the same element may be accessed many times.
--
-- Example 3
--   xs = for n. for ...
--   with_state ... \ref.
--     for i.
--       xs.i.(get ref)
--
-- - The outer indexing is represented by `Location IxAll One`
-- - The inner indexing is represented by `Location (Var i) (Location IxAll One)`
-- - Going above the `i` binder turns this into `Location IxAll (Location IxAll
--   One)`.  This is sound because there _is_ a proof that the accesses are
--   distinct across loop iterations -- each `xs.i` subarray is accessed in at
--   most one unknown location.
-- - `xs` will look inlinable.
--
-- TODO(precision) Could add a constructor to IxExpr to model "accessed at just
-- one unknown index", which might support trying harder to inline an array if
-- very few of its elements are actually read.  The case to imagine is a big
-- array whose elements have fairly uniform costs to compute, and which is
-- accessed at just a few (potentially colliding) indices.  Then inlining
-- duplicates work on some elements, but saves work overall.  The current
-- analysis would, however, conservatively avoid inlining such an array because
-- of the per-element work duplication.

-- === Boring instances ===

instance GenericE IxExpr where
  type RepE IxExpr = EitherE5
    (Name (AtomNameC SimpIR))
    (ListE IxExpr)
    (PairE (LiftE Int) IxExpr)
    (ListE (Name (AtomNameC SimpIR)))
    UnitE
  fromE = \case
    (Var n)            -> Case0   n
    (Product ixs)      -> Case1 $ ListE ixs
    (Inject i ix)      -> Case2 $ PairE (LiftE i) ix
    (Deterministic ns) -> Case3 $ ListE ns
    IxAll              -> Case4   UnitE
  {-# INLINE fromE #-}
  toE = \case
    (Case0 n) -> Var n
    (Case1 (ListE ixs)) -> Product ixs
    (Case2 (PairE (LiftE i) ix)) -> Inject i ix
    (Case3 (ListE ns)) -> Deterministic ns
    (Case4 UnitE) -> IxAll
    _ -> error "impossible"
  {-# INLINE toE #-}

instance HoistableE IxExpr
instance SinkableE IxExpr
instance RenameE IxExpr

instance GenericE Access where
  type RepE Access = EitherE4
    (PairE (ListE IxExpr) (LiftE Count))
    (ListE Access)
    (ListE Access)
    (Abs (NameBinder (AtomNameC SimpIR)) Access)
  fromE = \case
    (Location ixs ct) -> Case0 $ PairE (ListE ixs) (LiftE ct)
    (Any as) -> Case1 $ ListE as
    (All as) -> Case2 $ ListE as
    (ForEach b a) -> Case3 $ Abs b a
  {-# INLINE fromE #-}
  toE = \case
    (Case0 (PairE (ListE ixs) (LiftE ct))) -> Location ixs ct
    (Case1 (ListE as)) -> Any as
    (Case2 (ListE as)) -> All as
    (Case3 (Abs b a))  -> ForEach b a
    _ -> error "impossible"
  {-# INLINE toE #-}

instance HoistableE Access
instance SinkableE Access
instance RenameE Access

instance GenericE AccessInfo where
  type RepE AccessInfo = PairE (LiftE Count) Access
  fromE (AccessInfo s dyn) = PairE (LiftE s) dyn
  {-# INLINE fromE #-}
  toE (PairE (LiftE s) dyn) = AccessInfo s dyn
  {-# INLINE toE #-}

instance HoistableE AccessInfo
instance RenameE AccessInfo

instance Hashable UsageInfo
instance Store UsageInfo

-- === instances ===

instance Pretty UsageInfo where
  pretty (UsageInfo static (ixDepth, ct)) =
    "occurs in" <+> pretty static <+> "places, read"
    <+> pretty ct <+> "times, to depth" <+> pretty (show ixDepth)

instance Pretty Count where
  pretty = \case
    Bounded ct -> "<=" <+> pretty ct
    Unbounded -> "many"