module Language.PureScript.CoreFn.Laziness
  ( applyLazinessTransform
  ) where

import Protolude hiding (force)
import Protolude.Unsafe (unsafeHead)

import Control.Arrow ((&&&))
import qualified Data.Array as A
import Data.Coerce (coerce)
import Data.Graph (SCC(..), stronglyConnComp)
import Data.List (foldl1', (!!))
import qualified Data.IntMap.Monoidal as IM
import qualified Data.IntSet as IS
import qualified Data.Map.Monoidal as M
import Data.Semigroup (Max(..))
import qualified Data.Set as S

import Language.PureScript.AST.SourcePos
import qualified Language.PureScript.Constants.Libs as C
import Language.PureScript.CoreFn
import Language.PureScript.Crash
import Language.PureScript.Names
import Language.PureScript.PSString (mkString)

-- This module is responsible for ensuring that the bindings in recursive
-- binding groups are initialized in a valid order, introducing run-time
-- laziness and initialization checks as necessary.
--
-- PureScript is a call-by-value language with strict data constructors, this
-- transformation notwithstanding. The only laziness introduced here is in the
-- initialization of a binding. PureScript is uninterested in the order in
-- which bindings are written by the user. The compiler has always attempted to
-- emit the bindings in an order that makes sense for the backend, but without
-- this transformation, recursive bindings are emitted in an arbitrary order,
-- which can cause unexpected behavior at run time if a binding is dereferenced
-- before it has initialized.
--
-- To prevent unexpected errors, this transformation does a syntax-driven
-- analysis of a single recursive binding group to attempt to statically order
-- the bindings, and when that fails, falls back to lazy initializers that will
-- succeed or fail deterministically with a clear error at run time.
--
-- Example:
--
--   x = f \_ ->
--     x
--
-- becomes (with some details of the $runtime_lazy function elided):
-- 
--   -- the binding of x has been rewritten as a lazy initializer
--   $lazy_x = $runtime_lazy \_ ->
--     f \_ ->
--       $lazy_x 2  -- the reference to x has been rewritten as a force call
--   x = $lazy_x 1
--
-- Central to this analysis are the concepts of delay and force, which are
-- attributes given to every subexpression in the binding group. Delay and
-- force are defined by the following traversal. This traversal is used twice:
-- once to collect all the references made by each binding in the group, and
-- then again to rewrite some references to force calls. (The implications of
-- delay and force on initialization order are specified later.)

-- |
-- Visits every `Var` in an expression with the provided function, including
-- the amount of delay and force applied to that `Var`, and substitutes the
-- result back into the tree (propagating an `Applicative` effect).
--
-- Delay is a non-negative integer that represents the number of lambdas that
-- enclose an expression. Force is a non-negative integer that represents the
-- number of values that are being applied to an expression. Delay is always
-- statically determinable, but force can be *unknown*, so it's represented
-- here with a Maybe. In a function application `f a b`, `f` has force 2, but
-- `a` and `b` have unknown force--it depends on what `f` does with them.
--
-- The rules of assigning delay and force are simple:
--   * The expressions that are assigned to bindings in this group have
--     delay 0, force 0.
--   * In a function application, the function expression has force 1 higher
--     than the force of the application expression, and the argument
--     expression has unknown force.
--     * UNLESS this argument is being directly provided to a constructor (in
--       other words, the function expression is either a constructor itself or
--       a constructor that has already been partially applied), in which case
--       the force of both subexpressions is unchanged. We can assume that
--       constructors don't apply any additional force to their arguments.
--   * If the force of a lambda is zero, the delay of the body of the lambda is
--     incremented; otherwise, the force of the body of the lambda is
--     decremented. (Applying one argument to a lambda cancels out one unit of
--     delay.)
--   * In the argument of a Case and the bindings of a Let, force is unknown.
--   * Everywhere else, preserve the delay and force of the enclosing
--     expression.
--
-- Here are some illustrative examples of the above rules. We will use a
-- pseudocode syntax to annotate a subexpression with delay and force:
-- `expr#d!f` means `expr` has delay d and force f. `!*` is used to denote
-- unknown force.
--
--   x = y#0!0
--   x = y#0!2 a#0!* b#0!*
--   x = (\_ -> y#1!0)#0!0
--   x = \_ _ -> y#2!1 a#2!*
--   x = (\_ -> y#0!0)#0!1 z#0!*
--   x = Just { a: a#0!0, b: b#0!0 }
--   x = let foo = (y#1!* a b#1!*)#1!* in foo + 1
--
-- (Note that this analysis is quite ignorant of any actual control flow
-- choices made at run time. It doesn't even track what happens to a reference
-- after it has been locally bound by a Let or Case. Instead, it just assumes
-- the worst--once locally bound to a new name, it imagines that absolutely
-- anything could happen to that new name and thus to the underlying reference.
-- But the value-to-weight ratio of this approach is perhaps surprisingly
-- high.)
--
-- Every subexpression gets a delay and a force, but we are only interested
-- in references to other bindings in the binding group, so the traversal only
-- exposes `Var`s to the provided function.
--
onVarsWithDelayAndForce :: forall f. Applicative f => (Int -> Maybe Int -> Ann -> Qualified Ident -> f (Expr Ann)) -> Expr Ann -> f (Expr Ann)
onVarsWithDelayAndForce :: forall (f :: * -> *).
Applicative f =>
(Int -> Maybe Int -> Ann -> Qualified Ident -> f (Expr Ann))
-> Expr Ann -> f (Expr Ann)
onVarsWithDelayAndForce Int -> Maybe Int -> Ann -> Qualified Ident -> f (Expr Ann)
f = forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int
-> Maybe Int
-> (Bind Ann -> f (Bind Ann), Expr Ann -> f (Expr Ann))
go Int
0 forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just Int
0
  where
  go :: Int -> Maybe Int -> (Bind Ann -> f (Bind Ann), Expr Ann -> f (Expr Ann))
  go :: Int
-> Maybe Int
-> (Bind Ann -> f (Bind Ann), Expr Ann -> f (Expr Ann))
go Int
delay Maybe Int
force = (Bind Ann -> f (Bind Ann)
handleBind, Expr Ann -> f (Expr Ann)
handleExpr')
    where
    (Bind Ann -> f (Bind Ann)
handleBind, Expr Ann -> f (Expr Ann)
handleExpr, Binder Ann -> f (Binder Ann)
handleBinder, CaseAlternative Ann -> f (CaseAlternative Ann)
handleCaseAlternative) = forall (f :: * -> *) a.
Applicative f =>
(Bind a -> f (Bind a))
-> (Expr a -> f (Expr a))
-> (Binder a -> f (Binder a))
-> (CaseAlternative a -> f (CaseAlternative a))
-> (Bind a -> f (Bind a), Expr a -> f (Expr a),
    Binder a -> f (Binder a),
    CaseAlternative a -> f (CaseAlternative a))
traverseCoreFn Bind Ann -> f (Bind Ann)
handleBind Expr Ann -> f (Expr Ann)
handleExpr' Binder Ann -> f (Binder Ann)
handleBinder CaseAlternative Ann -> f (CaseAlternative Ann)
handleCaseAlternative
    handleExpr' :: Expr Ann -> f (Expr Ann)
handleExpr' = \case
      Var Ann
a Qualified Ident
i -> Int -> Maybe Int -> Ann -> Qualified Ident -> f (Expr Ann)
f Int
delay Maybe Int
force Ann
a Qualified Ident
i
      Abs Ann
a Ident
i Expr Ann
e -> forall a. a -> Ident -> Expr a -> Expr a
Abs Ann
a Ident
i forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a b. (a, b) -> b
snd (if Maybe Int
force forall a. Eq a => a -> a -> Bool
== forall a. a -> Maybe a
Just Int
0 then Int
-> Maybe Int
-> (Bind Ann -> f (Bind Ann), Expr Ann -> f (Expr Ann))
go (forall a. Enum a => a -> a
succ Int
delay) Maybe Int
force else Int
-> Maybe Int
-> (Bind Ann -> f (Bind Ann), Expr Ann -> f (Expr Ann))
go Int
delay forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. Enum a => a -> a
pred Maybe Int
force) Expr Ann
e
      -- A clumsy hack to preserve TCO in a particular idiom of unsafePartial once seen in Data.Map.Internal, possibly still used elsewhere.
      App Ann
a1 e1 :: Expr Ann
e1@(Var Ann
_ Qualified Ident
C.I_unsafePartial) (Abs Ann
a2 Ident
i Expr Ann
e2) -> forall a. a -> Expr a -> Expr a -> Expr a
App Ann
a1 Expr Ann
e1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> Ident -> Expr a -> Expr a
Abs Ann
a2 Ident
i forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr Ann -> f (Expr Ann)
handleExpr' Expr Ann
e2
      App Ann
a Expr Ann
e1 Expr Ann
e2 ->
        -- `handleApp` is just to handle the constructor application exception
        -- somewhat gracefully (i.e., without requiring a deep inspection of
        -- the function expression at every step). If we didn't care about
        -- constructors, this could have been simply:
        --   App a <$> snd (go delay (fmap succ force)) e1 <*> snd (go delay Nothing) e2
        Int -> [(Ann, Expr Ann)] -> Expr Ann -> f (Expr Ann)
handleApp Int
1 [(Ann
a, Expr Ann
e2)] Expr Ann
e1
      Case Ann
a [Expr Ann]
vs [CaseAlternative Ann]
alts -> forall a. a -> [Expr a] -> [CaseAlternative a] -> Expr a
Case Ann
a forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ Int
-> Maybe Int
-> (Bind Ann -> f (Bind Ann), Expr Ann -> f (Expr Ann))
go Int
delay forall a. Maybe a
Nothing) [Expr Ann]
vs forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse CaseAlternative Ann -> f (CaseAlternative Ann)
handleCaseAlternative [CaseAlternative Ann]
alts
      Let Ann
a [Bind Ann]
ds Expr Ann
e -> forall a. a -> [Bind a] -> Expr a -> Expr a
Let Ann
a forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ Int
-> Maybe Int
-> (Bind Ann -> f (Bind Ann), Expr Ann -> f (Expr Ann))
go Int
delay forall a. Maybe a
Nothing) [Bind Ann]
ds forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr Ann -> f (Expr Ann)
handleExpr' Expr Ann
e
      Expr Ann
other -> Expr Ann -> f (Expr Ann)
handleExpr Expr Ann
other

    handleApp :: Int -> [(Ann, Expr Ann)] -> Expr Ann -> f (Expr Ann)
handleApp Int
len [(Ann, Expr Ann)]
args = \case
      App Ann
a Expr Ann
e1 Expr Ann
e2 -> Int -> [(Ann, Expr Ann)] -> Expr Ann -> f (Expr Ann)
handleApp (Int
len forall a. Num a => a -> a -> a
+ Int
1) ((Ann
a, Expr Ann
e2) forall a. a -> [a] -> [a]
: [(Ann, Expr Ann)]
args) Expr Ann
e1
      Var a :: Ann
a@(SourceSpan
_, [Comment]
_, Maybe SourceType
_, Just Meta
meta) Qualified Ident
i | Meta -> Bool
isConstructorLike Meta
meta
        -> forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\f (Expr Ann)
e1 (Ann
a2, Expr Ann
e2) -> forall a. a -> Expr a -> Expr a -> Expr a
App Ann
a2 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Expr Ann)
e1 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr Ann -> f (Expr Ann)
handleExpr' Expr Ann
e2) (Int -> Maybe Int -> Ann -> Qualified Ident -> f (Expr Ann)
f Int
delay Maybe Int
force Ann
a Qualified Ident
i) [(Ann, Expr Ann)]
args
      Expr Ann
e -> forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\f (Expr Ann)
e1 (Ann
a2, Expr Ann
e2) -> forall a. a -> Expr a -> Expr a -> Expr a
App Ann
a2 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Expr Ann)
e1 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a b. (a, b) -> b
snd (Int
-> Maybe Int
-> (Bind Ann -> f (Bind Ann), Expr Ann -> f (Expr Ann))
go Int
delay forall a. Maybe a
Nothing) Expr Ann
e2) (forall a b. (a, b) -> b
snd (Int
-> Maybe Int
-> (Bind Ann -> f (Bind Ann), Expr Ann -> f (Expr Ann))
go Int
delay (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. Num a => a -> a -> a
+ Int
len) Maybe Int
force)) Expr Ann
e) [(Ann, Expr Ann)]
args
    isConstructorLike :: Meta -> Bool
isConstructorLike = \case
      IsConstructor{} -> Bool
True
      Meta
IsNewtype -> Bool
True
      Meta
_ -> Bool
False

-- Once we assign a delay and force value to every `Var` in the binding group,
-- we can consider how to order the bindings to allow them all to successfully
-- initialize. There is one principle here: each binding must be initialized
-- before the identifier being bound is ready for use. If the preorder thus
-- induced has cycles, those cycles need to be resolved with laziness. All of
-- the details concern what "ready for use" means.
--
-- The definition of delay and force suggests that "ready for use" depends on
-- those attributes. If a lambda is bound to the name x, then the references in
-- the lambda don't need to be initialized before x is initialized. This is
-- represented by the fact that those references have non-zero delay. But if
-- the expression bound to x is instead the application of a function y that is
-- also bound in this binding group, then not only does y need to be
-- initialized before x, so do some of the non-zero delay references in y. This
-- is represented by the fact that the occurrence of y in the expression bound
-- to x has non-zero force.
--
-- An example, reusing the pseudocode annotations defined above:
--
--   x _ = y#1!0
--   y = x#0!1 a
--
-- y doesn't need to be initialized before x is, because the reference to y in
-- x's initializer has delay 1. But y does need to be initialized before x is
-- ready for use with force 1, because force 1 is enough to overcome the delay
-- of that reference. And since y has a delay-0 reference to x with force 1, y
-- will need to be ready for use before it is initialized; thus, y needs to be
-- made lazy.
--
-- So just as function applications "cancel out" lambdas, a known applied force
-- cancels out an equal amount of delay, causing some references that may not
-- have been needed earlier to enter play. (And to be safe, we must assume that
-- unknown force cancels out *any* amount of delay.) There is another, subtler
-- aspect of this: if there are not enough lambdas to absorb every argument
-- applied to a function, those arguments will end up applied to the result of
-- the function. Likewise, if there is excess force left over after some of it
-- has been canceled by delay, that excess is carried to the references
-- activated. (Again, an unknown amount of force must be assumed to lead to an
-- unknown amount of excess force.)
--
-- Another example:
--
--   f = g#0!2 a b
--   g x = h#1!2 c x
--   h _ _ _ = f#3!0
--
-- Initializing f will lead to an infinite loop in this example. f invokes g
-- with two arguments. g absorbs one argument, and the second ends up being
-- applied to the result of h c x, resulting in h being invoked with three
-- arguments. Invoking h with three arguments results in dereferencing f, which
-- is not yet ready. To capture this loop in our analysis, we say that making
-- f ready for use with force 0 requires making g ready for use with force 2,
-- which requires making h ready for use with force 3 (two units of force from
-- the lexical position of h, plus one unit of excess force carried forward),
-- which cyclically requires f to be ready for use with force 0.
--
-- These preceding observations are captured and generalized by the following
-- rules:
--
--   USE-INIT: Before a reference to x is ready for use with any force, x must
--     be initialized.
--
--     We will make x lazy iff this rule induces a cycle--i.e., initializing x
--     requires x to be ready for use first.
--
--   USE-USE: Before a reference to x is ready for use with force f:
--     * if a reference in the initializer of x has delay d and force f',
--     * and either d <= f or f is unknown,
--     * then that reference must itself be ready for use with
--       force f – d + f' (or with unknown force if f or f' is unknown).
--
--   USE-IMMEDIATE: Initializing a binding x is equivalent to requiring a
--     reference to x to be ready for use with force 0, per USE-USE.
--     
--     Equivalently: before x is initialized, any reference in the initializer
--     of x with delay 0 and force f must be ready for use with force f.
--
-- Examples:
--
--   Assume x is bound in a recursive binding group with the below bindings.
--
--   All of the following initializers require x to be ready for use with some
--   amount of force, and therefore require x to be initialized first.
--
--   a = x#0!0
--   b = (\_ -> x#0!0) 1
--   c = foo x#0!*
--   d = (\_ -> foo x#0!*) 1
--
--   In the following initializers, before p can be initialized, x must be
--   ready for use with force f – d + f'. (And both x and q must be
--   initialized, of course; but x being ready for use with that force may
--   induce additional constraints.)
--
--   p = ... q#0!f ...
--   q = ... x#d!f' ... (where d <= f)
--
--   Excess force stacks, of course: in the following initializers, before r
--   can be initialized, x must be ready for use with force
--   f — d + f' — d' + f'':
--
--   r = ... s#0!f ...
--   s = ... t#d!f' ... (where d <= f)
--   t = ... x#d'!f'' ... (where d' <= f – d + f')
--
--
-- To satisfy these rules, we will construct a graph between (identifier,
-- delay) pairs, with edges induced by the USE-USE rule, and effectively run a
-- topsort to get the initialization preorder. For this part, it's simplest to
-- think of delay as an element of the naturals extended with a positive
-- infinity, corresponding to an unknown amount of force. (We'll do arithmetic
-- on these extended naturals as you would naively expect; we won't do anything
-- suspect like subtracting infinity from infinity.) With that in mind, we can
-- construct the graph as follows: for each reference from i1 to i2 with delay
-- d and force f, draw an infinite family of edges from (i1, d + n) to (i2, f +
-- n) for all 0 <= n <= ∞, where n represents the excess force carried over
-- from a previous edge. Unfortunately, as an infinite graph, we can't expect
-- the tools in Data.Graph to help us traverse it; we will have to be a little
-- bit clever.
--
-- The following data types and functions are for searching this infinite graph
-- and carving from it a finite amount of data to work with. Specifically, we
-- want to know for each identifier i, which other identifiers are
-- irreflexively reachable from (i, 0) (and thus must be initialized before i
-- is), and with what maximum force (in the event of a loop, not every
-- reference to i in the reachable identifier needs to be rewritten to a force
-- call; only the ones with delay up to the maximum force used during i's
-- initialization). We also want the option of aborting a given reachability
-- search, for one of two reasons.
--
--   * If we encounter a reference with unknown force, abort.
--   * If we encounter a cycle where force on a single identifier is
--     increasing, abort. (Because of USE-USE, as soon as an identifier is
--     revisited with greater force than its first visit, the difference is
--     carried forward as excess, so it is possible to retrace that path to get
--     an arbitrarily high amount of force.)
--
-- Both reasons mean that it is theoretically possible for the identifier in
-- question to need every other identifier in the binding group to be
-- initialized before it is. (Every identifier in a recursive binding group is
-- necessarily reachable from every other, ignoring delay and force, which is
-- what arbitrarily high force lets you do.)
--
-- In order to reuse parts of this reachability computation across identifiers,
-- we are going to represent it with a rose tree data structure interleaved with
-- a monad capturing the abort semantics. (The monad is Maybe, but we don't
-- need to know that here!)

type MaxRoseTree m a = m (IM.MonoidalIntMap (MaxRoseNode m a))
data MaxRoseNode m a = MaxRoseNode a (MaxRoseTree m a)

-- Dissecting this data structure:
--
-- m (...)
-- ^ represents whether to abort or continue the search
--
--   IM.MonoidalIntMap (...)
--   ^ the keys of this map are other identifiers reachable from the current
--     one (we'll map the identifiers in this binding group to Ints for ease of
--     computation)
--
--     the values of this map are:
--
--     MaxRoseNode a (...)
--     ^ this will store the force applied to the next identifier
--       (MaxRoseTree m a)
--       ^ and this, the tree of identifiers reachable from there
--
-- We're only interested in continuing down the search path that applies the
-- most force to a given identifier! So when we combine two MaxRoseTrees,
-- we want to resolve any key collisions in their MonoidalIntMaps with this
-- semigroup:

instance Ord a => Semigroup (MaxRoseNode m a) where
  l :: MaxRoseNode m a
l@(MaxRoseNode a
l1 MaxRoseTree m a
_) <> :: MaxRoseNode m a -> MaxRoseNode m a -> MaxRoseNode m a
<> r :: MaxRoseNode m a
r@(MaxRoseNode a
r1 MaxRoseTree m a
_) = if a
r1 forall a. Ord a => a -> a -> Bool
> a
l1 then MaxRoseNode m a
r else MaxRoseNode m a
l

-- And that's why this is called a MaxRoseTree.
--
-- Traversing this tree to get a single MonoidalIntMap with the entire closure
-- plus force information is fairly straightforward:

mrtFlatten :: (Monad m, Ord a) => MaxRoseTree m a -> m (IM.MonoidalIntMap (Max a))
mrtFlatten :: forall (m :: * -> *) a.
(Monad m, Ord a) =>
MaxRoseTree m a -> m (MonoidalIntMap (Max a))
mrtFlatten = (forall {k} (f :: k -> *) (a :: k). Ap f a -> f a
getAp forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a m. Monoid m => (Int -> a -> m) -> MonoidalIntMap a -> m
IM.foldMapWithKey (\Int
i (MaxRoseNode a
a m (MonoidalIntMap (MaxRoseNode m a))
inner) -> forall {k} (f :: k -> *) (a :: k). f a -> Ap f a
Ap forall a b. (a -> b) -> a -> b
$ (forall a. Int -> a -> MonoidalIntMap a
IM.singleton Int
i (forall a. a -> Max a
Max a
a) forall a. Semigroup a => a -> a -> a
<>) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a.
(Monad m, Ord a) =>
MaxRoseTree m a -> m (MonoidalIntMap (Max a))
mrtFlatten m (MonoidalIntMap (MaxRoseNode m a))
inner) forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<)

-- The use of the `Ap` monoid ensures that if any child of this tree aborts,
-- the entire tree aborts.
--
-- One might ask, why interleave the abort monad with the tree at all if we're
-- just going to flatten it out at the end? The point is to flatten it out at
-- the end, but *not* during the generation of the tree. Attempting to flatten
-- the tree as we generate it can result in an infinite loop, because a subtree
-- needs to be exhaustively searched for abort conditions before it can be used
-- in another tree. With this approach, we can use lazy trees as building
-- blocks and, as long as they get rewritten to be finite or have aborts before
-- they're flattened, the analysis still terminates.

-- |
-- Given a maximum index and a function that returns a map of edges to next
-- indices, returns an array for each index up to maxIndex of maps from the
-- indices reachable from the current index, to the maximum force applied to
-- those indices.
searchReachable
  :: forall m force
   . (Alternative m, Monad m, Enum force, Ord force)
  => Int
  -> ((Int, force) -> m (IM.MonoidalIntMap (Max force)))
  -> A.Array Int (m (IM.MonoidalIntMap (Max force)))
searchReachable :: forall (m :: * -> *) force.
(Alternative m, Monad m, Enum force, Ord force) =>
Int
-> ((Int, force) -> m (MonoidalIntMap (Max force)))
-> Array Int (m (MonoidalIntMap (Max force)))
searchReachable Int
maxIdx (Int, force) -> m (MonoidalIntMap (Max force))
lookupEdges = forall (m :: * -> *) a.
(Monad m, Ord a) =>
MaxRoseTree m a -> m (MonoidalIntMap (Max a))
mrtFlatten forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. HasCallStack => [a] -> a
unsafeHead forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Array Int [MaxRoseTree m force]
mem
  where
  -- This is a finite array of infinite lists, used to memoize all the search
  -- trees. `unsafeHead` is used above to pull the first tree out of each list
  -- in the array--the one corresponding to zero force, which is what's needed
  -- to initialize the corresponding identifier. (`unsafeHead` is safe here, of
  -- course: infinite lists.)
  mem :: A.Array Int [MaxRoseTree m force]
  mem :: Array Int [MaxRoseTree m force]
mem = forall i e. Ix i => (i, i) -> [e] -> Array i e
A.listArray (Int
0, Int
maxIdx)
    [ [(Int, force) -> MaxRoseTree m force -> MaxRoseTree m force
cutLoops forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. (Int -> a -> b) -> MonoidalIntMap a -> MonoidalIntMap b
IM.mapWithKey Int -> Max force -> MaxRoseNode m force
memoizedNode) forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, force) -> m (MonoidalIntMap (Max force))
lookupEdges forall a b. (a -> b) -> a -> b
$ (Int
i, force
f) | force
f <- [forall a. Enum a => Int -> a
toEnum Int
0..]]
    | Int
i <- [Int
0..Int
maxIdx]
    ]

  memoizedNode :: Int -> Max force -> MaxRoseNode m force
  memoizedNode :: Int -> Max force -> MaxRoseNode m force
memoizedNode Int
i (Max force
force) = forall (m :: * -> *) a. a -> MaxRoseTree m a -> MaxRoseNode m a
MaxRoseNode force
force forall a b. (a -> b) -> a -> b
$ Array Int [MaxRoseTree m force]
mem forall i e. Ix i => Array i e -> i -> e
A.! Int
i forall a. [a] -> Int -> a
!! forall a. Enum a => a -> Int
fromEnum force
force

  -- And this is the function that prevents the search from actually being
  -- infinite. It applies a filter to a `MaxRoseTree` at every level, looking for
  -- indices anywhere in the tree that match the current vertex. If a match is
  -- found with greater force than the current force, that part of the tree is
  -- rewritten to abort; otherwise, that part of the tree is rewritten to be
  -- empty (there's nothing new in that part of the search).
  --
  -- A new version of `cutLoops` is applied for each node in the search, so
  -- each edge in a search path will add another filter on a new index. Since
  -- there are a finite number of indices in our universe, this guarantees that
  -- the analysis terminates, because no single search path can have length
  -- greater than `maxIdx`.
  cutLoops :: (Int, force) -> MaxRoseTree m force -> MaxRoseTree m force
  cutLoops :: (Int, force) -> MaxRoseTree m force -> MaxRoseTree m force
cutLoops (Int
i, force
force) = MaxRoseTree m force -> MaxRoseTree m force
go
    where
    go :: MaxRoseTree m force -> MaxRoseTree m force
go = forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
(=<<) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a b.
Applicative t =>
(Int -> a -> t b) -> MonoidalIntMap a -> t (MonoidalIntMap b)
IM.traverseWithKey forall a b. (a -> b) -> a -> b
$ \Int
i' (MaxRoseNode force
force' MaxRoseTree m force
inner) ->
      forall (m :: * -> *) a. a -> MaxRoseTree m a -> MaxRoseNode m a
MaxRoseNode force
force' forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> if Int
i forall a. Eq a => a -> a -> Bool
== Int
i' then forall (f :: * -> *). Alternative f => Bool -> f ()
guard (force
force forall a. Ord a => a -> a -> Bool
>= force
force') forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. MonoidalIntMap a
IM.empty else forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ MaxRoseTree m force -> MaxRoseTree m force
go MaxRoseTree m force
inner

-- One last data structure to define and then it's on to the main event.
--
-- The laziness transform effectively takes a list of eager bindings (x = ...)
-- and splits some of them into lazy definitions ($lazy_x = ...) and lazy
-- bindings (x = $lazy_x ...). It's convenient to work with these three
-- declarations as the following sum type:

data RecursiveGroupItem e = EagerBinding Ann e | LazyDefinition e | LazyBinding Ann
  deriving forall a b. a -> RecursiveGroupItem b -> RecursiveGroupItem a
forall a b.
(a -> b) -> RecursiveGroupItem a -> RecursiveGroupItem b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> RecursiveGroupItem b -> RecursiveGroupItem a
$c<$ :: forall a b. a -> RecursiveGroupItem b -> RecursiveGroupItem a
fmap :: forall a b.
(a -> b) -> RecursiveGroupItem a -> RecursiveGroupItem b
$cfmap :: forall a b.
(a -> b) -> RecursiveGroupItem a -> RecursiveGroupItem b
Functor

-- |
-- Transform a recursive binding group, reordering the bindings within when a
-- correct initialization order can be statically determined, and rewriting
-- bindings and references to be lazy otherwise.
--
applyLazinessTransform :: ModuleName -> [((Ann, Ident), Expr Ann)] -> ([((Ann, Ident), Expr Ann)], Any)
applyLazinessTransform :: ModuleName
-> [((Ann, Ident), Expr Ann)] -> ([((Ann, Ident), Expr Ann)], Any)
applyLazinessTransform ModuleName
mn [((Ann, Ident), Expr Ann)]
rawItems = let

  -- Establish the mapping from names to ints.
  rawItemsByName :: M.MonoidalMap Ident (Ann, Expr Ann)
  rawItemsByName :: MonoidalMap Ident (Ann, Expr Ann)
rawItemsByName = forall k a. Ord k => [(k, a)] -> MonoidalMap k a
M.fromList forall a b. (a -> b) -> a -> b
$ (forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst forall (a :: * -> * -> *) b c c'.
Arrow a =>
a b c -> a b c' -> a b (c, c')
&&& forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first forall a b. (a, b) -> a
fst) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [((Ann, Ident), Expr Ann)]
rawItems

  maxIdx :: Int
maxIdx = forall k a. MonoidalMap k a -> Int
M.size MonoidalMap Ident (Ann, Expr Ann)
rawItemsByName forall a. Num a => a -> a -> a
- Int
1

  rawItemsByIndex :: A.Array Int (Ann, Expr Ann)
  rawItemsByIndex :: Array Int (Ann, Expr Ann)
rawItemsByIndex = forall i e. Ix i => (i, i) -> [e] -> Array i e
A.listArray (Int
0, Int
maxIdx) forall a b. (a -> b) -> a -> b
$ forall k a. MonoidalMap k a -> [a]
M.elems MonoidalMap Ident (Ann, Expr Ann)
rawItemsByName

  names :: S.Set Ident
  names :: Set Ident
names = forall k a. MonoidalMap k a -> Set k
M.keysSet MonoidalMap Ident (Ann, Expr Ann)
rawItemsByName

  -- Now do the first delay/force traversal of all the bindings to find
  -- references to other names in this binding group.
  --
  -- The parts of this type mean:
  -- D is the maximum force (or Nothing if unknown) with which the identifier C
  -- is referenced in any delay-B position inside the expression A.
  --
  -- where A, B, C, and D are as below:
  --                A           B (keys)           C (keys)           D
  findReferences :: Expr Ann -> IM.MonoidalIntMap (IM.MonoidalIntMap (Ap Maybe (Max Int)))
  findReferences :: Expr Ann -> MonoidalIntMap (MonoidalIntMap (Ap Maybe (Max Int)))
findReferences = (forall {k} a (b :: k). Const a b -> a
getConst forall b c a. (b -> c) -> (a -> b) -> a -> c
.) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *).
Applicative f =>
(Int -> Maybe Int -> Ann -> Qualified Ident -> f (Expr Ann))
-> Expr Ann -> f (Expr Ann)
onVarsWithDelayAndForce forall a b. (a -> b) -> a -> b
$ \Int
delay Maybe Int
force Ann
_ -> \case
    Qualified QualifiedBy
qb Ident
ident | forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (forall a. Eq a => a -> a -> Bool
== ModuleName
mn) (QualifiedBy -> Maybe ModuleName
toMaybeModuleName QualifiedBy
qb), Just Int
i <- Ident
ident forall a. Ord a => a -> Set a -> Maybe Int
`S.lookupIndex` Set Ident
names
      -> forall {k} a (b :: k). a -> Const a b
Const forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Int -> a -> MonoidalIntMap a
IM.singleton Int
delay forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Int -> a -> MonoidalIntMap a
IM.singleton Int
i forall a b. (a -> b) -> a -> b
$ Maybe Int -> Ap Maybe (Max Int)
coerceForce Maybe Int
force
    Qualified Ident
_ -> forall {k} a (b :: k). a -> Const a b
Const forall a. MonoidalIntMap a
IM.empty

  -- The parts of this type mean:
  -- D is the maximum force (or Nothing if unknown) with which the identifier C
  -- is referenced in any delay-B position inside the binding of identifier A.
  --
  -- where A, B, C, and D are as below:
  --                     A    B (keys)           C (keys)           D
  refsByIndex :: A.Array Int (IM.MonoidalIntMap (IM.MonoidalIntMap (Ap Maybe (Max Int))))
  refsByIndex :: Array Int (MonoidalIntMap (MonoidalIntMap (Ap Maybe (Max Int))))
refsByIndex = Expr Ann -> MonoidalIntMap (MonoidalIntMap (Ap Maybe (Max Int)))
findReferences forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Array Int (Ann, Expr Ann)
rawItemsByIndex

  -- Using the approach explained above, traverse the reference graph generated
  -- by `refsByIndex` and find all reachable names.
  --
  -- The parts of this type mean:
  -- D is the maximum force with which the identifier C is referenced,
  -- directly or indirectly, during the initialization of identifier A. B is
  -- Nothing if the analysis of A was inconclusive and A might need the entire
  -- binding group.
  -- 
  -- where A, B, C, and D are as below:
  --                           A    B      C (keys)           D
  reachablesByIndex :: A.Array Int (Maybe (IM.MonoidalIntMap (Max Int)))
  reachablesByIndex :: Array Int (Maybe (MonoidalIntMap (Max Int)))
reachablesByIndex = forall (m :: * -> *) force.
(Alternative m, Monad m, Enum force, Ord force) =>
Int
-> ((Int, force) -> m (MonoidalIntMap (Max force)))
-> Array Int (m (MonoidalIntMap (Max force)))
searchReachable Int
maxIdx forall a b. (a -> b) -> a -> b
$ \(Int
i, Int
force) ->
    forall {k} (f :: k -> *) (a :: k). Ap f a -> f a
getAp forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a m. Monoid m => (Int -> a -> m) -> MonoidalIntMap a -> m
IM.foldMapWithKey (forall a. Int -> MonoidalIntMap a -> MonoidalIntMap a
dropKeysAbove Int
force forall a b. (a -> b) -> a -> b
$ Array Int (MonoidalIntMap (MonoidalIntMap (Ap Maybe (Max Int))))
refsByIndex forall i e. Ix i => Array i e -> i -> e
A.! Int
i) forall a b. (a -> b) -> a -> b
$ \Int
delay ->
      forall a m. Monoid m => (Int -> a -> m) -> MonoidalIntMap a -> m
IM.foldMapWithKey forall a b. (a -> b) -> a -> b
$ \Int
i' Ap Maybe (Max Int)
force' ->
        forall {k} (f :: k -> *) (a :: k). f a -> Ap f a
Ap forall a b. (a -> b) -> a -> b
$ forall a. Int -> a -> MonoidalIntMap a
IM.singleton Int
i' forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> Max a
Max forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int
force forall a. Num a => a -> a -> a
- Int
delay forall a. Num a => a -> a -> a
+) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ap Maybe (Max Int) -> Maybe Int
uncoerceForce Ap Maybe (Max Int)
force'

  -- If `reachablesByIndex` is a sort of labeled relation, this function
  -- produces part of the reverse relation, but only for the edges from the
  -- given vertex.
  --
  -- The parts of this type mean:
  -- The identifier A is reachable from the identifier B with maximum force C
  -- (B is also the index provided to the function).
  --
  -- where A, B, and C are as below:
  --                      (B)    A                  B (singleton key)  C
  reverseReachablesFor :: Int -> IM.MonoidalIntMap (IM.MonoidalIntMap (Ap Maybe (Max Int)))
  reverseReachablesFor :: Int -> MonoidalIntMap (MonoidalIntMap (Ap Maybe (Max Int)))
reverseReachablesFor Int
i = case Array Int (Maybe (MonoidalIntMap (Max Int)))
reachablesByIndex forall i e. Ix i => Array i e -> i -> e
A.! Int
i of
    Maybe (MonoidalIntMap (Max Int))
Nothing -> forall a. [(Int, a)] -> MonoidalIntMap a
IM.fromAscList forall a b. (a -> b) -> a -> b
$ (, forall a. Int -> a -> MonoidalIntMap a
IM.singleton Int
i forall a b. (a -> b) -> a -> b
$ forall {k} (f :: k -> *) (a :: k). f a -> Ap f a
Ap forall a. Maybe a
Nothing) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int
0..Int
maxIdx]
    Just MonoidalIntMap (Max Int)
im -> forall a. Int -> a -> MonoidalIntMap a
IM.singleton Int
i forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (f :: k -> *) (a :: k). f a -> Ap f a
Ap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> Maybe a
Just forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MonoidalIntMap (Max Int)
im

  -- We can use `reachablesByIndex` to build a finite graph and topsort it;
  -- in the process, we'll pack the nodes of the graph with data we'll want
  -- next. Remember that if our reachability computation aborted, we have to
  -- assume that every other identifier is reachable from that one--hence the
  -- `maybe [0..maxIdx]`.
  sccs :: [SCC
   (MonoidalIntMap (MonoidalIntMap (Ap Maybe (Max Int))),
    (Ident, (Ann, Expr Ann)))]
sccs = forall key node. Ord key => [(node, key, [key])] -> [SCC node]
stronglyConnComp forall a b. (a -> b) -> a -> b
$ do
    (Int
i, Maybe (MonoidalIntMap (Max Int))
mbReachable) <- forall i e. Ix i => Array i e -> [(i, e)]
A.assocs Array Int (Maybe (MonoidalIntMap (Max Int)))
reachablesByIndex
    forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Int -> MonoidalIntMap (MonoidalIntMap (Ap Maybe (Max Int)))
reverseReachablesFor Int
i, (forall a. Int -> Set a -> a
S.elemAt Int
i Set Ident
names, Array Int (Ann, Expr Ann)
rawItemsByIndex forall i e. Ix i => Array i e -> i -> e
A.! Int
i)), Int
i, forall b a. b -> (a -> b) -> Maybe a -> b
maybe [Int
0..Int
maxIdx] (IntSet -> [Int]
IS.toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. MonoidalIntMap a -> IntSet
IM.keysSet) Maybe (MonoidalIntMap (Max Int))
mbReachable)

  (MonoidalIntMap (MonoidalIntMap (Ap Maybe (Max Int)))
replacements, [(Ident, RecursiveGroupItem (Expr Ann))]
items) = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap [SCC
   (MonoidalIntMap (MonoidalIntMap (Ap Maybe (Max Int))),
    (Ident, (Ann, Expr Ann)))]
sccs forall a b. (a -> b) -> a -> b
$ \case
    -- The easy case: this binding doesn't need to be made lazy after all!
    AcyclicSCC (MonoidalIntMap (MonoidalIntMap (Ap Maybe (Max Int)))
_, (Ident
ident, (Ann
a, Expr Ann
e))) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Ident
ident, forall e. Ann -> e -> RecursiveGroupItem e
EagerBinding Ann
a Expr Ann
e)]
    -- The tough case: we have a loop.
    -- We need to do two things here:
    --   * Collect the reversed reachables relation for each vertex in this
    --     loop; we'll use this to replace references with force calls
    --   * Copy the vertex list into two lists: a list of lazy definitions and
    --     a list of lazy bindings
    -- Both of these results are monoidal, so the outer `foldMap` will
    -- concatenate them pairwise.
    CyclicSCC [(MonoidalIntMap (MonoidalIntMap (Ap Maybe (Max Int))),
  (Ident, (Ann, Expr Ann)))]
vertices -> (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall a b. (a, b) -> a
fst [(MonoidalIntMap (MonoidalIntMap (Ap Maybe (Max Int))),
  (Ident, (Ann, Expr Ann)))]
vertices, forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall e. e -> RecursiveGroupItem e
LazyDefinition forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(MonoidalIntMap (MonoidalIntMap (Ap Maybe (Max Int))),
  (Ident, (Ann, Expr Ann)))]
vertices forall a. [a] -> [a] -> [a]
++ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall e. Ann -> RecursiveGroupItem e
LazyBinding forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(MonoidalIntMap (MonoidalIntMap (Ap Maybe (Max Int))),
  (Ident, (Ann, Expr Ann)))]
vertices)

  -- We have `replacements` expressed in terms of indices; we want to map it
  -- back to names before traversing the bindings again.
  replacementsByName :: M.MonoidalMap Ident (M.MonoidalMap Ident (Ap Maybe (Max Int)))
  replacementsByName :: MonoidalMap Ident (MonoidalMap Ident (Ap Maybe (Max Int)))
replacementsByName = forall k a. Eq k => [(k, a)] -> MonoidalMap k a
M.fromAscList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
map (forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Int -> Set a -> a
S.elemAt Set Ident
names) (forall k a. Eq k => [(k, a)] -> MonoidalMap k a
M.fromAscList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
map (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Int -> Set a -> a
S.elemAt Set Ident
names)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. MonoidalIntMap a -> [(Int, a)]
IM.toAscList)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. MonoidalIntMap a -> [(Int, a)]
IM.toAscList forall a b. (a -> b) -> a -> b
$ MonoidalIntMap (MonoidalIntMap (Ap Maybe (Max Int)))
replacements

  -- And finally, this is the second delay/force traversal where we take
  -- `replacementsByName` and use it to rewrite references with force calls,
  -- but only if the delay of those references is at most the maximum amount
  -- of force used by the initialization of the referenced binding to
  -- reference the outer binding. A reference made with a higher delay than
  -- that can safely continue to use the original reference, since it won't be
  -- needed until after the referenced binding is done initializing.
  replaceReferencesWithForceCall :: (Ident, RecursiveGroupItem (Expr Ann)) -> (Ident, RecursiveGroupItem (Expr Ann))
  replaceReferencesWithForceCall :: (Ident, RecursiveGroupItem (Expr Ann))
-> (Ident, RecursiveGroupItem (Expr Ann))
replaceReferencesWithForceCall pair :: (Ident, RecursiveGroupItem (Expr Ann))
pair@(Ident
ident, RecursiveGroupItem (Expr Ann)
item) = case Ident
ident forall k a. Ord k => k -> MonoidalMap k a -> Maybe a
`M.lookup` MonoidalMap Ident (MonoidalMap Ident (Ap Maybe (Max Int)))
replacementsByName of
    Maybe (MonoidalMap Ident (Ap Maybe (Max Int)))
Nothing -> (Ident, RecursiveGroupItem (Expr Ann))
pair
    Just MonoidalMap Ident (Ap Maybe (Max Int))
m -> let
      rewriteExpr :: Expr Ann -> Expr Ann
rewriteExpr = (forall a. Identity a -> a
runIdentity forall b c a. (b -> c) -> (a -> b) -> a -> c
.) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *).
Applicative f =>
(Int -> Maybe Int -> Ann -> Qualified Ident -> f (Expr Ann))
-> Expr Ann -> f (Expr Ann)
onVarsWithDelayAndForce forall a b. (a -> b) -> a -> b
$ \Int
delay Maybe Int
_ Ann
ann -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. \case
        Qualified QualifiedBy
qb Ident
ident' | forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (forall a. Eq a => a -> a -> Bool
== ModuleName
mn) (QualifiedBy -> Maybe ModuleName
toMaybeModuleName QualifiedBy
qb), forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (forall a. Ord a => a -> a -> Bool
>= forall a. a -> Max a
Max Int
delay) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (f :: k -> *) (a :: k). Ap f a -> f a
getAp) forall a b. (a -> b) -> a -> b
$ Ident
ident' forall k a. Ord k => k -> MonoidalMap k a -> Maybe a
`M.lookup` MonoidalMap Ident (Ap Maybe (Max Int))
m
          -> Ann -> Ident -> Expr Ann
makeForceCall Ann
ann Ident
ident'
        Qualified Ident
q -> forall a. a -> Qualified Ident -> Expr a
Var Ann
ann Qualified Ident
q
      in (Ident
ident, Expr Ann -> Expr Ann
rewriteExpr forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> RecursiveGroupItem (Expr Ann)
item)

  -- All that's left to do is run the above replacement on every item,
  -- translate items from our `RecursiveGroupItem` representation back into the
  -- form CoreFn expects, and inform the caller whether we made any laziness
  -- transformations after all. (That last bit of information is used to
  -- determine if the runtime factory function needs to be injected.)
  in (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Ident -> RecursiveGroupItem (Expr Ann) -> ((Ann, Ident), Expr Ann)
fromRGI forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Ident, RecursiveGroupItem (Expr Ann))
-> (Ident, RecursiveGroupItem (Expr Ann))
replaceReferencesWithForceCall forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Ident, RecursiveGroupItem (Expr Ann))]
items, Bool -> Any
Any forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ forall a. MonoidalIntMap a -> Bool
IM.null MonoidalIntMap (MonoidalIntMap (Ap Maybe (Max Int)))
replacements)

  where

  nullAnn :: Ann
nullAnn = SourceSpan -> Ann
ssAnn SourceSpan
nullSourceSpan
  runtimeLazy :: Expr Ann
runtimeLazy = forall a. a -> Qualified Ident -> Expr a
Var Ann
nullAnn forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. QualifiedBy -> a -> Qualified a
Qualified QualifiedBy
ByNullSourcePos forall a b. (a -> b) -> a -> b
$ InternalIdentData -> Ident
InternalIdent InternalIdentData
RuntimeLazyFactory
  runFn3 :: Expr Ann
runFn3 = forall a. a -> Qualified Ident -> Expr a
Var Ann
nullAnn forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. QualifiedBy -> a -> Qualified a
Qualified (ModuleName -> QualifiedBy
ByModuleName ModuleName
C.M_Data_Function_Uncurried) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Ident
Ident forall a b. (a -> b) -> a -> b
$ forall a. (Eq a, IsString a) => a
C.S_runFn forall a. Semigroup a => a -> a -> a
<> Text
"3"
  strLit :: Text -> Expr Ann
strLit = forall a. a -> Literal (Expr a) -> Expr a
Literal Ann
nullAnn forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. PSString -> Literal a
StringLiteral forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> PSString
mkString

  lazifyIdent :: Ident -> Ident
lazifyIdent = \case
    Ident Text
txt -> InternalIdentData -> Ident
InternalIdent forall a b. (a -> b) -> a -> b
$ Text -> InternalIdentData
Lazy Text
txt
    Ident
_ -> forall a. HasCallStack => String -> a
internalError String
"Unexpected argument to lazifyIdent"

  makeForceCall :: Ann -> Ident -> Expr Ann
  makeForceCall :: Ann -> Ident -> Expr Ann
makeForceCall (SourceSpan
ss, [Comment]
_, Maybe SourceType
_, Maybe Meta
_) Ident
ident
    -- We expect the functions produced by `runtimeLazy` to accept one
    -- argument: the line number on which this reference is made. The runtime
    -- code uses this number to generate a message that identifies where the
    -- evaluation looped.
    = forall a. a -> Expr a -> Expr a -> Expr a
App Ann
nullAnn (forall a. a -> Qualified Ident -> Expr a
Var Ann
nullAnn forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. QualifiedBy -> a -> Qualified a
Qualified QualifiedBy
ByNullSourcePos forall a b. (a -> b) -> a -> b
$ Ident -> Ident
lazifyIdent Ident
ident)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> Literal (Expr a) -> Expr a
Literal Ann
nullAnn forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Either Integer Double -> Literal a
NumericLiteral forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. a -> Either a b
Left forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Integral a => a -> Integer
toInteger forall b c a. (b -> c) -> (a -> b) -> a -> c
. SourcePos -> Int
sourcePosLine
    forall a b. (a -> b) -> a -> b
$ SourceSpan -> SourcePos
spanStart SourceSpan
ss

  fromRGI :: Ident -> RecursiveGroupItem (Expr Ann) -> ((Ann, Ident), Expr Ann)
  fromRGI :: Ident -> RecursiveGroupItem (Expr Ann) -> ((Ann, Ident), Expr Ann)
fromRGI Ident
i = \case
    EagerBinding Ann
a Expr Ann
e -> ((Ann
a, Ident
i), Expr Ann
e)
    -- We expect the `runtimeLazy` factory to accept three arguments: the
    -- identifier being initialized, the name of the module, and of course a
    -- thunk that actually contains the initialization code.
    LazyDefinition Expr Ann
e -> ((Ann
nullAnn, Ident -> Ident
lazifyIdent Ident
i), forall a. (a -> a -> a) -> [a] -> a
foldl1' (forall a. a -> Expr a -> Expr a -> Expr a
App Ann
nullAnn) [Expr Ann
runFn3, Expr Ann
runtimeLazy, Text -> Expr Ann
strLit forall a b. (a -> b) -> a -> b
$ Ident -> Text
runIdent Ident
i, Text -> Expr Ann
strLit forall a b. (a -> b) -> a -> b
$ ModuleName -> Text
runModuleName ModuleName
mn, forall a. a -> Ident -> Expr a -> Expr a
Abs Ann
nullAnn Ident
UnusedIdent Expr Ann
e])
    LazyBinding Ann
a -> ((Ann
a, Ident
i), Ann -> Ident -> Expr Ann
makeForceCall Ann
a Ident
i)

  dropKeysAbove :: Int -> IM.MonoidalIntMap a -> IM.MonoidalIntMap a
  dropKeysAbove :: forall a. Int -> MonoidalIntMap a -> MonoidalIntMap a
dropKeysAbove Int
n = forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a.
Int -> MonoidalIntMap a -> (MonoidalIntMap a, MonoidalIntMap a)
IM.split (Int
n forall a. Num a => a -> a -> a
+ Int
1)

  coerceForce :: Maybe Int -> Ap Maybe (Max Int)
  coerceForce :: Maybe Int -> Ap Maybe (Max Int)
coerceForce = coerce :: forall a b. Coercible a b => a -> b
coerce

  uncoerceForce :: Ap Maybe (Max Int) -> Maybe Int
  uncoerceForce :: Ap Maybe (Max Int) -> Maybe Int
uncoerceForce = coerce :: forall a b. Coercible a b => a -> b
coerce