Deforest rewriting improvements by ychenfo · Pull Request #499 · hkust-taco/mlscript

ychenfo · 2026-05-24T20:21:57Z

This PR fixes a deforestation bug where the rewriter wrongly turns implicit returns to explicit returns. Turning implicit returns to explicit ones is needed for branch-body and the continuation "rest" functions whose bodies are originated from top level blocks, but wrong for rewriting blocks like class constructors: it could produce explicit return super() and made code after super() unreachable, which causes bugs. This is now fixed.

This PR also avoids traversing the program when there is no fusion, and adds some high level comments for the rewriting process.

TODO here and in following PRs

Temporarily revert the use of irDefn to re-expose the bug
Update classCtorSymbol to use ClassCtorSymbol
Remove the special method-body deforestation logic (even tually, we'll do most optimization at the top level); also, this logic would likely break when we make the analysis understand methods
Fix this – only support it if doing so does not incur significant extra complexity
```
module Test with
  module Test2 with
    let x = 1
    fun foo(a) = if a is
      A(y) then x + y
Test.Test2.foo(A(2))
```
From the meeting discussion:
- reconstruct paths based on owner chains that are rooted in toplevel symbols
- this also changes in Rewrite's handling of FVs
As part of the compiler cache's Artifact computation, store the polymorphic analysis types of each top-level-like definition
Scope patmat branches, now that we use proper LetSplit and shouldn't share variables across splits as before

LPTK · 2026-05-25T09:19:26Z

-  // marks if a function or lambda is affine, i.e. called at most once.
-  // for functions with multiple parameter lists, `whichParamList` is the zero-based parameter-list index.
+  // marks if a function or lambda is one-shot, i.e. called at most once.
+  // `whichParamList` is the zero-based index of the parameter list for functions with multiple parameter lists.


I still don't get it, from this comment. Why is there a parameter list index?!

I've tried to further clarify this comment in b4e9fb6

LPTK · 2026-05-25T09:20:57Z

Turning implicit returns to explicit ones is needed for branch-body and the continuation "rest" functions whose bodies are originated from top level blocks, but wrong for rewriting blocks like class constructors: it could produce explicit return super() and made code after super() unreachable, which causes bugs. This is now fixed.

Wow... If I'd known the so-called "implicit returns" would cause such problems, I would have pushed for removing them earlier. They're mostly a legacy piece of design that should probably go. They should certainly not be the cause of such extra complexity.

# Conflicts: # hkmc2/shared/src/main/scala/hkmc2/codegen/deforest/Rewrite.scala

LPTK · 2026-05-25T14:32:45Z

I'd known the so-called "implicit returns" would cause such problems, I would have pushed for removing them earlier.

Ok, done. And I've merged this branch with the help of Codex, so please do review your own diff, making sure all the important relevant changes are still in.

Copilot

Pull request overview

This PR refines the deforestation (fusion) rewriting pass to avoid incorrectly converting non-returning contexts into value-returning return ... statements (notably in class/module constructor-like blocks), which could make subsequent constructor code unreachable (e.g., return super()).

Changes:

Make the rewriter conditional: skip program rewriting when the solver reports no fusions.
Introduce context-sensitive tail rewriting (returnsFromCurrentBlock) so “tail” fused calls are either emitted as Return(...) (function/lambda bodies) or as discarded statements (constructor/module bodies).
Add/clarify high-level comments and rename a comment in Annot.Affine to describe “one-shot” usage.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
hkmc2/shared/src/main/scala/hkmc2/semantics/Term.scala	Updates `Annot.Affine` doc comment to describe one-shot semantics more clearly.
hkmc2/shared/src/main/scala/hkmc2/codegen/deforest/Rewrite.scala	Adds context-sensitive tail rewriting to avoid value-returning `return ...` in constructor-like blocks; skips rewriting when there is no fusion.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

…mprovements Conflicts: hkmc2/shared/src/main/scala/hkmc2/codegen/deforest/Rewrite.scala hkmc2/shared/src/main/scala/hkmc2/codegen/flowAnalysis/FlowAnalysis.scala

LPTK

Is the PR otherwise ready to merge?

LPTK · 2026-05-30T04:39:59Z

 :deforest
 :expect 5
 :noInline
+:fixme


You should add the logic to treat self-references as free variables, right?

Will you fix this as part of this PR, or a follow-up? I'm fine with both options.

You should add the logic to treat self-references as free variables, right?

Yes.

Will you fix this as part of this PR, or a follow-up? I'm fine with both options.

I currently plan to fix this in a follow-up, maybe it's also better to fix this after #505 is merged?

Co-authored-by: Lionel Parreaux <lionel.parreaux@gmail.com>

LPTK · 2026-05-31T11:10:18Z

@ychenfo Pls fix the conflicts again.

…mprovements Conflicts: hkmc2/shared/src/main/scala/hkmc2/codegen/deforest/Rewrite.scala hkmc2/shared/src/test/mlscript/deforest/fusibility.mls

…l a class and a module

- remove error raising in extractor - uses `MemberRefTo` instead - update `Annot.Affine` documentation

LPTK

Net-negative PR despite added tests & documentation – nice!

ychenfo added 3 commits May 25, 2026 04:03

fix deforest rewriting about explicit/implicit return

35b4866

add high level comments for deforestation rewriting

b92bc53

white space

ac365d2

ychenfo mentioned this pull request May 24, 2026

Eta expansion #493

Merged

LPTK reviewed May 25, 2026

View reviewed changes

Merge branch 'hkmc2' into deforest-rewriting-improvements

5e7550e

# Conflicts: # hkmc2/shared/src/main/scala/hkmc2/codegen/deforest/Rewrite.scala

LPTK requested a review from Copilot May 25, 2026 14:32

Copilot started reviewing on behalf of LPTK May 25, 2026 14:33 View session

LPTK assigned ychenfo May 25, 2026

Copilot AI reviewed May 25, 2026

View reviewed changes

Comment thread hkmc2/shared/src/main/scala/hkmc2/codegen/deforest/Rewrite.scala Outdated

Comment thread hkmc2/shared/src/main/scala/hkmc2/codegen/deforest/Rewrite.scala

ychenfo added 5 commits May 29, 2026 02:40

temporarily revert the use of irDefn

685e55b

no need explicit/implicit return handling in deforest rewrite

cd93af5

remove the special method-body deforestation logic

2efcbc5

add todo related to the miscompilation of nested module

ae3d121

Merge remote-tracking branch 'origin/hkmc2' into deforest-rewriting-i…

38bcf2f

…mprovements Conflicts: hkmc2/shared/src/main/scala/hkmc2/codegen/deforest/Rewrite.scala hkmc2/shared/src/main/scala/hkmc2/codegen/flowAnalysis/FlowAnalysis.scala

LPTK requested changes May 30, 2026

View reviewed changes

ychenfo and others added 2 commits May 31, 2026 04:29

Update hkmc2/shared/src/test/mlscript/deforest/relaxedPrograms.mls

1d50c40

Co-authored-by: Lionel Parreaux <lionel.parreaux@gmail.com>

Remove WASM prelude quirk

3b039bb

LPTK force-pushed the deforest-rewriting-improvements branch from df3f951 to 3b039bb Compare May 31, 2026 06:47

ychenfo added 3 commits June 1, 2026 00:34

Merge remote-tracking branch 'origin/hkmc2' into deforest-rewriting-i…

7eb9e9c

…mprovements Conflicts: hkmc2/shared/src/main/scala/hkmc2/codegen/deforest/Rewrite.scala hkmc2/shared/src/test/mlscript/deforest/fusibility.mls

do not conflate object ref and class construction; make Prelude.Symbo…

15badc2

…l a class and a module

update comments

b4e9fb6

LPTK requested changes Jun 1, 2026

View reviewed changes

LPTK reviewed Jun 1, 2026

View reviewed changes

Comment thread hkmc2/shared/src/main/scala/hkmc2/semantics/Term.scala Outdated

address comments

209895d

- remove error raising in extractor - uses `MemberRefTo` instead - update `Annot.Affine` documentation

LPTK reviewed Jun 1, 2026

View reviewed changes

Comment thread hkmc2/shared/src/main/scala/hkmc2/codegen/flowAnalysis/FlowAnalysis.scala Outdated

Merge branch 'hkmc2' into deforest-rewriting-improvements

b11bba8

get rid of TrackedSymOf

7e95396

LPTK approved these changes Jun 1, 2026

View reviewed changes

LPTK merged commit 7e080c5 into hkust-taco:hkmc2 Jun 1, 2026
1 check passed

LPTK deleted the deforest-rewriting-improvements branch June 1, 2026 13:11

Conversation

ychenfo commented May 24, 2026 • edited by LPTK Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

TODO here and in following PRs

Uh oh!

LPTK May 25, 2026

Choose a reason for hiding this comment

Uh oh!

ychenfo May 31, 2026

Choose a reason for hiding this comment

Uh oh!

LPTK commented May 25, 2026

Uh oh!

LPTK commented May 25, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

LPTK left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

LPTK May 30, 2026

Choose a reason for hiding this comment

Uh oh!

ychenfo May 30, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

LPTK commented May 31, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

LPTK left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ychenfo commented May 24, 2026 •

edited by LPTK

Loading