Skip to content

Count free references incrementally in the IR optimizer #189

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Unisay merged 1 commit into from
Jul 6, 2026
Merged

Count free references incrementally in the IR optimizer #189

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
101 changes: 79 additions & 22 deletions lib/Language/PureScript/Backend/IR/Optimizer.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,6 @@ import Language.PureScript.Backend.IR.Types
, RewriteRuleM
, WasRewritten (..)
, alphaEq
, bindingExprs
, countFreeRef
, countFreeRefs
, getAnn
Expand Down Expand Up @@ -184,10 +183,51 @@ neverInlineNames UberModule {uberModuleBindings} =
, getAnn expr == Just Never
]

-- | Free-reference counts keyed by the referenced qualified name.
type FreeRefs = Map (Qualified Name) Natural

{- Note [Incremental free-reference counting]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
'withBinding' decides whether to inline a 'Standalone' binding partly on
whether it is referenced exactly once. Counting those references naively —
folding 'countFreeRefs' over the whole module for each binding — costs
O(bindings × module size) per 'optimizeModule' run, and the run itself
repeats inside the optimize+dce fixpoint (issue #142).

Instead we thread one 'FreeRefs' map through the fold, holding the exact
invariant:

counts = free references over the current accumulator
= (bindingExprs =<< bindings) <> map snd exports

'foldrM' visits bindings right-to-left, and a binding may only be referenced
by material to its right (Note [Sequential scoping of Let bindings]), so by
the time a binding is visited every one of its referencers is already in the
accumulator and its count in 'counts' is final. The map is built once from
the exports and updated by a known delta at each step:

* binding kept → its own expression joins the accumulator, so add its
free refs: @counts ⊕ countFreeRefs expr@.
* binding inlined → each of its @k@ occurrences is replaced by a copy of
the expression, so add @k@ copies of the expression's free refs and drop
the binding's own entry: @delete qn (counts ⊕ k · countFreeRefs expr)@.
Under the pass's unique-binders invariant a 'Standalone' RHS never names
its own binder, so @countFreeRefs expr@ cannot reintroduce @qn@.
* recursive group → its members are never inlined; their (mutually
recursive) refs join the accumulator like any kept binding.

Besides removing the quadratic cost, this counts against the live
(post-substitution) accumulator rather than a stale snapshot of the original
bindings, so the use-once decision no longer misjudges a binding whose
references changed earlier in the same run (issue #143).
-}
optimizeModule ∷ Set QName → UberModule → SupplyM (UberModule, WasRewritten)
optimizeModule neverNames UberModule {..} = runWriterT do
(bindings, exports) ←
foldrM withBinding ([], uberModuleExports) uberModuleBindings
-- See Note [Incremental free-reference counting]
let initialCounts =
Map.unionsWith (+) (countFreeRefs . snd <$> uberModuleExports)
(_finalCounts, bindings, exports) ←
foldrM withBinding (initialCounts, [], uberModuleExports) uberModuleBindings
uberModuleBindings' ←
traverse (traverse (traverse optimizeExp)) bindings
uberModuleExports' ← traverse (traverse optimizeExp) exports
Expand All @@ -206,39 +246,56 @@ optimizeModule neverNames UberModule {..} = runWriterT do
(e', rewritten) ← lift (optimizedExpressionM e)
e' <$ tell rewritten

-- See Note [Incremental free-reference counting]
withBinding
∷ Grouping (QName, Exp)
→ ([Grouping (QName, Exp)], [(Name, Exp)])
→ WriterT WasRewritten SupplyM ([Grouping (QName, Exp)], [(Name, Exp)])
withBinding binding (bindings, exports) =
→ (FreeRefs, [Grouping (QName, Exp)], [(Name, Exp)])
→ WriterT
WasRewritten
SupplyM
(FreeRefs, [Grouping (QName, Exp)], [(Name, Exp)])
withBinding binding (counts, bindings, exports) =
case binding of
Standalone (qname, expr0) → do
expr ← optimizeExp expr0
-- See Note [Inline annotations and inlining heuristics]
let isUsedOnce name =
1 == Map.findWithDefault 0 (qualifiedQName name) uberModuleFreeRefs
uberModuleFreeRefs ∷ Map (Qualified Name) Natural =
foldr
(\e m → Map.unionWith (+) m (countFreeRefs e))
mempty
uberModuleExprs
uberModuleExprs =
(bindingExprs =<< uberModuleBindings) <> map snd exports
let qn = qualifiedQName qname
occurrences = Map.findWithDefault 0 qn counts
isUsedOnce = occurrences == 1
if qname `Set.notMember` neverNames
&& (isInlinableExpr expr || isUsedOnce qname)
&& (isInlinableExpr expr || isUsedOnce)
then do
-- The binding is dropped from the module in favor of the
-- substituted copies: a rewrite even when it had no
-- occurrences left to substitute.
tell Rewritten
lift $
(,)
<$> substituteInBindings qname expr bindings
<*> substituteInExports qname expr exports
else pure (Standalone (qname, expr) : bindings, exports)
(bindings', exports') ←
lift $
(,)
<$> substituteInBindings qname expr bindings
<*> substituteInExports qname expr exports
-- Substituting drops the binding's own occurrences and pastes
-- one copy of its free refs per occurrence replaced.
let pastedRefs = fmap (* occurrences) (countFreeRefs expr)
counts' = Map.delete qn (Map.unionWith (+) counts pastedRefs)
pure (counts', bindings', exports')
else
-- The binding survives, so its refs now live in the accumulator.
pure
( Map.unionWith (+) counts (countFreeRefs expr)
, Standalone (qname, expr) : bindings
, exports
)
RecursiveGroup recGroup → do
recGroup' ← traverse (traverse optimizeExp) recGroup
pure (RecursiveGroup recGroup' : bindings, exports)
-- Recursive-group members are never inlined; their (mutually
-- recursive) refs now live in the accumulator.
let counts' =
foldl'
(\m (_, e) → Map.unionWith (+) m (countFreeRefs e))
counts
(toList recGroup')
pure (counts', RecursiveGroup recGroup' : bindings, exports)

-- Cross-entry substitution: the host entry's binders give no uniqueness
-- guarantee against the inlinee's, so every copy is freshened
Expand Down
38 changes: 37 additions & 1 deletion test/Language/PureScript/Backend/IR/Optimizer/Spec.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,11 @@ import Language.PureScript.Backend.IR.Names
, moduleNameFromString
)
import Language.PureScript.Backend.IR.Optimizer
( optimizedExpression
( optimizeModule
, optimizedExpression
, optimizedUberModule
)
import Language.PureScript.Backend.IR.Supply (runSupply)
import Language.PureScript.Backend.IR.Types
( Exp
, Grouping (..)
Expand Down Expand Up @@ -234,6 +236,40 @@ spec = describe "IR Optimizer" do
annotateShow optimized
fooKept === [QName mainModule (Name "foo")]

describe "counts free references against the live module (#143)" do
-- Within a single 'optimizeModule' run the use-once check must consult
-- the current (post-substitution) view of the module, not a stale
-- snapshot of the original bindings. Here `y` collapses to a bare
-- reference to `x` and is inlined into the export, leaving `x`
-- referenced exactly once; counting against the live view inlines `x`
-- too, whereas counting the two pre-collapse occurrences in the
-- original `y` wrongly keeps it. A single run is deliberate: the
-- optimize+dce fixpoint masks the misjudgment by self-correcting on a
-- later iteration (issue #143).
it "inlines a binding that becomes used-once mid-run" do
let main' = moduleNameFromString "Main"
extern = moduleNameFromString "Extern"
-- Non-inlinable and never rewritten by 'optimizeExp'.
xExpr = application (refImported extern (Name "f")) (literalInt 1)
yExpr =
ifThenElse
(literalBool True)
(refImported main' (Name "x"))
(refImported main' (Name "x"))
original =
Linker.UberModule
{ uberModuleForeigns = []
, uberModuleBindings =
[ Standalone (QName main' (Name "x"), xExpr)
, Standalone (QName main' (Name "y"), yExpr)
]
, uberModuleExports =
[(Name "main", refImported main' (Name "y"))]
}
optimized = fst (runSupply (optimizeModule mempty original))
Linker.uberModuleBindings optimized `shouldBe` []
Linker.uberModuleExports optimized `shouldBe` [(Name "main", xExpr)]

describe "inliner unlocks more optimizations" do
test "constant folding after inlining" do
name ← forAll Gen.name
Expand Down
58 changes: 29 additions & 29 deletions test/ps/output/Golden.ArrayOfUnits.Test/golden.ir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,7 @@ UberModule
( PropName "foldMap", Abs Nothing
( ParamNamed Nothing ( Name "dictMonoid" ) )
( Abs Nothing
( ParamNamed Nothing ( Name "f$1824" ) )
( ParamNamed Nothing ( Name "f$885" ) )
( App Nothing
( App Nothing
( App Nothing
Expand All @@ -88,9 +88,9 @@ UberModule
)
)
( Abs Nothing
( ParamNamed Nothing ( Name "x$1825" ) )
( ParamNamed Nothing ( Name "x$886" ) )
( Abs Nothing
( ParamNamed Nothing ( Name "acc$1826" ) )
( ParamNamed Nothing ( Name "acc$887" ) )
( App Nothing
( App Nothing
( ObjectProp Nothing
Expand All @@ -106,11 +106,11 @@ UberModule
( PropName "append" )
)
( App Nothing
( Ref Nothing ( Local ( Name "f$1824" ) ) )
( Ref Nothing ( Local ( Name "x$1825" ) ) )
( Ref Nothing ( Local ( Name "f$885" ) ) )
( Ref Nothing ( Local ( Name "x$886" ) ) )
)
)
( Ref Nothing ( Local ( Name "acc$1826" ) ) )
( Ref Nothing ( Local ( Name "acc$887" ) ) )
)
)
)
Expand Down Expand Up @@ -182,9 +182,9 @@ UberModule
( LiteralObject Nothing
[
( PropName "map", Abs Nothing
( ParamNamed Nothing ( Name "f$1109" ) )
( ParamNamed Nothing ( Name "f$707" ) )
( Abs Nothing
( ParamNamed Nothing ( Name "a$1110" ) )
( ParamNamed Nothing ( Name "a$708" ) )
( App Nothing
( App Nothing
( App Nothing
Expand Down Expand Up @@ -212,10 +212,10 @@ UberModule
( Imported ( ModuleName "Effect" ) ( Name "applicativeEffect" ) )
)
)
( Ref Nothing ( Local ( Name "f$1109" ) ) )
( Ref Nothing ( Local ( Name "f$707" ) ) )
)
)
( Ref Nothing ( Local ( Name "a$1110" ) ) )
( Ref Nothing ( Local ( Name "a$708" ) ) )
)
)
)
Expand All @@ -235,7 +235,7 @@ UberModule
[
( PropName "apply", Let Nothing
( Standalone
( Nothing, Name "bind$1090", App Nothing
( Nothing, Name "bind$687", App Nothing
( Ref Nothing ( Imported ( ModuleName "Control.Bind" ) ( Name "bind" ) ) )
( App Nothing
( ObjectProp Nothing
Expand All @@ -249,23 +249,23 @@ UberModule
) :| []
)
( Abs Nothing
( ParamNamed Nothing ( Name "f$1092" ) )
( ParamNamed Nothing ( Name "f$689" ) )
( Abs Nothing
( ParamNamed Nothing ( Name "a$1093" ) )
( ParamNamed Nothing ( Name "a$690" ) )
( App Nothing
( App Nothing
( Ref Nothing ( Local ( Name "bind$1090" ) ) )
( Ref Nothing ( Local ( Name "f$1092" ) ) )
( Ref Nothing ( Local ( Name "bind$687" ) ) )
( Ref Nothing ( Local ( Name "f$689" ) ) )
)
( Abs Nothing
( ParamNamed Nothing ( Name "f'$1094" ) )
( ParamNamed Nothing ( Name "f'$691" ) )
( App Nothing
( App Nothing
( Ref Nothing ( Local ( Name "bind$1090" ) ) )
( Ref Nothing ( Local ( Name "a$1093" ) ) )
( Ref Nothing ( Local ( Name "bind$687" ) ) )
( Ref Nothing ( Local ( Name "a$690" ) ) )
)
( Abs Nothing
( ParamNamed Nothing ( Name "a'$1095" ) )
( ParamNamed Nothing ( Name "a'$692" ) )
( App Nothing
( App Nothing
( Ref Nothing
Expand All @@ -290,8 +290,8 @@ UberModule
)
)
( App Nothing
( Ref Nothing ( Local ( Name "f'$1094" ) ) )
( Ref Nothing ( Local ( Name "a'$1095" ) ) )
( Ref Nothing ( Local ( Name "f'$691" ) ) )
( Ref Nothing ( Local ( Name "a'$692" ) ) )
)
)
)
Expand Down Expand Up @@ -369,9 +369,9 @@ UberModule
)
)
( Abs Nothing
( ParamNamed Nothing ( Name "x$1847" ) )
( ParamNamed Nothing ( Name "x$907" ) )
( Abs Nothing
( ParamNamed Nothing ( Name "b$1140$1831" ) )
( ParamNamed Nothing ( Name "b$895" ) )
( App Nothing
( App Nothing
( App Nothing
Expand Down Expand Up @@ -422,8 +422,8 @@ UberModule
)
( Abs Nothing ( ParamUnused Nothing )
( Abs Nothing
( ParamNamed Nothing ( Name "x$1843" ) )
( Ref Nothing ( Local ( Name "x$1843" ) ) )
( ParamNamed Nothing ( Name "x$902" ) )
( Ref Nothing ( Local ( Name "x$902" ) ) )
)
)
)
Expand All @@ -443,11 +443,11 @@ UberModule
]
)
)
( Ref Nothing ( Local ( Name "x$1847" ) ) )
( Ref Nothing ( Local ( Name "x$907" ) ) )
)
)
)
( Ref Nothing ( Local ( Name "b$1140$1831" ) ) )
( Ref Nothing ( Local ( Name "b$895" ) ) )
)
)
)
Expand Down Expand Up @@ -498,7 +498,7 @@ UberModule
( PropName "foldl" )
)
( Abs Nothing
( ParamNamed Nothing ( Name "c$558$1821" ) )
( ParamNamed Nothing ( Name "c$372$882" ) )
( Abs Nothing ( ParamUnused Nothing )
( App Nothing
( App Nothing
Expand All @@ -515,7 +515,7 @@ UberModule
( PropName "one" )
)
)
( Ref Nothing ( Local ( Name "c$558$1821" ) ) )
( Ref Nothing ( Local ( Name "c$372$882" ) ) )
)
)
)
Expand Down
Loading
Loading