Add Class Tag Optimization Pass by NeilKleistGao · Pull Request #555 · hkust-taco/mlscript

NeilKleistGao · 2026-08-28T09:01:59Z

Reuse web computation.
Add builtin intrinsic shape.match and annotation @matchShapes.
Insert tags for class instances in a web and transform shape.match into tag checks.

~~Problems to be addressed:~~

~~How to track linearity (discussed in the meeting and decided to insert tags for class instances for pattern matching for now).~~
~~Block IR uses a flattened pattern matching, which makes tag checks difficult. e.g.,~~

if ls is
    Cons(1, Cons(2, Cons(3, Cons(x, Cons(y, Cons(z, Nil)))))) then x + y + z

~~It is flattened into:~~

if ls is
  Cons then
    let x = ls.x
    let xs = ls.xs
    if x == 1 and xs is Cons then ...

Assume that Cons(1, Cons(2, Cons(3, Cons(dyn, Cons(dyn, Cons(dyn, Nil))))))'s tag is 0, we need to check if we want to replace the whole huge matching with the tag checking. It is not clear when to set the boundary to say that the inner irrelevant matchings are untouched. e.g.,

let tmp1 = ...
let tmp2 = ...
let tmp3 = ...
if tmp1 is ....

~~how should we know whether we are now checking against a nested pattern (e.g., C(D(...), E(...), F(...))), or this is just a user code?~~

TODOs in future:

We may need to make the tag a real field in classes if the JIT cannot optimize the code (already tracked in the code)

NeilKleistGao · 2026-09-02T09:56:55Z

I think one possible way to address the second problem is:

We first compute all possible tags for a given scrutinee.
e.g., for if x is ..., x's tag can range from 1 to 5, which means x can have 5 different shapes.
For a branch against a pattern, we filter out impossible tags. e.g., if tag 1 is for C(D(_)), and tag 2 is for C(E(1)), and other tags are not for C, then we transform if x is C into if x.__tag == 1 || x.__tag == 2
We track aliases for x's fields. e.g., let tmp = x.m, and when tmp becomes a nested matching's scrutinee, we check its shape tags against corresponding tags for x's fields. e.g., say D(_)'s shape tag is 10. If x.m's tag is 10, and tmp's tag is also 10, then we know that x.__tag == 1 implies tmp.__tag == 10. Then we split if x.__tag == 1 || x.__tag == 2 into if x.__tag == 1 and if x.__tag == 2, and omit check tmp.__tag == 10 in the first branch.

LPTK · 2026-09-02T14:46:16Z

Hmm that does seem quite complicated. I don't think we actually need to transform the pattern matches in such a fine-grained way. We should be able to replace the entire deep pattern match by a single tag-check in one go. Indeed, the extracted ctor fields are only used for matching the shape itself, and not for something else; we should be able to not extract them at all. This might require a bit of a variable use analaysis, though. Or, better, we should simply find a way of annotating those "shape" pattern matches, which always discard the extracted fields, so the IR pass does not have to make a guess.

Actually, now that I think of it, why don't we just generate code that already includes all the necessary info (in the form of annotations) for easily compiling the shape match away? Eg, something like:

@matchShapes(C(D(_)), C(E(1)), F...) if x is // this is a top-level shape matching node
  C then
    @shapeMatchingOnly let y = x.f1
    if y is
      @shapeMatchingOnly D then ...
      @shapeMatchingOnly E then ...
  F then
    ...

NeilKleistGao · 2026-09-03T09:05:50Z

Hmm that does seem quite complicated. I don't think we actually need to transform the pattern matches in such a fine-grained way. We should be able to replace the entire deep pattern match by a single tag-check in one go. Indeed, the extracted ctor fields are only used for matching the shape itself, and not for something else; we should be able to not extract them at all. This might require a bit of a variable use analaysis, though. Or, better, we should simply find a way of annotating those "shape" pattern matches, which always discard the extracted fields, so the IR pass does not have to make a guess.

Actually, now that I think of it, why don't we just generate code that already includes all the necessary info (in the form of annotations) for easily compiling the shape match away? Eg, something like:
@matchShapes(C(D(_)), C(E(1)), F...) if x is // this is a top-level shape matching node
  C then
    @shapeMatchingOnly let y = x.f1
    if y is
      @shapeMatchingOnly D then ...
      @shapeMatchingOnly E then ...
  F then
    ...

Actually, it seems that the @matchShapes annotation alone is enough. We do not flatten the UCS by ourselves, i.e., we only write (or generate)

@matchShapes(C(D(_)), C(E(1)), F(_))
if x is
  C(D(_)) then ...
  C(E(1)) then ...
  F(_) then ...

Those @shapeMatchingOnly annotations can be automatically done by checking sub matches.

LPTK · 2026-09-04T08:18:28Z

Suggestion from the meeting:

Instead of

@matchShapes(C(D(_)), C(E(1)), F(_)) if x is
  C(D(_)) then ...
  C(E(1)) then ...
  F(_) then ...

Use an intrinsic like

@matchShapes(C(D(_)), C(E(1)), F(_)) shapeMatch(x,
  () => ...
  () => ...
  () => ...
)

which can be easily desugared after shape analysis to

~>
if x.tag is
  42 then ...
  43 then ...
  44 then ...

Copilot

Copilot review overview

🟡 Changes recommended

Tag insertion currently fails with default-frozen JS objects and Wasm, while several shape-analysis paths can miscompile or fail to terminate.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 7 · 1 · 1

Open (9)

Recursive producer flows cause compiler stack overflow · New
Tag assignment fails after immutable instance freezing · New
Tuple shape validation ignores element shapes · New
Shape checks lack short-circuit evaluation · New
Copied match branches reuse bound symbols · New
Uneliminated shape.match intrinsic can emit undefined calls · New
Class-tag pass emits unsupported Wasm dynamic fields · New
Shape alternatives cause unbounded Cartesian product growth · New
No-effect warning suppressed for non-shape.match annotations · New

What changed in this PR

Adds an opt-in compiler pass that derives class-instance shapes, assigns runtime tags, and lowers shape.match calls into tag-based dispatch.

Changes:

Adds shape.match, @matchShapes, and class-tag configuration.
Extends flow analysis and compilation with tag insertion and match rewriting.
Adds diff tests for basic, nested, functional, invalid, and subsumption cases.

File	Description
`hkmc2DiffTests/src/test/scala/hkmc2/MLsDiffMaker.scala`	Adds the `:classTags` test directive.
`hkmc2/shared/src/test/mlscript/invalml/InvalMLPrelude.mls`	Declares the new annotation and intrinsic.
`hkmc2/shared/src/test/mlscript/decls/Prelude.mls`	Adds public prelude declarations.
`hkmc2/shared/src/test/mlscript/class-tags/Subsumption.mls`	Tests shape subsumption.
`hkmc2/shared/src/test/mlscript/class-tags/Nested.mls`	Tests nested shapes and tuples.
`hkmc2/shared/src/test/mlscript/class-tags/Func.mls`	Tests interprocedural flows.
`hkmc2/shared/src/test/mlscript/class-tags/Basic.mls`	Tests basic tag generation and dispatch.
`hkmc2/shared/src/test/mlscript/class-tags/BadShapes.mls`	Tests malformed shape diagnostics.
`hkmc2/shared/src/test/mlscript/class-tags/Annotations.mls`	Tests annotation handling.
`hkmc2/shared/src/main/scala/hkmc2/semantics/Term.scala`	Represents shape annotations semantically.
`hkmc2/shared/src/main/scala/hkmc2/semantics/Elaborator.scala`	Elaborates the annotation and intrinsic.
`hkmc2/shared/src/main/scala/hkmc2/Config.scala`	Adds pass configuration.
`hkmc2/shared/src/main/scala/hkmc2/codegen/Lowering.scala`	Recognizes the intrinsic during lowering.
`hkmc2/shared/src/main/scala/hkmc2/codegen/flowAnalysis/FlowAnalysis.scala`	Models shape-match data flow.
`hkmc2/shared/src/main/scala/hkmc2/codegen/EtaExpansion.scala`	Supplies elaboration context to flow analysis.
`hkmc2/shared/src/main/scala/hkmc2/codegen/DeadParamElim.scala`	Supplies elaboration context to flow analysis.
`hkmc2/shared/src/main/scala/hkmc2/codegen/CompilationPipeline.scala`	Registers the class-tag pass.
`hkmc2/shared/src/main/scala/hkmc2/codegen/ClassTagsTransformer.scala`	Implements shape collection, tagging, and rewriting.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

Add 14 reproduction cases for polymorphic tag loss, stale mutable shapes, unknown and null inputs, incomplete scope and web traversal, recursive and curried branches, surviving intrinsics, and Wasm tag fields. Preserve the expected results and current golden diagnostics with explicit fixme markers. Validation: ctest, focused class-tag and Wasm diff-tests, and hkmc2AllTests/test.

LPTK

[Astra] I reproduced 11 issues and committed 14 regression cases with golden outputs in ac25c9565. hkmc2AllTests/test passes with known defects explicitly marked :fixme; the compiler defects remain unresolved. The findings below include silent miscompilations in polymorphic allocation contexts, mutable field shapes, and mixed known/unknown scrutinees.

LPTK · 2026-09-28T08:22:40Z

+
+  private val producersInWeb = webs.iterator.flatMap(_.markedProducers).toSet
+
+  private val ctorsByResultId = producersInWeb.iterator.map(ctor => ctor.exprId -> ctor).toMap


[Astra] [P1] Preserve every polymorphic allocation context

With the default mono=false, multiple Ctor instances share an exprId. This map retains an arbitrary one, so wrap(A(0)) and wrap(B(0)) receive the same tag. The poly regression returns a difference of 0 instead of 1. Group and reconcile all contexts per allocation site; the entry collector at lines 46–48 has the same lossy mapping.

Reproduction: ReviewPolymorphism.mls.

Fixed: the previous version only considered the mono case where one expirId maintains only one Ctor. Now the map maintains a list of Ctors.

LPTK · 2026-09-28T08:22:41Z

+      .flatMap(_.srcs)
+      .collect:
+        case ctor: Ctor => ctor
+      .toList.distinct.flatMap: ctor =>
+        taggedShapesByProducer.getOrElse(ctor, Nil)


[Astra] [P1] Include unknown sources in the exhaustiveness check

Collecting only tagged constructors silently discards UnknownProd and untagged sources. A mixed local/external scrutinee consequently passes validation: unknownAlternative(0, new B(0)) returns unit instead of 2. Every possible source must have a compatible tag, otherwise compilation needs a diagnostic or semantic fallback.

Reproduction: unknownAlternative in ReviewShapes.mls.

The previous web computation failed to filter out variables that may receive external data. Now it is fixed.

LPTK · 2026-09-28T08:22:41Z

+      // * Generate branch based on the branch function
+      private def mkBranch(branch: FunDefn, resultSymbol: TempSymbol): Block =
+        SymbolRefresher(Map.empty).apply(applyFunBodyLikeBlock(branch.body)).mapReturn:
+          case Return(result) => Assign(resultSymbol, result, End())


[Astra] [P1] Prevent recursive branch expansion

mkBranch recursively transforms the branch body without tracking active expansions. A recursive branch that terminates after one runtime call causes a compiler StackOverflowError. Emitting ordinary branch calls would let the existing inliner handle recursion safely; otherwise expansion needs an explicit recursion guard. This is separate from cycle detection while computing producer shapes.

Reproduction: recursiveBranch in ReviewBranches.mls.

Now we restrict shape.match to take only anonymous lambda functions as parameters, which is enough for the specialization.

LPTK · 2026-09-28T08:22:41Z

+        (entries.producers.nonEmpty || entries.consumers.nonEmpty)
+          && !entries.producers.exists(coveredProducers)
+          && !entries.consumers.exists(coveredConsumers)
+      then
+        val web = mkWeb(entries)


[Astra] [P2] Retain uncovered entries in partially covered functions

One covered entry suppresses all entries belonging to the function. Visiting inspect first covers the shared Box web, causing independent A and B matches in first and second to be omitted and rejected. Process the uncovered entries instead of skipping the whole function.

Reproduction: inspect, first, and second in ReviewScopes.mls.

The previous web computation failed to connect some elements in the same function. Now it is fixed.

NeilKleistGao added 9 commits August 25, 2026 16:53

Draft dpf

dc4957a

WIP: Build web

207ab7a

WIP: Fix dup web

fa30dd9

WIP: Use pipeline

cf9f267

Merge branch 'hkmc2' of github.com:hkust-taco/mlscript into web🦚

cb66e8b

WIP: Fix merge errors

580483e

WIP: Retrieve shape information from flow analysis

6b30c90

WIP: Fix web construction

f91c6b1

WIP: Insert tags and minor fix

356f9b2

NeilKleistGao added 2 commits September 3, 2026 13:40

WIP: Remove union tag insertion

623db84

WIP: Add annotations for match

2e64870

WIP: Minor

5641244

NeilKleistGao added 13 commits September 8, 2026 10:50

WIP: Clean

2becb73

WIP: Add checks in flow analysis

0593648

WIP: Match branches

65cc0dd

WIP: Add tests

9846bef

WIP: Fix flow analysis and consumer comp

c44134d

WIP: Minor fix

eaf5508

WIP: Add tests

52cdd98

WIP: Add join point checks

2b329ce

WIP: Rename the pass

1595976

WIP: Check patterns not matched

6d279ae

WIP: Refactor

b5ca410

WIP: Refactor

973f81f

Misc changes

b9232cc

NeilKleistGao changed the title ~~Add Data Representation Flattening Pass~~ Add Class Tag Optimization Pass Sep 16, 2026

NeilKleistGao requested a review from LPTK September 21, 2026 05:23

LPTK requested a balanced review from Copilot September 21, 2026 11:07

Copilot started reviewing on behalf of LPTK September 21, 2026 11:08 View session

Copilot AI reviewed Sep 21, 2026

View reviewed changes

NeilKleistGao added 9 commits September 22, 2026 14:44

Merge from hkmc2

05664b1

Add cycle detection

e52b7e7

Add error for missing noFreeze config

eb94ba8

Add checks for tuple elements

ae58e70

Use SymbolRefresher

1f5ff87

Add error messages

00b14ba

Add warning messages

b63220d

Add short-circuit and comments

9d24fdb

Merge from hkmc2

dce95c4

LPTK requested a balanced review from Copilot September 25, 2026 13:06

Copilot AI reviewed Sep 25, 2026

LPTK reviewed Sep 28, 2026

View reviewed changes

NeilKleistGao and others added 13 commits September 28, 2026 17:36

WIP: Add errors for mutable features and add todo for tag field

9133c14

WIP: Add check for tag check

6a79c98

WIP: Fix web computation & test cases

cf21f47

WIP: Report for unannotated match

e30be24

WIP: Reject named functions as arms

324f553

WIP: Add a pass to reject unsupported definitions

c537eee

WIP: Fix web computation across functions

314c787

WIP: Fix poly

2a05f35

Add threshold for flatten

6f7b018

Do not insert tags for objects that are not matched by shape.match

3585f95

Relax some changes

93246cc

Merge from hkmc2

c748224

Merge branch 'hkmc2' into web🦚

e785fae


		private val producersInWeb = webs.iterator.flatMap(_.markedProducers).toSet

		private val ctorsByResultId = producersInWeb.iterator.map(ctor => ctor.exprId -> ctor).toMap

Hidden character warning

Conversation

NeilKleistGao commented Aug 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

NeilKleistGao commented Sep 2, 2026

Uh oh!

LPTK commented Sep 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

NeilKleistGao commented Sep 3, 2026

Uh oh!

LPTK commented Sep 4, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

LPTK left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

NeilKleistGao commented Aug 28, 2026 •

edited

Loading

LPTK commented Sep 2, 2026 •

edited

Loading