fix: train late-binding channel on its source column (#52)

Author: mcanouil

CHANGELOG.md

@@ -12,6 +12,7 @@
 - feat: `labs()` fields default to `auto`; pass `none` to suppress an axis or legend title and reclaim the space it reserved. (#12)
 - feat: `element-blank()` on a text surface (axis, plot, or legend title) collapses the space the text would reserve. (#12)
 - feat: `width`/`height` accept `auto` to fill the available space of a bounded container. (#10)
+- fix: a `stage`/`after-scale` channel now trains its scale on the marker's source column, so the closure receives the channel's scale-resolved value as documented. (#52)
 - fix: grouped geoms (e.g., `geom-smooth`) no longer panic when a grouping aesthetic is mapped via `after-scale` or `stage`. (#51)
 - fix: facet-grid legends reserve space for a secondary x-axis so a top legend no longer overlaps its ticks. (#46)
 - fix: `geom-errorbar()`/`geom-errorbarh()`/`geom-rug()` resolve an `after-scale`/`stage` colour mapping instead of panicking. (#45)

src/scale/train.typ

@@ -9,7 +9,9 @@
 #import "../data.typ": column
 #import "../utils/types.typ": infer-column-type, parse-number
 #import "../utils/typst-markup.typ": is-typst-markup
-#import "../utils/late-binding.typ": is-late-binding, late-binding-name
+#import "../utils/late-binding.typ": (
+  after-scale-source, is-late-binding, late-binding-name,
+)
 
 #let _resolve-mapping(layer, plot-mapping) = {
   if layer.at("inherit-aes", default: true) and plot-mapping != none {
@@ -260,19 +262,25 @@
     for a in aes-list {
       let raw = layer-mapping.at(a, default: none)
       if raw == none { continue }
-      // Late-binding markers (`after-stat`, `from-theme`, ...) carry no
-      // column to train against; evaluation happens in `_prepare-layer`
-      // for `from-theme` and post-stat for `after-stat`.
-      if is-late-binding(raw) { continue }
-      let col-name = mapping-ref-col(raw)
+      // Late-binding markers that carry a `source` column (an `after-scale`
+      // derived from `stage(start: ...)`) train on that column so the per-row
+      // resolver hands the closure the channel's scale-resolved value, as
+      // documented on `@after-scale`. Source-less markers (`after-stat`,
+      // `from-theme`, and a pure `after-scale` closure) are evaluated
+      // elsewhere; skip them.
+      let train-ref = if is-late-binding(raw) { after-scale-source(raw) } else {
+        raw
+      }
+      if train-ref == none or is-late-binding(train-ref) { continue }
+      let col-name = mapping-ref-col(train-ref)
       let entry = cache.at(a)
       entry.cols.push((
         name: col-name,
         values: column(layer-data, col-name),
-        forced-type: _resolve-forced-type(raw, layer-data, col-name),
+        forced-type: _resolve-forced-type(train-ref, layer-data, col-name),
         levels: layer-factor-levels.at(col-name, default: none),
       ))
-      if is-typst-markup(raw) { entry.typst-mark = true }
+      if is-typst-markup(train-ref) { entry.typst-mark = true }
       cache.insert(a, entry)
     }
   }

tests/unit/test-after-scale-trains.typ

@@ -0,0 +1,46 @@
+// A `stage`/`after-scale` marker that carries a `source` column trains the
+// channel scale on that column (post `apply-stages`), so the per-row resolver
+// hands the closure the scale-resolved value as `@after-scale` documents.
+// Regression for #50: previously `train` skipped every late-binding marker,
+// leaving the channel's domain empty and the closure receiving ink.
+
+#import "../../src/render.typ": _prepare-layer
+#import "../../src/scale/train.typ": train
+#import "../../src/aes.typ": aes
+#import "../../src/utils/late-binding.typ": after-scale, stage
+#import "../../src/geom/point.typ": geom-point
+
+#let raw = (
+  (x: 1, y: 1, g: "a"),
+  (x: 2, y: 2, g: "a"),
+  (x: 3, y: 3, g: "b"),
+  (x: 4, y: 4, g: "b"),
+)
+
+// A stage-derived after-scale marker trains its channel on the source column.
+#let mapping-stage = aes(
+  x: "x",
+  y: "y",
+  colour: stage(start: "g", after-scale: (c, _) => c),
+)
+#let prepared-stage = (_prepare-layer(geom-point(), mapping-stage, raw),)
+#let trained-stage = train(
+  layers: prepared-stage,
+  mapping: mapping-stage,
+  data: raw,
+)
+#assert.eq(trained-stage.colour.type, "discrete")
+#assert.eq(trained-stage.colour.domain, ("a", "b"))
+
+// A pure `after-scale` carries no source, so the channel scale stays untrained
+// (the closure is expected to compute the value from `ctx`).
+#let mapping-pure = aes(x: "x", y: "y", colour: after-scale((v, _) => v))
+#let prepared-pure = (_prepare-layer(geom-point(), mapping-pure, raw),)
+#let trained-pure = train(
+  layers: prepared-pure,
+  mapping: mapping-pure,
+  data: raw,
+)
+#assert.eq(trained-pure.at("colour", default: none), none)
+
+after-scale source training tests passed.