purescript-lua · Unisay · Jul 24, 2026 · Jul 24, 2026 · Jul 24, 2026
diff --git a/bench/goldens/fnew_Bench.EffectStep.txt b/bench/goldens/fnew_Bench.EffectStep.txt
@@ -1,9 +1,9 @@
 chunk: Bench.EffectStep.lua
 runtime: LuaJIT 2.1.1741730670
 main-chunk FNEW: 11
-function-body FNEW: 12
-total FNEW: 23
-prototypes: 24
+function-body FNEW: 11
+total FNEW: 22
+prototypes: 23
 function-body FNEW sites:
   Bench.EffectStep.lua:4
   Bench.EffectStep.lua:6
@@ -13,7 +13,6 @@ function-body FNEW sites:
   Bench.EffectStep.lua:12
   Bench.EffectStep.lua:12
   Bench.EffectStep.lua:33
-  Bench.EffectStep.lua:46
-  Bench.EffectStep.lua:45
-  Bench.EffectStep.lua:49
-  Bench.EffectStep.lua:58
+  Bench.EffectStep.lua:44
+  Bench.EffectStep.lua:47
+  Bench.EffectStep.lua:56
diff --git a/bench/goldens/tnew_Bench.EffectStep.txt b/bench/goldens/tnew_Bench.EffectStep.txt
@@ -3,6 +3,6 @@ runtime: LuaJIT 2.1.1741730670
 main-chunk TNEW+TDUP: 4
 function-body TNEW+TDUP: 1
 total TNEW+TDUP: 5
-prototypes: 24
+prototypes: 23
 function-body TNEW+TDUP sites:
   Bench.EffectStep.lua:9 TDUP
diff --git a/bench/goldens/trace_effect_step.txt b/bench/goldens/trace_effect_step.txt
@@ -4,9 +4,8 @@ workload: n=100000 reps=4 result=1500000
 aborts (distinct site -- reason):
   Bench.EffectStep.lua:10 -- NYI: bytecode FNEW
   Bench.EffectStep.lua:12 -- NYI: bytecode FNEW
-  Bench.EffectStep.lua:45 -- NYI: bytecode FNEW
-  Bench.EffectStep.lua:46 -- NYI: bytecode FNEW
-  Bench.EffectStep.lua:58 -- NYI: bytecode FNEW
+  Bench.EffectStep.lua:44 -- NYI: bytecode FNEW
+  Bench.EffectStep.lua:56 -- NYI: bytecode FNEW
   Bench.EffectStep.lua:9 -- NYI: bytecode FNEW
 bytecode end state (J*=compiled, I*=blacklisted):
   Bench.EffectStep.lua:10 IFUNCF
@@ -17,9 +16,8 @@ bytecode end state (J*=compiled, I*=blacklisted):
   Bench.EffectStep.lua:17 IFUNCF
   Bench.EffectStep.lua:36 IFUNCF
   Bench.EffectStep.lua:37 IFUNCF
-  Bench.EffectStep.lua:39 IFUNCF
   Bench.EffectStep.lua:9 JFUNCF
   effect_step.lua:11 JFUNCF
   effect_step.lua:13 JFORI
   effect_step.lua:13 JFORL
-counts: aborts=6 compiled=6 blacklisted=7
+counts: aborts=5 compiled=6 blacklisted=6
diff --git a/changelog.d/20260724_140000_unisay_fold_call_through_scope_call.md b/changelog.d/20260724_140000_unisay_fold_call_through_scope_call.md
@@ -0,0 +1,28 @@
+### Added
+
+- A call applied to the result of an IIFE folds into its returns. The
+  "pick a thunk, then run it" codegen shape — a scope call immediately
+  applied, `(function() … end)()()` or `(function() … end)()(b)` — kept a
+  closure allocation and an extra call per invocation that the
+  tail-position collapse could not touch, because the applied result, not
+  the scope call itself, sat in tail position. The new
+  `foldCallThroughScopeCall` rule rewrites
+  `(function() …; return e end)()(args)` to
+  `(function() …; return e(args) end)()`, pushing the application into the
+  tail returns — through a branching tail (`if … then return f else return
+  g end`) as well, where the thunk selection lives. Evaluation order is
+  preserved (leading statements, then the returned expression, then the
+  arguments, in both forms), and `return e(args)` is a tail call, so even
+  the activation depth at the moment the result runs is unchanged. The
+  rule declines on early returns, fall-off paths (the original would call
+  `nil` and error), multi-valued returns, arguments whose free names a
+  body local would capture, varargs among the arguments, and non-atomic
+  arguments on a branching tail (they would be syntactically duplicated
+  per return site). The exposed plain scope call is then spliced away by
+  `collapseTailScopeCall`, which now re-applies itself to the merged
+  body — the fold builds tails at depths the bottom-up driver has already
+  passed. Thirteen goldens shrink: `Golden.UncurryEffect.Test`'s
+  `countdown` loses both tail closures, `Golden.TailRecM2Shadow.Test`'s
+  `untilE` predicate drops its per-iteration selection closure, and the
+  long-bind family runs its final effect directly. Eval goldens are
+  unchanged.
diff --git a/lib/Language/PureScript/Backend/Lua/Optimizer.hs b/lib/Language/PureScript/Backend/Lua/Optimizer.hs
@@ -2,10 +2,11 @@ module Language.PureScript.Backend.Lua.Optimizer where
 
 import Control.Monad.Trans.Accum (Accum, add, execAccum)
 import Data.Map qualified as Map
+import Data.Set qualified as Set
 import Language.PureScript.Backend.Lua.Fixture qualified as Fixture
 import Language.PureScript.Backend.Lua.Limits (LuaLimits, workingLocalCeiling)
 import Language.PureScript.Backend.Lua.Linker.Foreign (chunkScopeUsesVararg)
-import Language.PureScript.Backend.Lua.Localize (localizeChunk)
+import Language.PureScript.Backend.Lua.Localize (localizeChunk, namesInBlock)
 import Language.PureScript.Backend.Lua.Name qualified as Lua
 import Language.PureScript.Backend.Lua.Promote (promoteChunk)
 import Language.PureScript.Backend.Lua.Traversal
@@ -69,6 +70,7 @@ rewriteRulesInOrder ∷ LuaLimits → [RewriteRule]
 rewriteRulesInOrder limits =
   [ reduceTableDefinitionAccessor
   , foldFieldProjectionThroughScopeCall limits
+  , foldCallThroughScopeCall
   , collapseTailScopeCall limits
   , foldNotEqual
   ]
@@ -170,23 +172,152 @@ foldFieldProjectionThroughScopeCall limits original
         _ → Nothing
     _ → Nothing
 
-  containsReturn ∷ Annotated Comments StatementF → Bool
-  containsReturn (Ann statement) = case statement of
-    Return {} → True
+{- | Whether a statement contains a 'Return' at the level of the enclosing
+function body — one that exits the activation. A 'Return' inside a nested
+'Function' (or 'LocalFunction') belongs to a different activation and does
+not count, while a 'Return' inside a loop or 'Do' block at body level does.
+-}
+containsReturn ∷ Annotated Comments StatementF → Bool
+containsReturn (Ann statement) = case statement of
+  Return {} → True
+  IfThenElse _predicate thenBlock elseBlock →
+    any containsReturn thenBlock || any containsReturn elseBlock
+  Do body → any containsReturn body
+  While _predicate body → any containsReturn body
+  Repeat body _predicate → any containsReturn body
+  ForNum _name _start _limit _step body → any containsReturn body
+  ForIn _names _exprs body → any containsReturn body
+  LocalFunction {} → False
+  Assign {} → False
+  Local {} → False
+  CallStatement {} → False
+  Break → False
+
+{- | Rewrites @(function() …; return e end)()(args)@ to
+@(function() …; return e(args) end)()@: a call applied to the result of a
+no-argument, immediately-invoked function is folded into its @return@s.
+The shape is how the code generator runs a selected thunk — "pick an
+effect, then run it" lowers to a scope call immediately applied — and
+folding the application inward exposes a plain scope call that
+'collapseTailScopeCall' can then splice away. The freshly built call is
+folded once more by this same rule, in case the returned expression is
+itself an applied scope call.
+
+Applying before versus after the call returns is observably the same: the
+leading statements run first either way, the returned expression is
+evaluated before the arguments in both forms, and @return e(args)@ is a
+tail call, so even the activation depth at the moment the result runs is
+unchanged.
+
+The tail chain the fold covers is a single-valued 'Return', or an
+'IfThenElse'/'Do' whose every branch ends in one, recursively — the thunk
+selection is exactly a branching tail. The rule declines when:
+
+* a leading statement (of the body or of any branch on the tail chain)
+  'containsReturn': such an early return leaves the call on a path the
+  fold does not cover;
+
+* the tail chain has a fall-off path (a branch not ending in a 'Return',
+  an empty body): falling off yields @nil@, which the original code then
+  calls — an error the folded code would not reproduce;
+
+* the tail 'Return' is not single-valued: the call consumes the first
+  value only after Lua's adjustment, but the fold cannot drop the other
+  results without dropping their effects;
+
+* the arguments mention @...@ in their own scope: moved inside the
+  no-parameter callee, it would be rebound or fail to load;
+
+* a free name of the arguments collides with a local the callee's body
+  declares ('declaredNamesInActivation'): moved inside, the argument
+  would resolve the name to the callee's local instead of the enclosing
+  scope's binding;
+
+* the tail chain branches and any argument is not a name or a literal:
+  the arguments are duplicated into every return site — evaluated at most
+  once, since a single branch runs, but syntactically repeated — and
+  atoms keep that duplication trivial.
+-}
+foldCallThroughScopeCall ∷ RewriteRule
+foldCallThroughScopeCall = \case
+  original@( FunctionCall
+               (Ann (FunctionCall (Ann (Function [] body)) []))
+               args
+             )
+      | not (chunkScopeUsesVararg argsBlock)
+      , Set.disjoint (declaredNamesInActivation body) (namesInBlock argsBlock)
+      , Just body' ← pushCallIntoTail body →
+          FunctionCall (Lua.ann (Function [] body')) []
+      | otherwise → original
+     where
+      argsBlock = [Lua.ann (Return args)]
+
+      atomicArgs ∷ Bool
+      atomicArgs = all (isAtom . Lua.unAnn) args
+       where
+        isAtom ∷ Exp → Bool
+        isAtom = \case
+          Nil → True
+          Boolean _ → True
+          Integer _ → True
+          Float _ → True
+          String _ → True
+          Var (Ann (VarName _)) → True
+          _ → False
+
+      pushCallIntoTail
+        ∷ [Annotated Comments StatementF]
+        → Maybe [Annotated Comments StatementF]
+      pushCallIntoTail block = case reverse block of
+        lastStatement : reverseLeading
+          | not (any containsReturn reverseLeading) →
+              pushIntoStatement lastStatement <&> \pushed →
+                reverse reverseLeading <> [pushed]
+        _ → Nothing
+
+      pushIntoStatement
+        ∷ Annotated Comments StatementF
+        → Maybe (Annotated Comments StatementF)
+      pushIntoStatement (c, statement) =
+        (c,) <$> case statement of
+          Return [returnedValue] →
+            Just . Return . pure . Lua.ann $
+              foldCallThroughScopeCall (FunctionCall returnedValue args)
+          IfThenElse p thenBlock elseBlock
+            | atomicArgs →
+                IfThenElse p
+                  <$> pushCallIntoTail thenBlock
+                  <*> pushCallIntoTail elseBlock
+          Do doBody → Do <$> pushCallIntoTail doBody
+          _ → Nothing
+  e → e
+
+{- | Every name declared at the activation level of a block: block-level
+declarations at any depth count, while nested function literals are
+separate scopes whose declarations are invisible outside. The name-set
+counterpart of 'activationLocalSlots'.
+-}
+declaredNamesInActivation ∷ [Annotated Comments StatementF] → Set Lua.Name
+declaredNamesInActivation = foldMap (declared . Lua.unAnn)
+ where
+  declared ∷ StatementF Comments → Set Lua.Name
+  declared = \case
+    Local names _values → Set.fromList (toList names)
+    LocalFunction fname _params _body → Set.singleton fname
+    ForNum n _start _limit _step body →
+      Set.insert n (declaredNamesInActivation body)
+    ForIn names _exprs body →
+      Set.fromList (toList names) <> declaredNamesInActivation body
     IfThenElse _predicate thenBlock elseBlock →
-      any containsReturn thenBlock || any containsReturn elseBlock
-    Do body → any containsReturn body
-    While _predicate body → any containsReturn body
-    Repeat body _predicate → any containsReturn body
-    ForNum _name _start _limit _step body → any containsReturn body
-    ForIn _names _exprs body → any containsReturn body
-    -- A nested (local) function is a different activation: its returns do
-    -- not exit this call.
-    LocalFunction {} → False
-    Assign {} → False
-    Local {} → False
-    CallStatement {} → False
-    Break → False
+      declaredNamesInActivation thenBlock
+        <> declaredNamesInActivation elseBlock
+    Do body → declaredNamesInActivation body
+    While _predicate body → declaredNamesInActivation body
+    Repeat body _predicate → declaredNamesInActivation body
+    Assign {} → mempty
+    Return {} → mempty
+    CallStatement {} → mempty
+    Break → mempty
 
 {- | Rewrites @function(…) …; return (function() <stmts> end)() end@ to
 @function(…) …; <stmts> end@: a no-argument, immediately-invoked function
@@ -226,6 +357,11 @@ Conditions checked:
   'localizeChunk' budgets its cache locals against what is actually
   declared after the splice, and only gains upvalue headroom from the
   disappearing proto.
+
+The rule re-applies itself to the merged function: the new tail may again
+be a scope-call return — 'foldCallThroughScopeCall' builds such tails at
+depths the bottom-up driver has already passed — and every round
+re-checks the conditions above.
 -}
 collapseTailScopeCall ∷ LuaLimits → RewriteRule
 collapseTailScopeCall limits = \case
@@ -235,7 +371,11 @@ collapseTailScopeCall limits = \case
     , let merged = leading <> spliced
     , length params + activationLocalSlots merged
         <= workingLocalCeiling limits →
-        Function params merged
+        -- The merged tail may itself be a scope-call return — one
+        -- 'foldCallThroughScopeCall' built after the bottom-up driver had
+        -- already passed that depth — so keep collapsing while the budget
+        -- admits it; every round consumes one nesting level.
+        collapseTailScopeCall limits (Function params merged)
   e → e
  where
   -- Splits a function body whose last statement is a single-valued