Note [Occurrence analysis for join points]

GHC/Core/Opt/OccurAnal.hs:602 compiler 2 tickets

Consider these two somewhat artificial programs (#22404)

  Program (P1)                      Program (P2)
     -------------------------------------
  let v = <small thunk> in          let v = <small thunk> in
                                    join j = case v of (a,b) -> a
  in case x of                      in case x of
        A -> case v of (a,b) -> a         A -> j
        B -> case v of (a,b) -> a         B -> j
        C -> case v of (a,b) -> b         C -> case v of (a,b) -> b
        D -> []                           D -> []

In (P1), `v` gets allocated, as a thunk, every time this code is executed.  But
notice that `v` occurs at most once in any case branch; the occurrence analyser
spots this and returns a OneOcc{ occ_n_br = 3 } for `v`.  Then the code in
GHC.Core.Opt.Simplify.Utils.postInlineUnconditionally inlines `v` at its three
use sites, and discards the let-binding.  That way, we avoid allocating `v` in
the A,B,C branches (though we still compute it of course), and branch D
doesn't involve <small thunk> at all.  This sometimes makes a Really Big
Difference.

In (P2) we have shared the common RHS of A, B, in a join point `j`.  We would
like to inline `v` in just the same way as in (P1).  But the usual strategy
for let bindings is conservative and uses `andUDs` to combine usage from j's
RHS to its body; as if `j` was called on every code path (once, albeit).  In
the case of (P2), we'll get ManyOccs for `v`.  Important optimisation lost!

Solving this problem makes the Simplifier less fragile.  For example,
the Simplifier might inline `j`, and convert (P2) into (P1)... or it might
not, depending in a perhaps-fragile way on the size of the join point.
I was motivated to implement this feature of the occurrence analyser
when trying to make optimisation join points simpler and more robust
(see e.g. #23627).

The occurrence analyser therefore has clever code that behaves just as
if you inlined `j` at all its call sites.  Here is a tricky variant
to keep in mind:

  Program (P3)
  
    join j = case v of (a,b) -> a
    in case f v of
          A -> j
          B -> j
          C -> []

If you mentally inline `j` you'll see that `v` is used twice on the path
through A, so it should have ManyOcc.  Bear this case in mind!

* We treat /non-recursive/ join points specially. Recursive join points are
  treated like any other letrec, as before.  Moreover, we only give this special
  treatment to /pre-existing/ non-recursive join points, not the ones that we
  discover for the first time in this sweep of the occurrence analyser.

* In occ_env, the new (occ_join_points :: IdEnv OccInfoEnv) maps
  each in-scope non-recursive join point, such as `j` above, to
  a "zeroed form" of its RHS's usage details. The "zeroed form"
    * deletes ManyOccs
    * maps a OneOcc to OneOcc{ occ_n_br = 0 }
  In our example, occ_join_points will be extended with
      [j :-> [v :-> OneOcc{occ_n_br=0}]]
  See addJoinPoint.

* At an occurrence of a join point, we do everything as normal, but add in the
  UsageDetails from the occ_join_points.  See mkOneOcc.

* Crucially, at the NonRec binding of the join point, in `occAnalBind`, we use
  `orUDs`, not `andUDs` to combine the usage from the RHS with the usage from
  the body.

Here are the consequences

* Because of the perhaps-surprising OneOcc{occ_n_br=0} idea of the zeroed
  form, the occ_n_br field of a OneOcc binder still counts the number of
  /actual lexical occurrences/ of the variable.  In Program P2, for example,
  `v` will end up with OneOcc{occ_n_br=2}, not occ_n_br=3.
  There are two lexical occurrences of `v`!
  (NB: `orUDs` adds occ_n_br together, so occ_n_br=1 is impossible, too.)

* In the tricky (P3) we'll get an `andUDs` of
    * OneOcc{occ_n_br=0} from the occurrences of `j`)
    * OneOcc{occ_n_br=1} from the (f v)
  These are `andUDs` together in `addOccInfo`, and hence
  `v` gets ManyOccs, just as it should.  Clever!

There are a couple of tricky wrinkles

(W1) Consider this example which shadows `j`:
          join j = rhs in
          in case x of { K j -> ..j..; ... }
     Clearly when we come to the pattern `K j` we must drop the `j`
     entry in occ_join_points.

     This is done by `drop_shadowed_joins` in `addInScope`.

(W2) Consider this example which shadows `v`:
          join j = ...v...
          in case x of { K v -> ..j..; ... }

     We can't make j's occurrences in the K alternative give rise to an
     occurrence of `v` (via occ_join_points), because it'll just be deleted by
     the `K v` pattern.  Yikes.  This is rare because shadowing is rare, but
     it definitely can happen.  Solution: when bringing `v` into scope at
     the `K v` pattern, chuck out of occ_join_points any elements whose
     UsageDetails mentions `v`.  Instead, just `andUDs` all that usage in
     right here.

     This requires work in two places.
     * In `preprocess_env`, we detect if the newly-bound variables intersect
       the free vars of occ_join_points.  (These free vars are conveniently
       simply the domain of the OccInfoEnv for that join point.) If so,
       we zap the entire occ_join_points.
     * In `postprcess_uds`, we add the chucked-out join points to the
       returned UsageDetails, with `andUDs`.

(W3) Consider this example, which shadows `j`, but this time in an argument
              join j = rhs
              in f (case x of { K j -> ...; ... })
     We can zap the entire occ_join_points when looking at the argument,
     because `j` can't posibly occur -- it's a join point!  And the smaller
     occ_join_points is, the better.  Smaller to look up in mkOneOcc, and
     more important, less looking-up when checking (W2).

     This is done in setNonTailCtxt.  It's important /not/ to do this for
     join-point RHS's because of course `j` can occur there!

     NB: this is just about efficiency: it is always safe /not/ to zap the
     occ_join_points.

(W4) What if the join point binding has a stable unfolding, or RULES?
     They are just alternative right-hand sides, and at each call site we
     will use only one of them. So again, we can use `orUDs` to combine
     usage info from all these alternatives RHSs.

Wrinkles (W1) and (W2) are very similar to Note [Binder swap] (BS3).

References 1

Referenced by 8