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)

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.

[Astra] Rechecked at e785faeb3: this remains a P1 miscompilation. The new filters reject the original direct case, but mkWeb follows constructor arguments and destinations without maintaining those filters. A safe Box entry pulls the excluded A producer and mixed local/external consumer back into the web; taggedShapesOfMatchScrutinee then drops the unknown source again.

With :classTags mono, :js, and :noFreeze:

class Box(val x) class A(val x) class B(val x) fun unknownThroughField(flag, external) = let local = new A(0) let holder = new Box(local) let _ = @matchShapes(Box(_)) shape.match of holder, () => 0 let value = if flag then local else external @matchShapes(A(_), B(_)) shape.match of value, () => 1, () => 2

unknownThroughField(false, new B(0)) returns 1 instead of 2, with no compilation diagnostic. Please enforce safety for the entire reachable web and validate every scrutinee source before lowering. Recorded in class-tags/ReviewFollowup.mls, unknownThroughField.

Regression evidence is committed locally in 2fbdb7344188d8303e81968994ae664a58c1380d (not pushed).

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.

LPTK · 2026-10-06T04:59:36Z

Looks like the merge needs a golden test output update.

Reproduce unsafe external-flow reentry, constructor field rewrites, and unlowered zero-branch shape matches. Preserve the shape-budget exhaustion reproducer behind :exit, verify anonymous recursive and curried-return branches, and refresh class-tag golden output after the upstream debug-printer changes.

LPTK

[Astra] Follow-up review of e785faeb399a76ec4a652d034d3c96b2afeb71b7: the review comments are not all satisfactorily addressed. I checked the code and reran the regression cases, including variants that exercise the new restrictions.

Of the eleven earlier Astra findings, seven are addressed within the newly restricted feature scope:

Finding	Verification
Polymorphic allocation contexts	Original runtime regression now returns the expected difference of 1.
Dynamic field/null tag reads	The original case is now rejected because its constructor argument can come from outside.
Recursive branch expansion	Named branches are rejected; an anonymous branch making an ordinary recursive call compiles and returns 42.
Main blocks and class initializers	Shape matches in those scopes now have explicit diagnostics.
Nested functions	Named nested-function matches now have an explicit unsupported-scope diagnostic.
Partially covered function entries	The `inspect`/`first`/`second` regression now lowers all matches.
Curried branch calling semantics	Named curried branches are rejected; an anonymous branch returning another lambda preserves the returned function.

Four Astra findings remain, plus the earlier budget comment:

P1: Indirect flow still admits unknown scrutinees; the new case returns 1 instead of 2.
P1: Constructor writes still produce stale tags; the new case returns 1 instead of 2.
P2: Malformed intrinsic calls can escape lowering; a zero-branch call compiles and crashes at runtime.
P2: Wasm remains ungated; Todos.mls still records unresolved-field backend errors.
P2: The shape budget is enforced after expansion; a 21-field case exhausts a capped JVM before the budget diagnostic.

The other earlier review concerns—producer cycles, frozen-object handling, tuple elements, short-circuiting, fresh symbols for copied branches, and misplaced annotations—have corresponding guards or implementation changes. Frozen objects are explicitly unsupported while the option requires noFreeze; that is a support restriction, not implementation of tags on frozen objects.

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

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

LPTK reviewed Oct 6, 2026

View reviewed changes


		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!

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!

LPTK commented Oct 6, 2026

Uh oh!

LPTK left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

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

LPTK left a comment •

edited

Loading