[ehigham]

Stack 10. Based on #15928.

Why

TypeCheck.checkSingleNode was one match of about 20.8 KB of bytecode. HotSpot never
compiles a method over 8000 bytes (DontCompileHugeMethods), so TypeCheck always ran in
the interpreter. A per-node tag fixes that and gives any pass that dispatches on BaseIR a
jump table instead of an instanceof chain.

Summary

 checkSingleNode(ir)
-  ir match                                   // ~140 instanceof tests, one 20.8 KB method
-    case I32(_) =>
-    case x @ Ref(name, _) => ...
+  (ir.tag.family.toInt: @switch) match       // jump table on 13 families
+    case Family.Core => checkCore(ir, env)
+      (ir.tag.toInt: @switch) match          // jump table on the family's nodes
+        case Tag.Ref =>
+          val x = ir.asInstanceOf[Ref]
+          ... x.name ...

ir-gen names each node's family where the node is defined, and generates the tags.
Hand-written nodes sit among the generated ones, so their tags follow the same order:

// ir-gen/src/Main.scala
r += node(Const, "I32", in("x", att("Int"))).withTraits(Atom)
r += handWritten(Table, "TableRead")   // class written by hand; only its tag is generated

// Tag_gen.scala (generated)
final class Tag private[ir] (private val repr: Short) extends AnyVal {
  def toInt: Int = repr.toInt
  def family: Family = new Family(toInt >>> Tag.FamilyShift)
}
object Tag { final val I32 = 1; final val Ref = 260; final val TableRead = 2584; ... }

// node classes
abstract class BaseIR(final val tag: Tag)
final case class I32(x: Int) extends IR(new Tag(Tag.I32)) with Atom

Relational, Table and Matrix hold the nodes a lowering pass removes. A node goes in the
family of the first invariant that forbids it. The other families are buckets that keep
each method small, and a node goes in the family of what it produces.

family	rule
Const	a constant value
Core	a ref or input, or binds, branches, loops, throws, calls a function or applies a built-in operator
Struct	produces or reads a struct or tuple
Stream	produces, reads, transforms or writes a stream
Collection	produces or reads an array, set or dict
Agg	uses an agg context or agg state
IO	reads or writes a partition, value or metadata
NDArray	produces, reads, transforms or writes an ndarray
Random	creates or splits RNG state
Relational	first forbidden by NoRelationalLets
Table	first forbidden by CompilableIR
Matrix	first forbidden by NoMatrixIR, so LowerMatrixIR removes the whole family
BlockMatrix	forbidden with Table, but reads or makes a block matrix and no Table

• Within a family, nodes are ordered so the cases TypeCheck and ChildEnv handle sit on
  consecutive tags. Every check* method compiles to a tableswitch.
• tag is a final constructor field, so reading it is a field load. An overridden
  def tag was a virtual call HotSpot could not inline, and made TypeCheck 11% slower on
  a 1000-field MakeStruct.
• Tag stores a Short, which fits in padding most node classes already have. 203 of 217
  classes keep their size. Weighted by node visits in the Python test suite, the average
  node grows from 42.4 to 42.5 bytes.
• The largest TypeCheck method, checkStream, is 4.2 KB.
• The NDArrayIR trait goes. ir-gen marks those nodes TypedIR[TNDArray] instead.

In IRPassBenchmark, against #15928, TypeCheck drops from 4.9 ms to 1.0 ms on a
3000-binding Block, from 5.3 ms to 0.58 ms on 2000 nested StreamMaps, from 6.6 ms to
0.92 ms on a balanced tree and from 1.5 ms to 0.26 ms on a 1000-field MakeStruct. The
gain comes from getting each method under 8000 bytes. A version that only split the
instanceof chain into four methods ran within noise of the tag version over 3 forks.

🤖 Generated with Claude Code

[coderabbitai]

Warning

Review limit reached

You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository.

Next included review available in 59 minutes.

Check out review usage here.

View limit details

Limit details: You’ve used all 4 included reviews currently available.

Learn how review limits work.

Review configuration:

Run configuration

• Source mode: Auto
• Repository: Unavailable
• Branch: Unavailable
• File: No configuration file from this repository applied
• Commit: Unavailable
• Layers: Organization UI
• Review profile: CHILL
• Plan: Advanced
• Run ID: c2f94258-620c-4ea8-a8e3-dbf0eb341cdd

Commits

Reviewing files that changed from the base of the PR and between 8a8bf3b and e04d4e7.

Files selected for processing: 10

• hail/hail/ir-gen/src/Main.scala
• hail/hail/src/is/hail/expr/ir/BaseIR.scala
• hail/hail/src/is/hail/expr/ir/BlockMatrixIR.scala
• hail/hail/src/is/hail/expr/ir/ChildEnv.scala
• hail/hail/src/is/hail/expr/ir/IR.scala
• hail/hail/src/is/hail/expr/ir/MatrixIR.scala
• hail/hail/src/is/hail/expr/ir/TableIR.scala
• hail/hail/src/is/hail/expr/ir/TypeCheck.scala
• hail/hail/src/is/hail/expr/ir/lowering/ForwardLets.scala
• hail/hail/src/is/hail/expr/ir/lowering/invariant/package.scala

No actionable comments were generated in the recent review.

Walkthrough

The IR generator now assigns families and emits node tags. IR nodes carry those tags, and TypeCheck dispatches validation by family and node tag. The change also removes the NDArrayIR trait and updates specified NDArray nodes to use TNDArray typing.

Changes: IR family tags and validation dispatch

Layer / File(s)	Summary
Define families and generate tags hail/hail/ir-gen/src/Main.scala	The generator assigns families to generated and hand-written IR nodes. It emits family and node tag constants in Tag_gen.scala.
Pass tags through IR node classes hail/hail/src/is/hail/expr/ir/BaseIR.scala, IR.scala, TableIR.scala, MatrixIR.scala, BlockMatrixIR.scala	BaseIR stores a tag. Generated and hand-written IR nodes pass node-specific tags to their parent classes. The NDArrayIR trait is removed, and specified NDArray nodes use TNDArray typing.
Dispatch validation by family hail/hail/src/is/hail/expr/ir/TypeCheck.scala	TypeCheck switches on family and node tags. Family-specific checks cover IR node validation, including stream folds and NDArray matrix requirements.

Estimated code review effort: 4 (Complex) | ~60 minutes

Change: Refactor

Sequence Diagram(s)

sequenceDiagram
  participant IRGenerator
  participant IRNode
  participant TypeCheck
  IRGenerator->>IRGenerator: Assign families and emit node tags
  IRGenerator->>IRNode: Generate node classes with tags
  IRNode->>TypeCheck: Provide family and node tag
  TypeCheck->>TypeCheck: Dispatch to family-specific checks

Loading

Possibly related PRs

• hail-is/hail#15913: Changes validation for aggregate and scan nodes that this PR redistributes into tagged family validators.
• hail-is/hail#15855: Adds aggregate and scan scope checks that this PR moves into per-family and per-tag methods.
• hail-is/hail#15928: Changes child traversal APIs used by TypeCheck; this PR retains that traversal contract while changing validation dispatch.

Merge Risk: Minimal · previous review 8a8bf

This change replaces TypeCheck's large pattern match with tag-based dispatch, to let the JVM compile it. No actionable merge-blocking risk was identified in the supplied context.

Pre-merge checks: 3 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Description check		The description explains the motivation (the 20.8 KB method exceeds HotSpot's 8000-byte compile limit), the design (families, tags as a final constructor field, Short storage), the NDArrayIR trait rem…	Add the template's "Security Assessment" section. Choose whether the change potentially impacts the Hail Batch instance, then select one Impact Rating. Write an Impact Description, with a one or two sentence justification for a none or low …

Passed checks: 3 passed

Check name	Status	Explanation
Title check		The title "[query] Dispatch TypeCheck on a per-node tag" matches the main change. The PR replaces the large instanceof chain in TypeCheck.checkSingleNode with two-level dispatch on per-node tags group…
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Full details: Title check

Explanation

The title "[query] Dispatch TypeCheck on a per-node tag" matches the main change. The PR replaces the large instanceof chain in TypeCheck.checkSingleNode with two-level dispatch on per-node tags grouped into families. The title is specific and concise, and it names the primary change from the developer's view.

Full details: Description check

Explanation

The description explains the motivation (the 20.8 KB method exceeds HotSpot's 8000-byte compile limit), the design (families, tags as a final constructor field, Short storage), the NDArrayIR trait removal, and benchmark results. It does not follow the repository template. It omits the required "Change Description" heading and the "Fixes #<issue_number>" line, and it has no "Security Assessment" section. That section requires an Impact Rating, an Impact Description, and an Appsec Review checkbox, which are absent. The description is mostly complete on technical content but is missing required template sections.

Resolution

Add the template's "Security Assessment" section. Choose whether the change potentially impacts the Hail Batch instance, then select one Impact Rating. Write an Impact Description, with a one or two sentence justification for a none or low rating, or a description of impact and mitigations for a medium or high rating. Complete the Appsec Review checkbox if the rating is medium or high. Also add the "Change Description" heading with the "Fixes #<issue_number>" line, or delete that line if no issue applies. Keep the existing explanation under the Change Description heading.

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

Share

• X
• Mastodon
• Reddit
• LinkedIn

Recent review info

Run configuration

• Source mode: Auto
• Repository: Unavailable
• Branch: Unavailable
• File: No configuration file from this repository applied
• Commit: Unavailable
• Layers: Organization UI
• Review profile: CHILL
• Plan: Advanced
• Run ID: 1f38f276-38e5-4d27-973a-a8939fc690ad

Commits

Reviewing files that changed from the base of the PR and between 922e8b3 and 8a8bf3b.

Files selected for processing: 1

• hail/hail/ir-gen/src/Main.scala

Included review availability: This review used your included allowance. Your plan provides up to 4 included reviews per hour; 2 remain after this review.