[Feature] Reduction-aware inlining

compiler/src/PureScript/Backend/Wasm/MiddleEnd/Optimize/DictElim.purs

@@ -18,6 +18,7 @@ module PureScript.Backend.Wasm.MiddleEnd.Optimize.DictElim
   ( buildCtx
   , simplifyModule
   , summarize
+  , normalFormSizeCap
   ) where
 
 import Prelude
@@ -132,8 +133,30 @@ summarize keepKeys m = m { decls = Array.filter keep m.decls }
 useNbE :: Boolean
 useNbE = true
 
+-- | A reduced declaration larger than this many IR nodes is re-reduced with the inline
+-- | context emptied (ADR 0035, the Layer C size guard). Inlining exists to *shrink* code by
+-- | firing redexes; when a binding instead inlines into a normal form orders of magnitude
+-- | larger than any real declaration — the canonical case is the `genericShow` dictionary of a
+-- | large derived-`Generic` ADT inlined into a `show`, which produces no redex, only bulk — it is
+-- | pure code-size blow-up. Falling back to the un-inlined form keeps that dictionary an ordinary
+-- | call, the same (correct) shape `--no-opt` emits, and is what bounds NbE when a program is
+-- | itself `show`-heavy (notably the compiler compiling itself). The threshold sits far above any
+-- | genuine declaration (tens of thousands of nodes) and far below the observed blow-ups
+-- | (5×10⁵–2×10⁶), so only pathological declarations fall back; the guard measures the *actual*
+-- | reduced size, so a large-but-shared normal form (which quote CSEs back down) is kept inlined.
+normalFormSizeCap :: Int
+normalFormSizeCap = 200_000
+
 reduce :: Ctx -> M.Expr -> M.Expr
-reduce ctx = if useNbE then normalize ctx else simplifyExpr ctx
+reduce ctx e =
+  let
+    r = reduce1 ctx e
+  in
+    if exprSize r > normalFormSizeCap && ctxInlines ctx then reduce1 (ctx { inline = Map.empty, instanceFields = Map.empty }) e
+    else r
+  where
+  reduce1 c = if useNbE then normalize c else simplifyExpr c
+  ctxInlines c = not (Map.isEmpty c.inline) || not (Map.isEmpty c.instanceFields)
 
 simplifyModule :: Ctx -> M.Module -> M.Module
 simplifyModule ctx m = m { decls = map go m.decls }

compiler/src/PureScript/Backend/Wasm/MiddleEnd/Optimize/Specialize.purs

@@ -36,7 +36,6 @@ import Prelude
 import Control.Monad.State (State, gets, modify_, runState)
 import Data.Array as Array
 import Data.Either (Either(..))
-import Data.Foldable (foldl)
 import Data.Map (Map)
 import Data.Map as Map
 import Data.Maybe (Maybe(..), fromMaybe)
@@ -48,6 +47,7 @@ import Data.Tuple (Tuple(..), snd)
 import PureScript.Backend.Wasm.MiddleEnd.FreeVars (binderVars, freeVars)
 import PureScript.Backend.Wasm.MiddleEnd.IR as M
 import PureScript.Backend.Wasm.MiddleEnd.Optimize.Simplify (substMany)
+import PureScript.Backend.Wasm.MiddleEnd.Serialize.Hash (hashString)
 import PureScript.CoreFn (Literal(..), ModuleName, Qualified(..))
 
 -- | A top-level function that is a candidate callee: its parameter list, body, and
@@ -327,42 +327,25 @@ lambdaFrees = case _ of
   e -> freeVars [] e
 
 -- a structural key for the lambda with its free variables abstracted, so two
--- lambdas equal up to their captures share a specialization
+-- lambdas equal up to their captures share a specialization. The free vars are
+-- abstracted to positional markers in a single capture-avoiding pass (rather than one
+-- `substVar` per free), and the canonical form is then **hashed** rather than used
+-- verbatim: the raw `show` of a large lambda body is a multi-kilobyte string, and it was
+-- both built (genericShow) and *retained as a `Map` key* for every specialization
+-- candidate — quadratic-ish key comparisons plus heavy GC that dominated the optimization of
+-- large higher-order modules. A 16-hex digest keeps the dedup `Map` keys tiny and fixed-size.
+-- A clash would mis-share two distinct specializations, but the digest is the same one already
+-- trusted for the build-cache identity, so an accidental clash is astronomically unlikely.
 canonicalKey :: Array String -> M.Expr -> String
 canonicalKey frees lam =
-  show (foldlWithIndexArr (\i e f -> substVar f (M.Var (Qualified Nothing ("#" <> show i))) e) lam frees)
+  hashString (show (substMany markers lam))
+  where
+  markers =
+    Map.fromFoldable
+      (Array.mapWithIndex (\i f -> Tuple f (M.Var (Qualified Nothing ("#" <> show i)))) frees)
 
 -- substitution / helpers ------------------------------------------------------
 
--- replace free occurrences of local `name` with `repl`, stopping at shadowing
-substVar :: String -> M.Expr -> M.Expr -> M.Expr
-substVar name repl = go
-  where
-  go = case _ of
-    e@(M.Var (Qualified Nothing n)) -> if n == name then repl else e
-    e@(M.Var _) -> e
-    M.Lit lit -> M.Lit (mapLit go lit)
-    e@(M.Constructor _ _ _) -> e
-    M.Accessor l e -> M.Accessor l (go e)
-    M.Update e cf kvs -> M.Update (go e) cf (map (map go) kvs)
-    M.Abs ps b -> if Array.elem name ps then M.Abs ps b else M.Abs ps (go b)
-    M.App f a -> mkApp (go f) (map go a)
-    M.Perform e -> M.Perform (go e)
-    M.Case ss alts -> M.Case (map go ss) (map goAlt alts)
-    M.Let bs body ->
-      if Array.elem name (bs >>= boundNames) then M.Let bs body
-      else M.Let (map goBind bs) (go body)
-  goAlt alt =
-    if Array.elem name (alt.binders >>= binderVars) then alt
-    else alt
-      { result = case alt.result of
-          Right e -> Right (go e)
-          Left gs -> Left (map (\g -> { guard: go g.guard, expression: go g.expression }) gs)
-      }
-  goBind = case _ of
-    M.NonRec meta i e -> M.NonRec meta i (go e)
-    M.Rec rs -> M.Rec (map (\r -> r { expr = go r.expr }) rs)
-
 -- keep `App` heads flat (never an `App` of an `App`)
 mkApp :: M.Expr -> Array M.Expr -> M.Expr
 mkApp head args
@@ -404,9 +387,6 @@ litExprs = case _ of
   LitObject kvs -> map snd kvs
   _ -> []
 
-foldlWithIndexArr :: forall a b. (Int -> b -> a -> b) -> b -> Array a -> b
-foldlWithIndexArr f z arr = foldl (\acc (Tuple i a) -> f i acc a) z (Array.mapWithIndex Tuple arr)
-
 key :: ModuleName -> String -> String
 key modName ident = joinWith "." modName <> "." <> ident
 

compiler/test/NbeStress.purs

@@ -0,0 +1,118 @@
+-- | The NbE exponential regression guard (bug A, ADR 0035 §8).
+-- |
+-- | It drives `Semantics.normalize` directly against a synthetic **diamond inline DAG**: an
+-- | inline set `b0 … bd` where each `bᵢ = f b₍ᵢ₋₁₎ b₍ᵢ₋₁₎` references the previous binding twice.
+-- | Normalizing `bd` re-evaluates the shared leaf `b0` once per path = Θ(2ᵈ) on the *unfixed*
+-- | reducer (M1 in eval, M2 in quote), so the result size **doubles every depth** and the wall
+-- | clock explodes around d≈22-25. After the ADR-0035 sharing fix (Layer A eval memo + Layer B
+-- | quote CSE) the result stays **O(d)** (≈ 4d) and the loop runs instantly — verified without
+-- | building the whole compiler to wasm.
+-- |
+-- | `spec` is the routine `test:unit` guard: the linear-bound diamond check above, plus a guard for
+-- | the **Layer C size cap** (`DictElim.simplifyModule` falls back to the un-inlined form when a
+-- | declaration inlines genuine, unshareable bulk — the `genericShow`-into-`show` blow-up that hung
+-- | self-compilation). `main` sweeps deeper diamond depths for manual inspection:
+-- |
+-- |     spago test -p compiler -m Test.NbeStress
+module Test.NbeStress where
+
+import Prelude
+
+import Data.Array as Array
+import Data.Foldable (for_)
+import Data.Map (Map)
+import Data.Map as Map
+import Data.Maybe (Maybe(..))
+import Data.Set as Set
+import Data.Tuple (Tuple(..))
+import Effect (Effect)
+import Effect.Console (error) as Console
+import PureScript.Backend.Wasm.MiddleEnd.IR as M
+import PureScript.Backend.Wasm.MiddleEnd.Optimize.Analysis (exprSize)
+import PureScript.Backend.Wasm.MiddleEnd.Optimize.DictElim (normalFormSizeCap, simplifyModule)
+import PureScript.Backend.Wasm.MiddleEnd.Optimize.Semantics (normalize)
+import PureScript.CoreFn (Literal(..), Qualified(..))
+import Test.Spec (Spec, describe, it)
+import Test.Spec.Assertions (shouldEqual)
+
+maxDepth :: Int
+maxDepth = 20
+
+-- A reference to inline binding `bᵢ` (module `M`), keyed so `qkey` = `"M.bᵢ"`.
+ref :: Int -> M.Expr
+ref i = M.Var (Qualified (Just [ "M" ]) ("b" <> show i))
+
+-- An opaque (non-inline) head, so the application stays a neutral whose two operands are each
+-- re-evaluated — the binary fan-out that makes the DAG a diamond rather than a chain.
+opaque :: M.Expr
+opaque = M.Var (Qualified (Just [ "Ext" ]) "f")
+
+-- The inline set for depth `d`: `b0` is a neutral leaf; `bᵢ = f b₍ᵢ₋₁₎ b₍ᵢ₋₁₎`.
+diamondInline :: Int -> Map String M.Expr
+diamondInline d =
+  Map.fromFoldable
+    ( Array.cons (Tuple "M.b0" (M.Var (Qualified (Just [ "Ext" ]) "leaf")))
+        (map (\i -> Tuple ("M.b" <> show i) (M.App opaque [ ref (i - 1), ref (i - 1) ])) (Array.range 1 d))
+    )
+
+ctxFor :: Int -> { newtypeCtors :: Set.Set String, dataCtors :: Set.Set String, inline :: Map String M.Expr, instanceFields :: Map String (Array (Tuple String M.Expr)), effectfulForeigns :: Set.Set String, impureBindings :: Set.Set String, memEffBindings :: Set.Set String }
+ctxFor d =
+  { newtypeCtors: Set.empty
+  , dataCtors: Set.empty
+  , inline: diamondInline d
+  , instanceFields: Map.empty
+  , effectfulForeigns: Set.empty
+  , impureBindings: Set.empty
+  , memEffBindings: Set.empty
+  }
+
+-- An **un-shareable** term of `n` distinct-position leaves (a flat array of int literals): unlike a
+-- diamond, `quote` cannot CSE it, so its normal form really is ~`n` nodes. Inlining a binding bound
+-- to this produces genuine bulk with no reduction — the `genericShow`-into-`show` pathology in
+-- miniature — which is exactly what the size cap must catch.
+flatBig :: Int -> M.Expr
+flatBig n = M.Lit (LitArray (Array.replicate n (M.Lit (LitInt 0))))
+
+-- a one-declaration module `U.user = <body>`, the unit `simplifyModule` reduces.
+userModule :: M.Expr -> M.Module
+userModule body = { name: [ "U" ], decls: [ M.NonRec Nothing "user" body ] }
+
+-- the reduced size of `U.user` after `simplifyModule`.
+userSize :: M.Module -> Int
+userSize m = case Array.head m.decls of
+  Just (M.NonRec _ _ e) -> exprSize e
+  _ -> -1
+
+spec :: Spec Unit
+spec = do
+  -- | Routine guard (ADR 0035 §8): a depth-20 diamond normalizes to a *linear*-size term. On the
+  -- | pre-sharing reducer this is 2^20 ≈ a million nodes (seconds of work + a huge tree); the
+  -- | Layer A memo + Layer B quote CSE keep it ≈ 4·20 + 3 = 83. The generous `< 1000` bound passes
+  -- | instantly when sharing holds and fails (or times out) the moment the exponential returns.
+  describe "Semantics.normalize — NbE exponential guard (ADR 0035)" do
+    it "normalizes a depth-20 diamond inline DAG to linear size, not O(2^d)" do
+      let d = 20
+      (exprSize (normalize (ctxFor d) (ref d)) < 1000) `shouldEqual` true
+
+  -- | The Layer C size cap: inlining that blows the normal form past `normalFormSizeCap` falls back
+  -- | to the un-inlined form (the binding stays a call). This is what bounds NbE when a declaration
+  -- | inlines genuine, unshareable bulk — the `genericShow` dictionary of a large derived-`Generic`
+  -- | ADT inlined into `show`, the case that hung the compiler compiling itself. The companion case
+  -- | proves the cap discriminates: a reduced form *under* the cap is still inlined.
+  describe "DictElim.simplifyModule — code-size cap (ADR 0035 Layer C lite)" do
+    it "falls back to the un-inlined call when inlining blows the size cap" do
+      let big = M.Var (Qualified (Just [ "M" ]) "big")
+      let ctx = (ctxFor 0) { inline = Map.singleton "M.big" (flatBig (normalFormSizeCap + 10)) }
+      -- inlined it would be > cap; the cap forces the un-inlined form, so `user` stays a small call.
+      (userSize (simplifyModule ctx (userModule big)) < normalFormSizeCap) `shouldEqual` true
+    it "still inlines a binding whose reduced form is under the cap (the cap discriminates)" do
+      let small = M.Var (Qualified (Just [ "M" ]) "small")
+      let ctx = (ctxFor 0) { inline = Map.singleton "M.small" (flatBig 50) }
+      -- well under the cap: inlining proceeds, so `user` grows to the inlined array (≫ a bare call).
+      (userSize (simplifyModule ctx (userModule small)) > 50) `shouldEqual` true
+
+main :: Effect Unit
+main = do
+  Console.error "NbE diamond stress — result size should be O(d); on the unfixed reducer it doubles per depth:"
+  for_ (Array.range 1 maxDepth) \d ->
+    Console.error ("  d=" <> show d <> "  normalized-size=" <> show (exprSize (normalize (ctxFor d) (ref d))))

compiler/test/Unit/Compiler.purs

@@ -29,6 +29,7 @@ import Test.Unit.PureScript.Backend.Wasm.Ulib.Interface as UlibInterface
 import Test.Unit.PureScript.Backend.Wasm.MiddleEnd.Transl as Transl
 import Test.Unit.PureScript.CoreFn as CoreFn
 import Test.Unit.PureScript.ExternsFile as ExternsFile
+import Test.NbeStress as NbeStress
 
 main :: Effect Unit
 main = runSpecAndExitProcess [ consoleReporter ] do
@@ -37,6 +38,7 @@ main = runSpecAndExitProcess [ consoleReporter ] do
   Externs.spec
   Caf.spec
   Codegen.spec
+  NbeStress.spec
   Lower.spec
   Match.spec
   Transl.spec

docs/design-decisions/0020-reduction-aware-inliner.md

@@ -22,6 +22,15 @@
 > (it recomputes values on every `eval`/`quote` traversal), independent of the fusion motivation
 > here. ADR 0035 sequences the fix: a behaviour-neutral **sharing/memo** pass (the scalability gate)
 > first, then this ADR's reduction-aware inline-or-share **policy** on top.
+>
+> **Update (2026-06-17):** [ADR 0035](0035-sharing-nbe-reduction-aware-inlining.md)'s **sharing pass
+> (Layers A + B) landed**, removing the NbE recomputation exponential. ~~and opened the scalability gate — the optimized self-compile now reduces
+> through `Optimize.Specialize` instead of spinning.~~ _(A+B alone did **not** clear `Optimize.Specialize`;
+> completing the optimized self-compile also required a Layer-C-lite `normalFormSizeCap` code-size cap
+> and the `Optimize.Specialize` dedup-key fix — see [ADR 0035](0035-sharing-nbe-reduction-aware-inlining.md).)_
+> **This ADR's reduction-aware *decision* is
+> ADR 0035 Layer C, now deferred**: with the exponential gone, it is an optimization-quality
+> improvement (the fusion win) rather than a scalability blocker.
 
 ## Context