[Perf] compiler: eliminate the --no-opt self-compilation memory floor & stack-unsafety (#40)

compiler/spago.yaml

@@ -17,6 +17,7 @@ package:
     - foldable-traversable
     - foreign
     - foreign-object
+    - free
     - functions
     - identity
     - integers
@@ -27,6 +28,7 @@ package:
     - node-buffer
     - prelude
     - strings
+    - tailrec
     - transformers
     - tuples
   test:

compiler/src/PureScript/Backend/Wasm/Codegen.purs

@@ -28,8 +28,11 @@ module PureScript.Backend.Wasm.Codegen
 import Prelude
 
 import Binaryen as B
+import Control.Monad.Rec.Class (Step(..), tailRecM)
 import Data.Array as Array
-import Data.Foldable (foldr, traverse_)
+import Data.Foldable (foldl, foldr, traverse_)
+import Data.List (List(..), (:))
+import Data.List as List
 import Data.Map (Map)
 import Data.Map as Map
 import Data.Maybe (Maybe(..), fromMaybe, maybe)
@@ -62,7 +65,7 @@ import PureScript.Backend.Wasm.Lower.Reps (primRep)
 buildModule :: Program -> Effect CompiledModule
 buildModule prog = do
   st <- initCodegen prog
-  traverse_ (addFunc st.ctx) prog.funcs
+  forEachArr_ prog.funcs (addFunc st.ctx)
   cafInit <- finalizeCodegen st prog
   pure { mod: st.ctx.mod, foreignModules: foreignModuleNames prog, cafInit }
 
@@ -98,7 +101,7 @@ initCodegen prog = do
 finalizeCodegen :: CodegenState -> Program -> Effect (Maybe B.Function)
 finalizeCodegen st prog = do
   cafInit <- addCafInit st.ctx st.cplan
-  traverse_ (addExportWrapper st.ctx prog.exportSigs) prog.funcs
+  forEachArr_ prog.funcs (addExportWrapper st.ctx prog.exportSigs)
   pure cafInit
 
 -- | The result of building a Binaryen module from the IR, with what packaging needs after:
@@ -220,13 +223,12 @@ buildDataTypes _ sigs0 = do
   B.typeBuilderSetStructType tb 0 [ { ty: B.i32, mutable: false } ]
   B.typeBuilderSetOpen tb 0
   baseTmp <- B.typeBuilderGetTempHeapType tb 0
-  traverse_
+  forEachArr_ (Array.mapWithIndex Tuple nonEmpty)
     ( \(Tuple i sig) -> do
         B.typeBuilderSetStructType tb (i + 1)
           (Array.cons { ty: B.i32, mutable: false } (map (\rep -> { ty: fieldWasmType rep, mutable: false }) sig))
         B.typeBuilderSetSubType tb (i + 1) baseTmp
     )
-    (Array.mapWithIndex Tuple nonEmpty)
   hts <- B.typeBuilderBuildAndDispose tb (1 + n)
   case Array.uncons hts of
     Just { head: baseHt, tail: structHts } -> do
@@ -268,7 +270,7 @@ addCounterGlobal ctx = do
   B.addGlobal ctx.mod counterGlobalName B.i32 true initE
 
 addNullaryGlobals :: Ctx -> Set Int -> Effect Unit
-addNullaryGlobals ctx tags = traverse_ addOne (Set.toUnfoldable tags :: Array Int)
+addNullaryGlobals ctx tags = forEachArr_ (Set.toUnfoldable tags :: Array Int) addOne
   where
   addOne tag = do
     tagE <- B.i32Const ctx.mod tag
@@ -288,7 +290,7 @@ cafInitName = "$caf_init"
 -- | a throwaway constant the init function overwrites at instantiation; an `eqref`
 -- | global uses a dummy boxed `0` (a GC const expression, like a nullary constructor).
 addCafGlobals :: Ctx -> Map FuncName Rep -> Effect Unit
-addCafGlobals ctx globals = traverse_ addOne (Map.toUnfoldable globals :: Array (Tuple FuncName Rep))
+addCafGlobals ctx globals = forEachArr_ (Map.toUnfoldable globals :: Array (Tuple FuncName Rep)) addOne
   where
   addOne (Tuple name rep) = do
     initE <- defaultConst ctx rep
@@ -348,24 +350,38 @@ addInternStr ctx exportIt labels = do
     dyn <- B.call ctx.mod internDynamicHelperName [ key ] B.i32
     base <- B.i32Const ctx.mod (Array.length labels)
     B.i32Add ctx.mod base dyn
-  body <- foldr step (pure miss) labels
+  -- `ifChain` (a tail loop) rather than `foldr step (pure miss)`: the whole-program label
+  -- set is large, and the fold's nested Effect binds overflow the host JS stack on a
+  -- self-sized program.
+  prepared <- traverse prepare labels
+  body <- ifChain ctx prepared miss
   _ <- B.addFunction ctx.mod internStrName (B.createType [ B.eqref ]) B.i32 [] body
   -- exported when a record/object foreign needs name→id resolution from the JS
   -- marshalling glue (ADR 0014); otherwise internal (and Binaryen-pruned if unused)
   when exportIt (void (B.addFunctionExport ctx.mod internStrName "internStr"))
   pure unit
   where
-  step (Tuple label labelId) accM = do
-    acc <- accM
+  prepare (Tuple label labelId) = do
     key <- B.localGet ctx.mod 0 B.eqref
     labelConst <- genAtom ctx (ALitString label)
     cond <- B.call ctx.mod strEqHelperName [ key, labelConst ] B.i32
     idExpr <- B.i32Const ctx.mod labelId
-    B.if_ ctx.mod cond idExpr acc
+    pure (Tuple cond idExpr)
 
 funcNameStr :: FuncName -> String
 funcNameStr (FuncName n) = n
 
+-- | Stack-safe `traverse_` over an array: the `Data.Foldable` one builds a deeply nested
+-- | `f x0 *> (f x1 *> …)` whose Effect run recurses to the array's length, overflowing the
+-- | host JS stack on whole-program-sized arrays (every function / global). `tailRecM` drives
+-- | it by index on the heap instead.
+forEachArr_ :: forall a. Array a -> (a -> Effect Unit) -> Effect Unit
+forEachArr_ arr f = tailRecM go 0
+  where
+  go i = case Array.index arr i of
+    Nothing -> pure (Done unit)
+    Just x -> f x *> pure (Loop (i + 1))
+
 -- | Add an internal function. Parameters take their declared representation (a
 -- | lifted code function's first parameter is `(ref $Clo)`); `Let`-bound locals
 -- | are all `eqref`. A code function added with `(ref $Clo, eqref) -> eqref`
@@ -480,13 +496,21 @@ addExportWrapper ctx exportSigs fn = case fn.export of
 -- | Generate a function body. `Let`s become `local.set` statements sequenced in
 -- | a `block` whose value is the tail (`Return` atom or `Switch`).
 genBody :: Ctx -> AnfExpr -> Effect B.Expression
-genBody ctx = go []
+-- Statements accumulate in a `List`, prepended most-recent-first (O(1) per `Let`) and
+-- reversed once at `seal` — an `Array` accumulator (`snoc`) copies the whole prefix per
+-- binding, i.e. O(n²) on the long `Let` chains a large function's ANF body becomes.
+--
+-- The `Let`/`LetRec`/`LetJoin` spine is walked with `tailRecM` rather than self-recursion:
+-- a large function's ANF body is a deeply nested spine, and the Effect recursion compiles to
+-- a non-tail-call chain (`__do`) that overflows the host JS stack on a self-sized program.
+-- `tailRecM` (MonadRec Effect) runs the spine on the heap, in constant stack.
+genBody ctx = tailRecM go <<< { statements: Nil, expr: _ }
   where
-  go statements = case _ of
+  go { statements, expr } = case expr of
     -- A returned atom is coerced to the function's result representation.
-    Return atom -> seal statements =<< genAtomAs ctx ctx.funcResult atom
-    Switch scrutAtom branches dflt -> seal statements =<< genSwitch ctx scrutAtom branches dflt
-    LitSwitch scrutAtom branches dflt -> seal statements =<< genLitSwitch ctx scrutAtom branches dflt
+    Return atom -> Done <$> (seal statements =<< genAtomAs ctx ctx.funcResult atom)
+    Switch scrutAtom branches dflt -> Done <$> (seal statements =<< genSwitch ctx scrutAtom branches dflt)
+    LitSwitch scrutAtom branches dflt -> Done <$> (seal statements =<< genLitSwitch ctx scrutAtom branches dflt)
     -- A direct call whose result is immediately returned is a *tail* call: emit
     -- `return_call` so a tail-recursive chain runs in constant stack. Only valid when
     -- the callee's result rep matches this function's (the frame is replaced, so no
@@ -499,29 +523,31 @@ genBody ctx = go []
       , Just sig <- Map.lookup name ctx.sigs
       , sig.result == ctx.funcResult -> do
           operands <- traverse (\(Tuple rep a) -> genAtomAs ctx rep a) (Array.zip sig.params args)
-          seal statements =<< B.returnCall ctx.mod (funcNameStr name) operands (repType ctx sig.result)
+          Done <$> (seal statements =<< B.returnCall ctx.mod (funcNameStr name) operands (repType ctx sig.result))
     -- Store the rhs into its slot, boxing/unboxing if the slot's chosen rep differs
     -- from the rhs's natural rep.
     Let (Slot index) _ rhs k -> do
       e <- genRhs ctx rhs >>= coerce ctx (rhsRep ctx rhs) (slotRep ctx index)
       stmt <- B.localSet ctx.mod index e
-      go (Array.snoc statements stmt) k
+      pure (Loop { statements: stmt : statements, expr: k })
     LetRec recBinds k -> do
       let groupSlots = map (\(RecBind (Slot s) _ _) -> s) recBinds
       allocs <- traverse (allocRecClosure ctx groupSlots) recBinds
       patches <- traverse (patchRecClosure ctx groupSlots) recBinds
-      go (statements <> allocs <> Array.concat patches) k
+      pure (Loop { statements: foldl (flip (:)) statements (allocs <> Array.concat patches), expr: k })
     -- A join point (ADR 0022): generate the `producer` as a value-producing block
     -- (its tails yield `rep`, and `return_call` is disabled so it cannot escape the
     -- function), store it into the join slot, then continue the (single) continuation.
     LetJoin (Slot slot) rep producer k -> do
       producerExpr <- genBody (ctx { funcResult = rep, tailPos = false }) producer
       stmt <- B.localSet ctx.mod slot producerExpr
-      go (Array.snoc statements stmt) k
-  -- the body / branch block produces the function's result (a tail position)
+      pure (Loop { statements: stmt : statements, expr: k })
+  -- the body / branch block produces the function's result (a tail position); `statements`
+  -- is most-recent-first, so `value : statements` reversed is emission order with the value last.
   seal statements value =
-    if Array.null statements then pure value
-    else B.block ctx.mod (Array.snoc statements value) (repType ctx ctx.funcResult)
+    case statements of
+      Nil -> pure value
+      _ -> B.block ctx.mod (Array.fromFoldable (List.reverse (value : statements))) (repType ctx ctx.funcResult)
 
 -- | Is this captured atom a forward reference to another member of the same
 -- | `LetRec` group (and thus a slot to back-patch)?
@@ -562,40 +588,55 @@ patchRecClosure ctx groupSlots (RecBind (Slot slot) _ env) =
 
 -- | A `Switch` becomes a chain of `if (tag == k) <branch> else …`, ending in the
 -- | default block or `unreachable`. The tag is read afresh per comparison.
+-- | Assemble an `if … else …` chain stack-safely from already-generated `(cond, then)`
+-- | pairs and an innermost `else` (`base`): fold from the last branch inward, so the
+-- | accumulating `else` is built by a tail loop rather than the non-tail recursion (each
+-- | level binding the recursive result before `B.if_`) that overflows the host JS stack on
+-- | a switch with many branches in a self-sized program.
+ifChain :: Ctx -> Array (Tuple B.Expression B.Expression) -> B.Expression -> Effect B.Expression
+ifChain ctx branches base = tailRecM step { acc: base, rest: Array.reverse branches }
+  where
+  step { acc, rest } = case Array.uncons rest of
+    Nothing -> pure (Done acc)
+    Just { head: Tuple cond thenE, tail } -> do
+      acc' <- B.if_ ctx.mod cond thenE acc
+      pure (Loop { acc: acc', rest: tail })
+
+dfltExpr :: Ctx -> Maybe AnfExpr -> Effect B.Expression
+dfltExpr ctx = case _ of
+  Just d -> genBody ctx d
+  Nothing -> B.unreachable ctx.mod
+
 genSwitch :: Ctx -> Atom -> Array Branch -> Maybe AnfExpr -> Effect B.Expression
-genSwitch ctx scrutAtom branches dflt = chain branches
+genSwitch ctx scrutAtom branches dflt = do
+  prepared <- traverse prepare branches
+  base <- dfltExpr ctx dflt
+  ifChain ctx prepared base
   where
   readTag = do
     s <- genAtomAs ctx Boxed scrutAtom
     c <- B.refCast ctx.mod s ctx.dataBase.ref
     B.structGet ctx.mod 0 c B.i32 false
-  chain bs = case Array.uncons bs of
-    Nothing -> case dflt of
-      Just d -> genBody ctx d
-      Nothing -> B.unreachable ctx.mod
-    Just { head: Branch tag body, tail } -> do
-      tagExpr <- readTag
-      k <- B.i32Const ctx.mod tag
-      cond <- B.i32Eq ctx.mod tagExpr k
-      thenE <- genBody ctx body
-      elseE <- chain tail
-      B.if_ ctx.mod cond thenE elseE
+  prepare (Branch tag body) = do
+    tagExpr <- readTag
+    k <- B.i32Const ctx.mod tag
+    cond <- B.i32Eq ctx.mod tagExpr k
+    thenE <- genBody ctx body
+    pure (Tuple cond thenE)
 
 -- | A `LitSwitch` becomes a chain of `if (scrutinee == literal) <branch> else …`.
 -- | The equality test unboxes the scrutinee per literal kind: `Int`/`Char` and
 -- | `Boolean` compare as `i32`, `Number` as `f64`.
 genLitSwitch :: Ctx -> Atom -> Array LitBranch -> Maybe AnfExpr -> Effect B.Expression
-genLitSwitch ctx scrutAtom branches dflt = chain branches
+genLitSwitch ctx scrutAtom branches dflt = do
+  prepared <- traverse prepare branches
+  base <- dfltExpr ctx dflt
+  ifChain ctx prepared base
   where
-  chain bs = case Array.uncons bs of
-    Nothing -> case dflt of
-      Just d -> genBody ctx d
-      Nothing -> B.unreachable ctx.mod
-    Just { head: LitBranch pat body, tail } -> do
-      cond <- litTest pat
-      thenE <- genBody ctx body
-      elseE <- chain tail
-      B.if_ ctx.mod cond thenE elseE
+  prepare (LitBranch pat body) = do
+    cond <- litTest pat
+    thenE <- genBody ctx body
+    pure (Tuple cond thenE)
   litTest = case _ of
     PInt n -> do
       s <- genAtomAs ctx I32 scrutAtom

compiler/src/PureScript/Backend/Wasm/Lower.purs

@@ -52,7 +52,7 @@ import Data.Traversable (traverse)
 import Data.Tuple (Tuple(..), fst, snd)
 import Foreign.Object (Object)
 import Foreign.Object as Object
-import PureScript.Backend.Wasm.Intrinsics (qualifiedIntrinsic, foreignIntrinsic)
+import PureScript.Backend.Wasm.Intrinsics (Intrinsic(MkEffectFn), qualifiedIntrinsic, foreignIntrinsic)
 import PureScript.Backend.Wasm.Lower.Collect (collectCtors, collectDictCtors, collectEnumCtors, collectFuncs, collectLabels, functionDecls, reachableFunctions)
 import PureScript.Backend.Wasm.Lower.Env (Env)
 import PureScript.Backend.Wasm.Lower.IR (Atom(..), AnfExpr(..), ForeignImport, FuncName(..), IRFunc, MarshalKind(..), Program, RecBind(..), Rep(..), Rhs(..), Slot(..), VarRef(..))
@@ -428,7 +428,7 @@ lowerCoreLetK env binds body finish = case Array.uncons binds of
       lowerCoreLetK (env { locals = Object.insert ident atom env.locals }) tail body finish
   Just { head: Rec recBinds, tail } -> case recBinds of
     [ r ]
-      | M.Abs params recBody <- r.expr
+      | M.Abs params recBody <- recBindFunctionForm r.expr
       , Just { head: param, tail: rest } <- Array.uncons params -> do
           { codeName, captures } <- liftLambda (Just r.ident) env param (reAbs rest recBody)
           bindRhs (RMkClosure codeName captures) \fAtom ->
@@ -451,32 +451,72 @@ lowerCoreLetK env binds body finish = case Array.uncons binds of
 -- | already bound to its slot, so sibling references become forward references for
 -- | the `LetRec` to patch.
 lowerRecBind :: Env -> Tuple M.RecBinding Slot -> Lower RecBind
-lowerRecBind env (Tuple rb slot) = case rb.expr of
+lowerRecBind env (Tuple rb slot) = case recBindFunctionForm rb.expr of
   M.Abs params recBody
     | Just { head: param, tail: rest } <- Array.uncons params -> do
         { codeName, captures } <- liftLambda Nothing env param (reAbs rest recBody)
         pure (RecBind slot codeName captures)
   -- A point-free recursive *function* (e.g. purescript-run's `loop = resume f pure`): not a
-  -- syntactic lambda, but a known callable applied below its arity, so it is in fact a function.
-  -- Eta-expand it to a saturated lambda (`\x -> e x`) — sound by the eta law, since it has positive
-  -- residual arity — and lower through the normal Abs path above. Genuine recursive *values*
-  -- (residual arity 0, e.g. a self-referential `Tuple`/`data`) do not match and fall through to the
-  -- error; a cyclic top-level value is instead handled by CAF globalization (ADR 0006, `FibAnd`).
-  _
-    | Just residual <- recBindResidualArity env rb.expr
-    , residual >= 1 ->
-        lowerRecBind env (Tuple (rb { expr = etaExpand rb.expr residual }) slot)
+  -- syntactic lambda, but a known callable whose (partial or saturated-returning-a-function)
+  -- application is in fact a function. Eta-expand it (`\x -> e x`, sound by the eta law) and lower
+  -- through the normal Abs path above. A saturated *constructor* application is genuine recursive
+  -- data, not a function, so `recBindEtaArity` returns `Nothing` and it falls through to the error
+  -- (a cyclic top-level value is instead handled by CAF globalization — ADR 0006, `FibAnd`).
+  peeled
+    | Just etaN <- recBindEtaArity env peeled ->
+        lowerRecBind env (Tuple (rb { expr = etaExpand peeled etaN }) slot)
   _ -> throw (UnsupportedExpr ("a recursive let binding must be a function: " <> rb.ident))
 
+-- | Normalise a recursive `let` binding's defining expression toward the syntactic lambda it
+-- | denotes, so the `LetRec` machinery recognises it as a function:
+-- |
+-- |   * peel `Data.Function.Uncurried.mkFnN` / `Effect.Uncurried.mkEffectFnN`, which are the
+-- |     identity (the uncurried value *is* the curried `$Clo`, ADR 0018; an *applied* `mkFnN`
+-- |     lowers to the `RPrim MkEffectFn` no-op) — e.g. `Data.Map.Internal`'s `mkFn2`-wrapped folds;
+-- |   * float a `let` whose body is a function inward — `let H in \xs -> b` becomes
+-- |     `\xs -> let H in b`. Sound for this strict, pure IR: `H` was bound outside the lambda so
+-- |     it cannot capture `xs`, and re-evaluating its (pure) bindings per call only forgoes
+-- |     sharing — e.g. the `let goLit … in \v -> case v of …` a `where`-helper-rich recursive
+-- |     worker compiles to.
+-- |
+-- | The middle-end normalises both, so only the `--no-opt` path reaches lowering with one intact.
+recBindFunctionForm :: M.Expr -> M.Expr
+recBindFunctionForm = case _ of
+  M.App (M.Var q) [ inner ]
+    | Just (Tuple MkEffectFn _) <- qualifiedIntrinsic (qualifiedKeyOf q) -> recBindFunctionForm inner
+  M.Let binds body
+    | M.Abs params inner <- recBindFunctionForm body -> M.Abs params (M.Let binds inner)
+  other -> other
+
 -- | The residual arity of a non-lambda binding RHS: a known callable (function / constructor /
 -- | intrinsic / foreign) applied to fewer arguments than its arity still denotes a function, of
 -- | the leftover arity. `Nothing` when the head's arity is unknown — then we cannot prove it is a
 -- | function, so the caller keeps the conservative "must be a function" error.
-recBindResidualArity :: Env -> M.Expr -> Maybe Int
-recBindResidualArity env = case _ of
-  v@(M.Var _) -> headArity env v
-  M.App h args -> (_ - Array.length args) <$> headArity env h
+-- | The number of parameters to eta-expand a non-lambda recursive binding by, turning it into a
+-- | syntactic function — or `Nothing` if the binding denotes a value rather than a function.
+-- |
+-- | A *partial* application is always a function: eta by its residual arity (`Cons 1`, `resume f`).
+-- | A *saturated* (or over-) application splits on the head. A constructor builds data, so a
+-- | saturated `Ctor …` is a genuine recursive *value* — a cyclic *local* data binding would diverge
+-- | under strict evaluation, so the only real case is a top-level CAF (handled by globalization,
+-- | ADR 0006); reject it here. A function's result is itself a value that, for a non-diverging local
+-- | recursive `let`, must be a function (e.g. `loop = resume f pure`, whose result type is a
+-- | function), so eta-expand by one to apply it — `lowerApp`'s over-application path does the rest.
+recBindEtaArity :: Env -> M.Expr -> Maybe Int
+recBindEtaArity env = case _ of
+  M.Var q -> etaArity q 0
+  M.App (M.Var q) args -> etaArity q (Array.length args)
   _ -> Nothing
+  where
+  etaArity q@(Qualified (Just _) ident) nargs
+    | Just info <- Object.lookup (qualifiedKeyOf q) env.ctors =
+        let residual = info.arity - nargs in if residual >= 1 then Just residual else Nothing
+    | Just arity <- funcArity q ident = Just (max 1 (arity - nargs))
+  etaArity _ _ = Nothing
+  funcArity q ident =
+    Object.lookup (qualifiedKeyOf q) env.knownFuncs
+      <|> (snd <$> (qualifiedIntrinsic (qualifiedKeyOf q) <|> foreignIntrinsic ident))
+      <|> ((Array.length <<< _.params) <$> Object.lookup (qualifiedKeyOf q) env.foreignSigs)
 
 -- | The declared arity of an application head, resolved the same way `lowerApp` dispatches a call:
 -- | a constructor, a top-level/specialized function, an intrinsic, or a foreign import.