Skip to content

Commit 41401fb

Browse files
UnisayUnisay
authored
Merge pull request #244 from purescript-lua/issue-231/complexity-capture-lattice
feat(optimizer): gate inlining with Complexity/Capture lattices
2 parents 72f88a5 + ed79493 commit 41401fb

25 files changed

Lines changed: 1367 additions & 1209 deletions

File tree

‎changelog.d/20260712_090000_unisay_complexity_capture_inlining.md‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
### Changed
2+
3+
- Inlining decisions now weigh expression cost and reference position, not
4+
only use counts and a flat size ceiling (#231). A `Complexity` lattice
5+
(`Trivial < Deref < KnownSize < NonTrivial`) prices duplicating an
6+
expression, and a `Capture` lattice (`CaptureNone < CaptureBranch <
7+
CaptureClosure`) locates a binding's references, admitting two tiers
8+
through the shared inlining guard: a small projection chain over
9+
write-once tables pastes at any use count, and a small closed lambda
10+
whose uses all sit outside branches and closures beta-reduces at every
11+
site. A non-trivial body is never duplicated into a branch or closure.

‎lib/Language/PureScript/Backend/IR/MagicDo.hs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,8 @@ nullary thunk (@function() … end@, run by calling it), we can recognise their
1313
1414
> function() local x = m1(); local _ = m2(); …; return last() end
1515
16-
which is flat regardless of length. This mirrors the magic-do pass of the
17-
upstream JS backend and @purs-backend-es@.
16+
which is flat regardless of length. This is the classic magic-do
17+
transformation of PureScript backends.
1818
1919
== Why a rewrite into existing 'Let'\/'Abs', not a new IR node
2020

‎lib/Language/PureScript/Backend/IR/Optimizer.hs‎

Lines changed: 164 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import Data.List qualified as List
77
import Data.List.NonEmpty qualified as NE
88
import Data.Map qualified as Map
99
import Data.Set qualified as Set
10+
import Data.Text qualified as Text
1011
import Language.PureScript.Backend.IR.DCE (eliminateDeadCode)
1112
import Language.PureScript.Backend.IR.FlattenDeepBinds (flattenDeepBindsM)
1213
import Language.PureScript.Backend.IR.FloatIn (floatIn)
@@ -35,15 +36,18 @@ import Language.PureScript.Backend.IR.Supply (SupplyM, freshName, runSupply)
3536
import Language.PureScript.Backend.IR.Types
3637
( AlgebraicType (SumType)
3738
, Ann
39+
, Capture (..)
3840
, Exp
3941
, Grouping (..)
4042
, Parameter (..)
4143
, PrimOp (..)
4244
, RawExp (..)
4345
, RewriteRuleM
46+
, Usage (..)
4447
, WasRewritten (..)
4548
, alphaEq
4649
, countFreeRef
50+
, countFreeRefUsage
4751
, countFreeRefs
4852
, ctorId
4953
, freshenBinders
@@ -465,6 +469,61 @@ inlineSizeBudget = 64
465469
expSize ∷ RawExp ann → Natural
466470
expSize e = 1 + sum (expSize <$> toListOf subexpressions e)
467471

472+
{- | The largest expression the Deref and KnownSize inlining tiers paste
473+
(Note [Complexity and Capture gate inlining]), sized in IR nodes like
474+
'inlineSizeBudget' but far below it: these tiers admit duplication at
475+
every use site, so growth scales with the use count.
476+
-}
477+
smallInlineBudget ∷ Natural
478+
smallInlineBudget = 16
479+
480+
{- | How costly an expression is to duplicate, ordered by escalation.
481+
Combines by taking the worse classification.
482+
See Note [Complexity and Capture gate inlining].
483+
-}
484+
data Complexity = Trivial | Deref | KnownSize | NonTrivial
485+
deriving stock (Show, Eq, Ord)
486+
487+
instance Semigroup Complexity where
488+
(<>) = max
489+
490+
instance Monoid Complexity where
491+
mempty = Trivial
492+
493+
{- | Bottom-up cost classification. 'Trivial': a reference or a
494+
scalar/empty literal.
495+
'Deref': a chain of cheap reads (projection, index, length, tag) over a
496+
Trivial base. 'KnownSize': an abstraction or a non-empty literal — a
497+
bounded allocation. Everything that computes is 'NonTrivial', and
498+
unlisted constructors deliberately land there, so a new node kind is
499+
conservative by default. A string literal above 128 characters counts
500+
as an allocation rather than a scalar: 'expSize' sees one node, but
501+
duplicating the payload is not free.
502+
-}
503+
complexityOf ∷ RawExp ann → Complexity
504+
complexityOf = \case
505+
Ref {} → Trivial
506+
LiteralInt {} → Trivial
507+
LiteralFloat {} → Trivial
508+
LiteralChar {} → Trivial
509+
LiteralBool {} → Trivial
510+
LiteralString _ann s
511+
| Text.length s > 128 → KnownSize
512+
| otherwise → Trivial
513+
LiteralArray _ann exprs
514+
| null exprs → Trivial
515+
| otherwise → KnownSize <> foldMap complexityOf exprs
516+
LiteralObject _ann props
517+
| null props → Trivial
518+
| otherwise → KnownSize <> foldMap (complexityOf . snd) props
519+
ObjectProp _ann base _prop → Deref <> complexityOf base
520+
ArrayIndex _ann base _idx → Deref <> complexityOf base
521+
ArrayLength _ann base → Deref <> complexityOf base
522+
ReflectCtor _ann base → Deref <> complexityOf base
523+
DataArgumentByIndex _ann _idx base → Deref <> complexityOf base
524+
AbsN _ann _params body → KnownSize <> complexityOf body
525+
_ → NonTrivial
526+
468527
{- | Pure wrapper for tests and standalone use: runs the rewrite with
469528
its own supply and no inlining environment. Production code uses
470529
'optimizedExpressionM' so all passes share one supply.
@@ -988,9 +1047,12 @@ inlineSaturatedCall env expr = case unwindApp expr of
9881047
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
9891048
'betaReduce' and 'inlineLocalBinding' decide whether to paste an expression
9901049
into its use sites by the same test: paste only when re-evaluating it cannot
991-
multiply work — the expression is trivial ('isInlinableExpr': a Ref, a
992-
literal, or an @inline-always node) or it is used at most once. Otherwise the
993-
expression stays behind a single 'Let' binding.
1050+
multiply work — the expression is cheap to re-evaluate ('isInlinableExpr': a
1051+
Ref, a literal, an @inline always@ node, or the Deref tier of
1052+
Note [Complexity and Capture gate inlining]), it is used at most once, or it
1053+
is a small closed abstraction whose uses all sit outside branches and
1054+
closures ('isDuplicatableClosedAbs', the KnownSize tier of the same Note).
1055+
Otherwise the expression stays behind a single 'Let' binding.
9941056
9951057
The two rules must agree, because they hand work to each other. When
9961058
'betaReduce' declines to substitute a redex it rewrites it to
@@ -1006,8 +1068,9 @@ fixpoint in one bottom-up pass instead of oscillating.
10061068
Fires only at exact arity: any other argument count on a literal lambda
10071069
head is ill-formed ('WellApplied', Note [n-ary abstraction]). The unary
10081070
redex is the singleton case. Each pair reduces by the shared guard: the
1009-
argument is substituted when trivial or used at most once, and bound by
1010-
a 'Let' otherwise; an argument at a 'ParamUnused' position is dropped
1071+
argument is substituted when cheap, used at most once, or a small
1072+
duplicatable closed abstraction, and bound by a 'Let' otherwise; an
1073+
argument at a 'ParamUnused' position is dropped
10111074
with its evaluation — the same call DCE makes when it drops an unused
10121075
binding unconditionally. See Note [IR is assumed well-typed].
10131076
-}
@@ -1044,7 +1107,10 @@ betaReduce = \case
10441107
ParamUnused _ann → pure (body, letBinds)
10451108
ParamNamed paramAnn name
10461109
-- See Note [Beta reduction and local inlining share an inlining guard]
1047-
| isInlinableExpr arg || countFreeRef (Local name) body <= 1 →
1110+
| usage ← countFreeRefUsage (Local name) body
1111+
, isInlinableExpr arg
1112+
|| usageTotal usage <= 1
1113+
|| isDuplicatableClosedAbs arg usage →
10481114
-- The λ is consumed by the rewrite, so the first inserted
10491115
-- occurrence of the argument may keep its binder names
10501116
-- ('substituteMoveM').
@@ -1229,19 +1295,72 @@ inlineLocalBinding rhsRefCounts grouping (body, inlined) =
12291295
-- names it (the count map records no zero entries, and the binding's
12301296
-- own RHS is excluded by the self-reference guard above).
12311297
isInlinableExpr inlinee
1232-
|| (occurrences == 1 && name `Map.notMember` rhsRefCounts) →
1298+
|| ( (occurrences == 1 || isDuplicatableClosedAbs inlinee usage)
1299+
&& name `Map.notMember` rhsRefCounts
1300+
) →
12331301
-- The binding survives until DCE drops it, so the inserted copy
12341302
-- must not reuse its binder names ('substituteCopyM').
12351303
(,Rewritten) <$> substituteCopyM name inlinee body
12361304
| otherwise → pure (body, inlined)
12371305
where
1306+
usage ∷ Usage
1307+
usage = countFreeRefUsage name body
12381308
occurrences ∷ Natural
1239-
occurrences = countFreeRef name body
1309+
occurrences = usageTotal usage
1310+
1311+
{- Note [Complexity and Capture gate inlining]
1312+
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1313+
Two small lattices refine the inlining heuristics beyond exact use
1314+
counts and flat size ceilings (issue #231):
1315+
1316+
Complexity = Trivial < Deref < KnownSize < NonTrivial ('complexityOf')
1317+
Capture = CaptureNone < CaptureBranch < CaptureClosure
1318+
('countFreeRefUsage')
1319+
1320+
'Complexity' prices duplicating an expression; 'Capture' locates a
1321+
binding's use sites relative to it. Together they admit two tiers past
1322+
the use-count rule:
1323+
1324+
* Deref tier ('isInlinableExpr'): an expression of complexity at most
1325+
'Deref' and size under 'smallInlineBudget' pastes at any use count.
1326+
Re-reading a projection chain is semantics-preserving because
1327+
PureScript records and module tables are write-once; the
1328+
re-evaluation cost is a field read. The tier reaches all three
1329+
guard sites — 'withBinding', 'betaReduce', 'inlineLocalBinding' —
1330+
through the shared predicate.
1331+
1332+
* KnownSize tier ('isDuplicatableClosedAbs', the two local sites
1333+
only): a closed 'AbsN' under 'smallInlineBudget' whose bound name's
1334+
uses all sit at 'CaptureNone' is substituted even when used many
1335+
times, so a locally-bound combinator beta-reduces at every site.
1336+
The 'CaptureNone' condition refuses to move the lambda literal into
1337+
a branch or closure: paying a closure allocation per call of the
1338+
surrounding function is the LuaJIT trace-abort pathology of issue
1339+
#204, so closedness and small size alone do not admit duplication.
1340+
1341+
A 'NonTrivial' body is admitted by neither tier, so it is never
1342+
duplicated — under a branch, a closure, or anywhere else. It still
1343+
inlines when used at most once: that is relocation, not duplication,
1344+
and the FloatIn pass moves work into single branches deliberately.
1345+
1346+
'withBinding' gets no KnownSize tier on purpose: an uber-module
1347+
binding's uses sit inside other top-level functions — 'CaptureClosure'
1348+
by construction — and saturated call sites of top-level lambdas are
1349+
already served by the call-site inliner ('inlineSaturatedCall').
1350+
1351+
Left for the follow-up analyses that need them: per-use-kind counters
1352+
(call/access/case) and an admission for a projection used only in call
1353+
position.
1354+
-}
12401355

12411356
-- See Note [Inline annotations and inlining heuristics]
1357+
-- and Note [Complexity and Capture gate inlining]
12421358
isInlinableExpr ∷ Exp → Bool
12431359
isInlinableExpr expr =
1244-
hasInlineAnnotation expr || isRef expr || isNonRecursiveLiteral expr
1360+
hasInlineAnnotation expr
1361+
|| isRef expr
1362+
|| isNonRecursiveLiteral expr
1363+
|| isCheapProjection expr
12451364
where
12461365
isRef ∷ RawExp a → Bool
12471366
isRef = \case
@@ -1254,3 +1373,39 @@ isInlinableExpr expr =
12541373
Just Always → True
12551374
Just Never → False
12561375
Nothing → False
1376+
1377+
-- The Deref tier. The explicit disjuncts above are subsumed for the
1378+
-- shapes they share (a Ref, a short scalar), but kept: they also admit
1379+
-- what the tier prices differently (a long string literal).
1380+
isCheapProjection ∷ Exp → Bool
1381+
isCheapProjection e =
1382+
complexityOf e <= Deref && expSize e < smallInlineBudget
1383+
1384+
{- | The KnownSize tier of Note [Complexity and Capture gate inlining]:
1385+
a small closed abstraction whose bound name is only used outside
1386+
branches and closures may be pasted at every use site.
1387+
-}
1388+
isDuplicatableClosedAbs ∷ Exp → Usage → Bool
1389+
isDuplicatableClosedAbs rhs Usage {usageCapture} =
1390+
isAbs rhs
1391+
&& usageCapture == CaptureNone
1392+
&& expSize rhs < smallInlineBudget
1393+
&& isClosedExp rhs
1394+
where
1395+
isAbs ∷ RawExp a → Bool
1396+
isAbs = \case
1397+
AbsN {} → True
1398+
_ → False
1399+
1400+
{- | No free local references. Imported references do not count: they
1401+
are valid at any position in the module, so they survive being
1402+
pasted anywhere.
1403+
-}
1404+
isClosedExp ∷ RawExp ann → Bool
1405+
isClosedExp =
1406+
not . any isLocalName . Map.keys . countFreeRefs
1407+
where
1408+
isLocalName ∷ Qualified Name → Bool
1409+
isLocalName = \case
1410+
Local _ → True
1411+
_ → False

‎lib/Language/PureScript/Backend/IR/Types.hs‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -777,6 +777,91 @@ countFreeRefs = fmap getSum . MMap.toMap . countFreeRefs' mempty
777777
countFreeRef ∷ Qualified Name → RawExp ann → Natural
778778
countFreeRef name = Map.findWithDefault 0 name . countFreeRefs
779779

780+
{- | Where a reference sits relative to the expression root: reached
781+
unconditionally, only inside an 'IfThenElse' arm, or from inside a
782+
nested 'AbsN'. Ordered by escalation; combines by taking the strongest
783+
context.
784+
-}
785+
data Capture = CaptureNone | CaptureBranch | CaptureClosure
786+
deriving stock (Show, Eq, Ord)
787+
788+
instance Semigroup Capture where
789+
(<>) = max
790+
791+
instance Monoid Capture where
792+
mempty = CaptureNone
793+
794+
{- | Aggregate of one name's free references: how many, and the
795+
strongest 'Capture' context any of them sits under.
796+
-}
797+
data Usage = Usage {usageTotal ∷ Natural, usageCapture ∷ Capture}
798+
deriving stock (Show, Eq)
799+
800+
instance Semigroup Usage where
801+
Usage t1 c1 <> Usage t2 c2 = Usage (t1 + t2) (c1 <> c2)
802+
803+
instance Monoid Usage where
804+
mempty = Usage 0 mempty
805+
806+
{- | 'countFreeRef' enriched with the 'Capture' context of the counted
807+
references: the traversal threads the strongest wrapper crossed between
808+
the expression root and each reference site — an 'AbsN' body raises the
809+
context to 'CaptureClosure', an 'IfThenElse' arm (not the condition) to
810+
'CaptureBranch'. A 'Let' defers nothing: its RHSs and body evaluate
811+
when the 'Let' does, so the context passes through unchanged and only
812+
the bound-name set advances (Note [Sequential scoping of Let bindings]).
813+
-}
814+
countFreeRefUsage ∷ Qualified Name → RawExp ann → Usage
815+
countFreeRefUsage name = go CaptureNone mempty
816+
where
817+
go ∷ Capture → Set Name → RawExp ann → Usage
818+
go cap bound = \case
819+
Ref _ann qname
820+
| qname == name → case qname of
821+
Local n | Set.member n bound → mempty
822+
_ → Usage 1 cap
823+
| otherwise → mempty
824+
AbsN _ann params body →
825+
go (cap <> CaptureClosure) (foldl' bindParam bound params) body
826+
where
827+
bindParam ∷ Set Name → Parameter ann → Set Name
828+
bindParam names = \case
829+
ParamNamed _paramAnn n → Set.insert n names
830+
ParamUnused _paramAnn → names
831+
IfThenElse _ann cond thenBranch elseBranch →
832+
go cap bound cond
833+
<> go (cap <> CaptureBranch) bound thenBranch
834+
<> go (cap <> CaptureBranch) bound elseBranch
835+
-- See Note [Sequential scoping of Let bindings]
836+
Let _ann binds body → fold (usageInBody : usagesInBinds)
837+
where
838+
usageInBody = go cap boundAfterBinds body
839+
(boundAfterBinds, usagesInBinds) =
840+
foldl' withGrouping (bound, []) (toList binds)
841+
withGrouping
842+
∷ (Set Name, [Usage])
843+
→ Grouping (ann, Name, RawExp ann)
844+
→ (Set Name, [Usage])
845+
withGrouping (names, usages) = \case
846+
Standalone (_nameAnn, boundName, expr) →
847+
(Set.insert boundName names, go cap names expr : usages)
848+
RecursiveGroup recBinds →
849+
( namesAfterGroup
850+
, ( toList recBinds <&> \(_nameAnn, _boundName, expr) →
851+
go cap namesAfterGroup expr
852+
)
853+
<> usages
854+
)
855+
where
856+
namesAfterGroup =
857+
foldr
858+
(\(_nameAnn, boundName, _expr) → Set.insert boundName)
859+
names
860+
recBinds
861+
-- No other constructor binds names or defers/conditions evaluation,
862+
-- so both context components pass through:
863+
other → foldMapOf subexpressions (go cap bound) other
864+
780865
{- | Structural equality modulo the names of locally-bound binders.
781866
782867
Two expressions are alpha-equivalent when they differ at most in the

0 commit comments

Comments
 (0)