polydat_core/compile/cone.rs
1// Copyright 2024-2026 Jonathan Shook
2// SPDX-License-Identifier: Apache-2.0
3
4//! SRD-105 — cone-level JIT inside the interpreter kernel.
5//!
6//! At assembly time, maximal cones of JIT-eligible nodes with
7//! scalar boundaries collapse into one synthetic `JitConeNode`
8//! each, compiled to native code via the existing P3 codegen. The
9//! cone node is an ordinary `PolydatNode`: the walker, scope
10//! chains, shared cells, None propagation, node_clean caching, and
11//! the enrich-and-re-raise panic contract all see a plain node.
12//!
13//! Boundary marshalling covers every one-slot immediate and every
14//! `Ref2` kind, borrowed into its pair for the call and copied out
15//! after it; interior fusion follows whatever the P3 classifier
16//! accepts. Extraction is recoverable:
17//! member nodes move into the cone only after codegen succeeds, so
18//! any JIT failure leaves the graph exactly as the interpreter
19//! would have compiled it.
20
21/// How much of the interpreter's graph is fused into native cones: the
22/// interpreter engine's one knob, carried by
23/// [`Engine::Interpreter`](crate::Engine::Interpreter) and settable per
24/// assembler with `set_jit_mode`. It is a property of the kernel being
25/// built, never of the process: two hosts in one process compiling
26/// under different modes get the kernels they each asked for.
27#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Hash)]
28pub enum JitMode {
29 /// Pure interpreter, no native code: the differential baseline.
30 Off,
31 /// Cone extraction with the cost model (fused cones of >= 2 nodes):
32 /// what a host gets when it names none.
33 #[default]
34 Auto,
35 /// Every eligible node joins a cone (threshold 1). Used by the
36 /// differential battery and for isolating marshalling regressions.
37 Force,
38}
39
40#[cfg(not(feature = "jit"))]
41pub(crate) fn extract_jit_cones(_dag: &mut super::assembly::ResolvedDag, _mode: JitMode) {}
42
43#[cfg(feature = "jit")]
44pub(crate) use jit_impl::extract_jit_cones;
45
46#[cfg(feature = "jit")]
47mod jit_impl {
48 use super::JitMode;
49 use crate::ast::{NodeMeta, PolydatNode, Port, PortType, Purity, Slot, SlotShape, Value};
50 use crate::compile::assembly::{PolydatAssembler, ResolvedDag};
51 use crate::compile::jit::{JitOp, classify_node_typed};
52 use crate::kernel::{InputDef, InputKind, WireSource};
53 use std::collections::HashMap;
54
55 /// A fused subgraph compiled to native code, standing in the
56 /// program as one ordinary node (SRD-105). The node is shared by
57 /// every state of the program; the slot buffer its native code
58 /// runs over, and the scratch entries its members' kits write
59 /// into, belong to the state that evaluates it, which hands them
60 /// in through [`PolydatNode::eval_in`] (axiom S3).
61 pub(crate) struct JitConeNode {
62 meta: NodeMeta,
63 code_fn: crate::compile::jit::NativeFn,
64 total_slots: usize,
65 /// The members' scratch entries, after the slot buffer in the
66 /// cone's scratch layout, with the validator's pairs.
67 scratch: crate::compile::jit::ScratchPlan,
68 /// Where each member lives, for the failure path (A7): the
69 /// member that failed is named as the program names it, with
70 /// its outputs under the program's names; the cone is no frame.
71 attribution: std::sync::Arc<crate::compile::Attribution>,
72 /// First buffer slot per boundary input, in port order.
73 in_slots: Vec<usize>,
74 /// Buffer slot per output port, in `meta.outs` order.
75 out_slots: Vec<usize>,
76 in_types: Vec<PortType>,
77 out_types: Vec<PortType>,
78 /// The original member nodes — kept alive for the LUT /
79 /// constant memory the native code references, and walked
80 /// by identity hashing (`fusion_subgraph`).
81 members: Vec<Box<dyn PolydatNode>>,
82 /// Local member wiring (`Input(i)` = this node's i-th
83 /// outer input; `NodeOutput(j, p)` = member j) — the
84 /// stored subgraph identity hashing recurses through.
85 sub_wiring: Vec<Vec<WireSource>>,
86 /// Per output port: (local member index, member port).
87 out_ports: Vec<(usize, usize)>,
88 /// The finalized code and the kits it calls, kept alive for
89 /// the life of the program.
90 _module: crate::compile::jit::JitCode,
91 /// Whether the code calls a helper, and so runs under the
92 /// catch; code with no call runs bare.
93 fallible: bool,
94 }
95
96 impl PolydatNode for JitConeNode {
97 fn meta(&self) -> &NodeMeta {
98 &self.meta
99 }
100
101 fn fusion_subgraph(&self) -> Option<crate::ast::FusionSubgraph<'_>> {
102 Some(crate::ast::FusionSubgraph {
103 members: &self.members,
104 wiring: &self.sub_wiring,
105 out_ports: &self.out_ports,
106 })
107 }
108
109 /// The state owns the cone's slot buffer and its members'
110 /// scratch entries (axiom S3): one `Slots` entry, then the
111 /// entries the members' kits declared, handed in at every
112 /// evaluation.
113 fn scratch_layout(&self) -> Vec<crate::ast::ScratchElem> {
114 let mut layout = vec![crate::ast::ScratchElem::Slots];
115 layout.extend(self.scratch.elems.iter().copied());
116 layout
117 }
118
119 fn eval_in(
120 &self,
121 scratch: &mut [crate::ast::ScratchBuf],
122 inputs: &[Value],
123 outputs: &mut [Value],
124 ) {
125 let (slots, members) = scratch.split_at_mut(1);
126 let crate::ast::ScratchBuf::Slots(buf) = &mut slots[0] else {
127 unreachable!("a cone's scratch is its slot buffer");
128 };
129 self.eval_with(buf, members, inputs, outputs)
130 }
131
132 /// An evaluation without a state's scratch (a node evaluated
133 /// on its own): a buffer and entries of the call's own.
134 fn eval(&self, inputs: &[Value], outputs: &mut [Value]) {
135 let mut buf = Vec::new();
136 let mut members: Vec<crate::ast::ScratchBuf> = self
137 .scratch
138 .elems
139 .iter()
140 .map(|e| crate::ast::ScratchBuf::new(*e))
141 .collect();
142 self.eval_with(&mut buf, &mut members, inputs, outputs)
143 }
144 }
145
146 impl JitConeNode {
147 /// Evaluate over `buf` and the members' scratch: the boundary
148 /// inputs are borrowed into their slots for the duration of the
149 /// call, the native code runs, and every output is copied out
150 /// as an owned `Value` (the interpreter never holds a reference
151 /// into a buffer).
152 fn eval_with(
153 &self,
154 buf: &mut Vec<u64>,
155 members: &mut [crate::ast::ScratchBuf],
156 inputs: &[Value],
157 outputs: &mut [Value],
158 ) {
159 buf.clear();
160 buf.resize(self.total_slots + 1, 0);
161 for (i, v) in inputs.iter().enumerate() {
162 let start = self.in_slots[i];
163 if crate::compile::marshal::encode_slots(v, self.in_types[i], &mut buf[start..])
164 .is_none()
165 {
166 panic!(
167 "cone `{}` boundary input [{i}] expected {:?}, got {:?}",
168 self.meta.name,
169 self.in_types[i],
170 v.port_type()
171 );
172 }
173 }
174 // Native code names the member it is in before each helper
175 // call (the slot past the layout); a failure is re-raised
176 // attributed to that member with the program's context and
177 // output names, and the interpreter re-raises it as is (A7).
178 let code_fn = self.code_fn;
179 let cp = buf.as_ptr();
180 let mp = buf.as_mut_ptr();
181 let sc = members.as_mut_ptr();
182 if !self.fallible {
183 // Code that calls no helper cannot fail: it runs bare.
184 unsafe { (code_fn)(cp, mp, sc) };
185 } else {
186 buf[self.total_slots] = u64::MAX;
187 let capture = crate::kernel::engines::EvalPanicCaptureGuard::arm();
188 let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
189 crate::compile::jit::invoke_with_catch(move || unsafe {
190 (code_fn)(cp, mp, sc);
191 })
192 }));
193 drop(capture);
194 if let Err(payload) = outcome {
195 let step = buf[self.total_slots] as usize;
196 self.attribution.reraise(payload, step, buf, None);
197 }
198 }
199 #[cfg(debug_assertions)]
200 for &(slot, idx) in &self.scratch.refs {
201 let (p, l) = members[idx].ptr_len();
202 assert!(
203 buf[slot] == p && buf[slot + 1] == l,
204 "S9 ref-validator: cone `{}` slot pair ({slot}, {}) does not name \
205 scratch[{idx}]",
206 self.meta.name,
207 slot + 1
208 );
209 }
210 for (k, slot) in self.out_slots.iter().enumerate() {
211 outputs[k] = crate::compile::marshal::decode_output(buf, *slot, self.out_types[k]);
212 }
213 }
214 }
215
216 /// A planned-but-rejected cone is diagnosable state, never
217 /// silent (audit channel, Debug level — rejections are normal
218 /// cost-model outcomes, not user-facing failures).
219 fn audit_skip(member_count: usize, reason: &str) {
220 crate::library::support::audit::debug(&format!(
221 "jit cone: leaving a {member_count}-member component on the interpreter: {reason}"
222 ));
223 }
224
225 /// Marshalable boundary types: every one-slot immediate, encoded
226 /// as the bits its `Wire` impl injects (a signed narrow carrier
227 /// sign-extended, an unsigned or float one as its bits;
228 /// type_system_alignment.md §2), and every `Ref2` kind, borrowed
229 /// into its pair for the call and copied out after it
230 /// (compiled_handles.md §4). The 128-bit immediates stay out until
231 /// they have a boundary encoding of their own.
232 fn scalar_ok(ty: PortType) -> bool {
233 use crate::ast::SlotColor;
234 match ty.slot_color() {
235 SlotColor::Imm1 | SlotColor::Ref2 => true,
236 SlotColor::Imm2 => false,
237 }
238 }
239
240 /// A node may join a cone iff the P3 classifier can lower it with
241 /// its wire types known, it is pure, and every wire port is a
242 /// single-slot value this push can marshal. The SRD-74 None rule
243 /// is applied by the caller, which knows where each input comes
244 /// from.
245 fn node_eligible(node: &dyn PolydatNode, wire_types: &[PortType]) -> bool {
246 matches!(node.purity(), Purity::Pure)
247 && !matches!(classify_node_typed(node, wire_types), JitOp::Fallback)
248 && node.meta().outs.iter().all(|p| scalar_ok(p.typ))
249 && wire_types.iter().all(|t| scalar_ok(*t))
250 && node.meta().wire_inputs().iter().all(|p| scalar_ok(p.typ))
251 }
252
253 /// SRD 11's three evaluation lifecycles, read from the one
254 /// classifier the program carries, so that extraction can
255 /// restrict fusion to per-cycle work. Const and scope-init
256 /// subgraphs belong to the fold passes (which evaluate them
257 /// once); fusing them would demote them to per-pull native
258 /// evaluation and — for multi-output cones — block
259 /// `fold_init_constants`' single-output replacement, breaking
260 /// `get_constant` consumers like `eval_const_expr`.
261 ///
262 /// This was a second copy of the walk, which had drifted: it
263 /// seeded from the inputs and the declared purity but knew
264 /// nothing of the `volatile` output modifier and did not
265 /// propagate volatility downstream, so a node a program declared
266 /// volatile could read here as const and be fused into a cone the
267 /// fold then evaluated once.
268 ///
269 /// Returns each node's lifecycle and whether it is volatile.
270 fn classify_lifecycles(dag: &ResolvedDag) -> (Vec<crate::kernel::EvalLifecycle>, Vec<bool>) {
271 let classes = crate::kernel::PolydatProgram::classify_lifecycle(
272 &dag.nodes,
273 &dag.wiring,
274 &dag.input_defs,
275 &dag.output_map,
276 &dag.output_modifiers,
277 );
278 (classes.lifecycle, classes.nondeterministic)
279 }
280
281 /// Dedup/lookup key for a boundary wire source.
282 fn src_key(src: &WireSource) -> (u8, usize, usize) {
283 match src {
284 WireSource::Input(i) => (0, *i, 0),
285 WireSource::NodeOutput(j, p) => (1, *j, *p),
286 }
287 }
288
289 struct ConePlan {
290 /// Member node indices, ascending (inherits topo order).
291 members: Vec<usize>,
292 /// Boundary input sources, deduped, in first-use order.
293 boundary_in: Vec<WireSource>,
294 in_types: Vec<PortType>,
295 /// Boundary output ports `(member_idx, port)`, first-use order.
296 boundary_out: Vec<(usize, usize)>,
297 out_types: Vec<PortType>,
298 }
299
300 /// Replace eligible cones in `dag` with compiled cone nodes.
301 /// On any per-cone failure the cone's members stay interpreter
302 /// nodes; the DAG is always left valid and topologically sorted.
303 pub(crate) fn extract_jit_cones(dag: &mut ResolvedDag, mode: JitMode) {
304 let min_members = match mode {
305 JitMode::Off => return,
306 JitMode::Auto => 2,
307 JitMode::Force => 1,
308 };
309 let n = dag.nodes.len();
310 if n == 0 {
311 return;
312 }
313
314 let (lifecycles, volatile) = classify_lifecycles(dag);
315 // Eligibility in topological order, because the SRD-74 None
316 // rule for a None-tolerant node depends on its sources: the
317 // kernel guard makes a fused cone None whenever a boundary
318 // input is None, so a node that would have seen the None and
319 // produced a value (`tile_encode` writes `null`, `to_json`
320 // keeps going) may join only when every input is an intra-cone
321 // wire from an eligible node, where no None can arrive. Every
322 // other node is guarded the same way fused or not.
323 let mut eligible: Vec<bool> = vec![false; n];
324 for i in 0..n {
325 if lifecycles[i] != crate::kernel::EvalLifecycle::Dynamic {
326 continue;
327 }
328 let nd = dag.nodes[i].as_ref();
329 if !node_eligible(nd, &crate::compile::assembly::wire_types_of(dag, i)) {
330 continue;
331 }
332 if !crate::compile::none_rule_admits(
333 nd.accepts_none_inputs(),
334 &dag.wiring[i],
335 &eligible,
336 ) {
337 continue;
338 }
339 eligible[i] = true;
340 }
341
342 // Connected components over eligible-to-eligible wires, by the
343 // rule every fusing engine shares (compile::fusion_units). A
344 // volatile node never shares a cone with a node that is not: a
345 // cone runs whole, so every read that re-evaluates the volatile
346 // node would re-run the cached work upstream of it too.
347 let preds: Vec<Vec<usize>> = dag
348 .wiring
349 .iter()
350 .map(|w| {
351 w.iter()
352 .filter_map(|src| match src {
353 WireSource::NodeOutput(j, _) => Some(*j),
354 WireSource::Input(_) => None,
355 })
356 .collect()
357 })
358 .collect();
359 let class: Vec<u64> = volatile.iter().map(|&v| v as u64).collect();
360 let components = crate::compile::fusion_units::components(&preds, &eligible, &class);
361
362 // Consumer adjacency over the ORIGINAL node graph — the
363 // convexity walk below routes through it.
364 let mut consumers: Vec<Vec<usize>> = vec![Vec::new(); n];
365 for (i, ps) in preds.iter().enumerate() {
366 for &j in ps {
367 consumers[j].push(i);
368 }
369 }
370
371 let mut nodes_opt: Vec<Option<Box<dyn PolydatNode>>> = std::mem::take(&mut dag.nodes)
372 .into_iter()
373 .map(Some)
374 .collect();
375 let mut cones: Vec<(ConePlan, JitConeNode)> = Vec::new();
376
377 for members in &components {
378 if members.len() < min_members {
379 continue;
380 }
381 // Connected components are not necessarily CONVEX: an
382 // eligible→ineligible→eligible sandwich whose ends
383 // connect through some other eligible path lands both
384 // ends in one component while the middle stays kept.
385 // Fusing that component makes the kept middle both a
386 // consumer of the cone and one of its producers — a
387 // cycle in the spliced graph (the rebuild topo-sort
388 // assert). Detection: walk the consumer graph from the
389 // members' external consumers, only through
390 // non-members (other cones' members are ordinary route
391 // nodes here, which also covers cross-cone quotient
392 // cycles); reaching a member proves an external path
393 // re-enters this cone. Per the module's fallback rule,
394 // such a component stays on the interpreter.
395 if !crate::compile::fusion_units::is_convex(members, &consumers) {
396 audit_skip(
397 members.len(),
398 "non-convex component (an external path re-enters the cone)",
399 );
400 continue;
401 }
402 let Some(plan) = plan_cone(dag, members, &nodes_opt) else {
403 // plan_cone audit-logs its own rejection reason;
404 // the component stays on the interpreter.
405 continue;
406 };
407 match build_cone(dag, &plan, &mut nodes_opt) {
408 Ok(cone) => {
409 // Formation is diagnosable state too — the B2
410 // sweep and cone-aware bench reporting key on
411 // this line to verify extraction actually ran.
412 crate::library::support::audit::debug(&format!(
413 "jit cone: fused {} members ({} boundary in, {} out): {}",
414 plan.members.len(),
415 plan.boundary_in.len(),
416 plan.boundary_out.len(),
417 cone.meta().name,
418 ));
419 cones.push((plan, cone));
420 }
421 // Members were restored by build_cone; the cone
422 // stays on the interpreter (SRD-105 fallback rule:
423 // a JIT failure never fails a compile). Eligibility
424 // prescreens classification, so a codegen error
425 // here is unexpected — surface it.
426 Err(e) => {
427 crate::library::support::audit::warn(&format!(
428 "jit cone: codegen failed for a {}-member cone — staying on the interpreter: {e}",
429 plan.members.len(),
430 ));
431 }
432 }
433 }
434
435 if cones.is_empty() {
436 dag.nodes = nodes_opt.into_iter().map(Option::unwrap).collect();
437 return;
438 }
439 rebuild(dag, nodes_opt, cones);
440 }
441
442 /// Compute the cone's boundaries; `None` rejects the component
443 /// (dead outputs, oversized boundary, unmarshalable edge type).
444 fn plan_cone(
445 dag: &ResolvedDag,
446 members: &[usize],
447 nodes: &[Option<Box<dyn PolydatNode>>],
448 ) -> Option<ConePlan> {
449 let is_member = |j: usize| members.binary_search(&j).is_ok();
450
451 let mut boundary_in: Vec<WireSource> = Vec::new();
452 let mut in_types: Vec<PortType> = Vec::new();
453 let mut seen_in: HashMap<(u8, usize, usize), usize> = HashMap::new();
454 for &m in members {
455 let member = nodes[m].as_ref()?;
456 let member_ports: Vec<PortType> =
457 member.meta().wire_inputs().iter().map(|p| p.typ).collect();
458 let wire_types: Vec<PortType> = dag.wiring[m]
459 .iter()
460 .map(|src| match src {
461 WireSource::Input(i) => Some(dag.input_defs[*i].port_type),
462 WireSource::NodeOutput(j, p) => Some(nodes[*j].as_ref()?.meta().outs[*p].typ),
463 })
464 .collect::<Option<_>>()?;
465 // A node that lowers as a slot call runs the kit built for
466 // its wire types (compiled_handles.md §6), so its advertised
467 // port types do not bind its wires: a variadic that inspects
468 // `Value`s at P1 reads each wire as the wire is. A named
469 // native lowering takes its ports as declared.
470 let typed_by_wires = matches!(
471 classify_node_typed(member.as_ref(), &wire_types),
472 JitOp::SlotCall { .. }
473 );
474 for (k, src) in dag.wiring[m].iter().enumerate() {
475 let ty = wire_types[k];
476 // Inside a cone every wire is exactly its port's type.
477 if !typed_by_wires
478 && let Some(expected) = member_ports.get(k)
479 && *expected != ty
480 {
481 audit_skip(
482 members.len(),
483 &format!(
484 "input [{k}] of `{}` is a {ty:?} wire on a {expected:?} port",
485 member.meta().name
486 ),
487 );
488 return None;
489 }
490 let intra = matches!(src, WireSource::NodeOutput(j, _) if is_member(*j));
491 // SRD-74: a None-tolerant member must not sit on the
492 // boundary, where a None could reach it (see the
493 // eligibility pass); a component split can put it there.
494 if !intra && member.accepts_none_inputs() {
495 audit_skip(
496 members.len(),
497 &format!(
498 "`{}` tolerates None inputs and input [{k}] is a boundary wire",
499 member.meta().name
500 ),
501 );
502 return None;
503 }
504 if intra {
505 continue;
506 }
507 let key = src_key(src);
508 if seen_in.contains_key(&key) {
509 continue;
510 }
511 if !scalar_ok(ty) {
512 audit_skip(
513 members.len(),
514 &format!("boundary input of type {ty:?} is not marshalable"),
515 );
516 return None;
517 }
518 seen_in.insert(key, boundary_in.len());
519 boundary_in.push(src.clone());
520 in_types.push(ty);
521 }
522 }
523 // SRD-105: cones are bounded at 64 boundary inputs. The bound
524 // is a size cap on a cone's boundary, kept from when a
525 // provenance mask was one word (`ProvMask` is now multi-word);
526 // a cone over it is skipped rather than re-split.
527 if boundary_in.len() > 64 {
528 audit_skip(
529 members.len(),
530 &format!(
531 "{} boundary inputs exceeds the 64-input bound (no re-split implemented — catchup item B2)",
532 boundary_in.len()
533 ),
534 );
535 return None;
536 }
537 // A cone with no boundary inputs is a compile-time
538 // constant: it would evaluate exactly once (node_clean)
539 // and belongs to const folding, not per-cycle fusion.
540 // It also breaks lifecycle analysis (a no-input node
541 // claiming per-cycle outputs). Leave it interpreted.
542 if boundary_in.is_empty() {
543 // Normal outcome for const subgraphs — the fold passes
544 // own them; not worth an audit line.
545 return None;
546 }
547
548 let mut boundary_out: Vec<(usize, usize)> = Vec::new();
549 let mut seen_out: HashMap<(usize, usize), usize> = HashMap::new();
550 let mut note_out = |j: usize, p: usize| {
551 if let std::collections::hash_map::Entry::Vacant(e) = seen_out.entry((j, p)) {
552 e.insert(boundary_out.len());
553 boundary_out.push((j, p));
554 }
555 };
556 for (i, wiring) in dag.wiring.iter().enumerate() {
557 if is_member(i) {
558 continue;
559 }
560 for src in wiring {
561 if let WireSource::NodeOutput(j, p) = src
562 && is_member(*j)
563 {
564 note_out(*j, *p);
565 }
566 }
567 }
568 for (j, p) in dag.output_map.values() {
569 if is_member(*j) {
570 note_out(*j, *p);
571 }
572 }
573 if boundary_out.is_empty() {
574 // Dead subgraph (no observable outputs) — DCE
575 // territory, not worth an audit line.
576 return None;
577 }
578 let out_types: Vec<PortType> = boundary_out
579 .iter()
580 .map(|(j, p)| nodes[*j].as_ref().map(|nd| nd.meta().outs[*p].typ))
581 .collect::<Option<_>>()?;
582 if out_types.iter().any(|t| !scalar_ok(*t)) {
583 audit_skip(members.len(), "a boundary output type is not marshalable");
584 return None;
585 }
586
587 Some(ConePlan {
588 members: members.to_vec(),
589 boundary_in,
590 in_types,
591 boundary_out,
592 out_types,
593 })
594 }
595
596 /// A boundary input's declared default, of its own type; the cone
597 /// is always evaluated with its inputs bound, so the default is
598 /// never read, but the definition is typed like any input's.
599 fn default_for(ty: PortType) -> Value {
600 match ty {
601 PortType::F64 => Value::F64(0.0),
602 PortType::Bool => Value::Bool(false),
603 PortType::Str => Value::Str("".into()),
604 PortType::Bytes => Value::Bytes(Vec::new().into()),
605 PortType::Json => Value::Json(std::sync::Arc::new(serde_json::Value::Null)),
606 PortType::U64 => Value::U64(0),
607 _ => Value::None,
608 }
609 }
610
611 /// Attempt native compilation of the planned cone. Codegen runs
612 /// before the members leave the graph permanently: on any error
613 /// they are restored and the caller keeps the interpreter form.
614 fn build_cone(
615 dag: &ResolvedDag,
616 plan: &ConePlan,
617 nodes: &mut [Option<Box<dyn PolydatNode>>],
618 ) -> Result<JitConeNode, String> {
619 let local: HashMap<usize, usize> = plan
620 .members
621 .iter()
622 .enumerate()
623 .map(|(l, &g)| (g, l))
624 .collect();
625 let in_pos: HashMap<(u8, usize, usize), usize> = plan
626 .boundary_in
627 .iter()
628 .enumerate()
629 .map(|(i, s)| (src_key(s), i))
630 .collect();
631
632 let sub_wiring: Vec<Vec<WireSource>> = plan
633 .members
634 .iter()
635 .map(|&m| {
636 dag.wiring[m]
637 .iter()
638 .map(|src| match src {
639 WireSource::NodeOutput(j, p) if local.contains_key(j) => {
640 WireSource::NodeOutput(local[j], *p)
641 }
642 other => WireSource::Input(in_pos[&src_key(other)]),
643 })
644 .collect()
645 })
646 .collect();
647 let sub_input_defs: Vec<InputDef> = plan
648 .in_types
649 .iter()
650 .enumerate()
651 .map(|(i, ty)| InputDef {
652 name: format!("c{i}"),
653 default: default_for(*ty),
654 port_type: *ty,
655 kind: InputKind::Coordinate,
656 type_origin: crate::kernel::TypeOrigin::Declared,
657 converts_to: None,
658 })
659 .collect();
660 let mut sub_output_map: HashMap<String, (usize, usize)> = HashMap::new();
661 let mut sub_output_order: Vec<String> = Vec::new();
662 for (k, (j, p)) in plan.boundary_out.iter().enumerate() {
663 let name = format!("o{k}");
664 sub_output_map.insert(name.clone(), (local[j], *p));
665 sub_output_order.push(name);
666 }
667
668 let taken: Vec<Box<dyn PolydatNode>> = plan
669 .members
670 .iter()
671 .map(|&m| nodes[m].take().expect("cone member present"))
672 .collect();
673 let member_label = cone_label(&taken);
674
675 let mut sub = ResolvedDag {
676 nodes: taken,
677 wiring: sub_wiring,
678 input_defs: sub_input_defs,
679 coord_count: plan.boundary_in.len(),
680 output_map: sub_output_map,
681 output_order: sub_output_order,
682 cursor_schemas: Vec::new(),
683 source: String::new(),
684 // A member's failure is reported against the program the
685 // cone stands in, as the same node's failure is reported on
686 // every other engine (A7); the cone is not a frame of its own.
687 context: dag.context.clone(),
688 output_modifiers: HashMap::new(),
689 const_outputs: std::collections::HashSet::new(),
690 const_inits: Vec::new(),
691 // A cone is a fragment of the program that stands in the
692 // tree's ledger already, not a program of its own: its
693 // kernel is recorded nowhere.
694 ledger: crate::kernel::CompileLedger::new(),
695 };
696
697 let restore = |sub_nodes: Vec<Box<dyn PolydatNode>>,
698 nodes: &mut [Option<Box<dyn PolydatNode>>]| {
699 for (&m, nd) in plan.members.iter().zip(sub_nodes) {
700 nodes[m] = Some(nd);
701 }
702 };
703
704 let layout = match PolydatAssembler::build_jit_layout(&sub) {
705 Ok(l) => l,
706 Err(e) => {
707 restore(sub.nodes, nodes);
708 return Err(e);
709 }
710 };
711 let (coord_slots, total_slots, jit_steps, jit_outputs, scratch, _volatile) = layout;
712 // Boundary inputs occupy the first slots, each as wide as its
713 // type.
714 let mut in_slots = Vec::with_capacity(plan.in_types.len());
715 let mut next = 0usize;
716 for ty in &plan.in_types {
717 in_slots.push(next);
718 next += ty.slot_width();
719 }
720 debug_assert_eq!(coord_slots, next);
721 let compiled = crate::compile::jit::compile_jit_entry(&jit_steps, Some(total_slots));
722 let (code_fn, code) = match compiled {
723 Ok(parts) => parts,
724 Err(e) => {
725 restore(sub.nodes, nodes);
726 return Err(e);
727 }
728 };
729
730 let out_slots: Vec<usize> = (0..plan.boundary_out.len())
731 .map(|k| jit_outputs[&format!("o{k}")])
732 .collect();
733 // Port metadata mirrors the fused subgraph rather than
734 // being synthesized: outputs clone the member's original
735 // port (lifecycle analysis and downstream diagnostics see
736 // what the interpreter form would have declared); inputs
737 // clone the source port where one exists (graph inputs are
738 // per-cycle by definition).
739 let meta = NodeMeta {
740 name: member_label,
741 ins: plan
742 .boundary_in
743 .iter()
744 .zip(&plan.in_types)
745 .enumerate()
746 .map(|(i, (src, ty))| {
747 // Boundary producers are ineligible nodes by
748 // definition, so they are never cone members
749 // and always present in the slot vec.
750 let mut port = match src {
751 WireSource::NodeOutput(j, p) => nodes[*j]
752 .as_ref()
753 .map(|nd| nd.meta().outs[*p].clone())
754 .unwrap_or_else(|| Port::new("", *ty)),
755 WireSource::Input(_) => Port::new("", *ty),
756 };
757 port.name = format!("c{i}");
758 port.constraint = None;
759 Slot::Wire(port)
760 })
761 .collect(),
762 outs: plan
763 .boundary_out
764 .iter()
765 .enumerate()
766 .map(|(k, (j, p))| {
767 let mut port = sub.nodes[local[j]].meta().outs[*p].clone();
768 port.name = format!("o{k}");
769 port.constraint = None;
770 port
771 })
772 .collect(),
773 };
774 let out_ports: Vec<(usize, usize)> = plan
775 .boundary_out
776 .iter()
777 .map(|(j, p)| (local[j], *p))
778 .collect();
779 // A member's failure names the member's outputs as the program
780 // names them (A7), not as the cone numbers them: the boundary
781 // outputs take the program's names for the attribution.
782 let mut named = sub.output_map.clone();
783 for (k, (j, p)) in plan.boundary_out.iter().enumerate() {
784 let names: Vec<String> = dag
785 .output_map
786 .iter()
787 .filter(|(_, v)| **v == (*j, *p))
788 .map(|(n, _)| n.clone())
789 .collect();
790 if !names.is_empty()
791 && let Some(target) = named.remove(&format!("o{k}"))
792 {
793 for n in names {
794 named.insert(n, target);
795 }
796 }
797 }
798 let numbered = std::mem::replace(&mut sub.output_map, named);
799 let attribution = std::sync::Arc::new(PolydatAssembler::attribution_of(&sub));
800 sub.output_map = numbered;
801 Ok(JitConeNode {
802 attribution,
803 in_slots,
804 meta,
805 code_fn,
806 total_slots,
807 out_slots,
808 in_types: plan.in_types.clone(),
809 out_types: plan.out_types.clone(),
810 members: sub.nodes,
811 sub_wiring: sub.wiring,
812 out_ports,
813 scratch,
814 fallible: code.fallible(),
815 _module: code,
816 })
817 }
818
819 /// Diagnostic name carrying the fused members, so an enriched
820 /// eval panic attributes the interior functions.
821 fn cone_label(members: &[Box<dyn PolydatNode>]) -> String {
822 const SHOWN: usize = 6;
823 let names: Vec<&str> = members
824 .iter()
825 .take(SHOWN)
826 .map(|n| n.meta().name.as_str())
827 .collect();
828 let suffix = if members.len() > SHOWN {
829 format!("+{} more", members.len() - SHOWN)
830 } else {
831 String::new()
832 };
833 format!("jit_cone[{}{}]", names.join("+"), suffix)
834 }
835
836 /// Splice the compiled cones into the DAG and restore
837 /// topological order.
838 fn rebuild(
839 dag: &mut ResolvedDag,
840 nodes_opt: Vec<Option<Box<dyn PolydatNode>>>,
841 cones: Vec<(ConePlan, JitConeNode)>,
842 ) {
843 let old_n = nodes_opt.len();
844 // (old_idx, port) → (cone_ordinal, cone_out_port)
845 let mut cone_port: HashMap<(usize, usize), (usize, usize)> = HashMap::new();
846 for (ci, (plan, _)) in cones.iter().enumerate() {
847 for (k, (j, p)) in plan.boundary_out.iter().enumerate() {
848 cone_port.insert((*j, *p), (ci, k));
849 }
850 }
851
852 let mut kept_map: HashMap<usize, usize> = HashMap::new();
853 let mut new_nodes: Vec<Box<dyn PolydatNode>> = Vec::new();
854 let mut new_wiring: Vec<Vec<WireSource>> = Vec::new();
855 for (old, slot) in nodes_opt.into_iter().enumerate() {
856 if let Some(node) = slot {
857 kept_map.insert(old, new_nodes.len());
858 new_nodes.push(node);
859 new_wiring.push(dag.wiring[old].clone());
860 }
861 }
862 let cone_base = new_nodes.len();
863 let mut cone_plans: Vec<ConePlan> = Vec::with_capacity(cones.len());
864 for (plan, cone) in cones {
865 new_nodes.push(Box::new(cone));
866 new_wiring.push(plan.boundary_in.clone());
867 cone_plans.push(plan);
868 }
869
870 let remap = |src: &WireSource| -> WireSource {
871 match src {
872 WireSource::Input(i) => WireSource::Input(*i),
873 WireSource::NodeOutput(j, p) => {
874 if let Some(&nj) = kept_map.get(j) {
875 WireSource::NodeOutput(nj, *p)
876 } else {
877 let (ci, k) = cone_port[&(*j, *p)];
878 WireSource::NodeOutput(cone_base + ci, k)
879 }
880 }
881 }
882 };
883 for wiring in new_wiring.iter_mut() {
884 for src in wiring.iter_mut() {
885 *src = remap(src);
886 }
887 }
888 let mut new_output_map: HashMap<String, (usize, usize)> = HashMap::new();
889 for (name, (j, p)) in dag.output_map.iter() {
890 let (nj, np) = match remap(&WireSource::NodeOutput(*j, *p)) {
891 WireSource::NodeOutput(a, b) => (a, b),
892 WireSource::Input(_) => unreachable!("outputs map to nodes"),
893 };
894 new_output_map.insert(name.clone(), (nj, np));
895 }
896
897 // Kahn topo sort — consumers of cone interiors may sit at
898 // indices below the spliced cone node.
899 let m = new_nodes.len();
900 let mut indegree = vec![0usize; m];
901 let mut dependents: Vec<Vec<usize>> = vec![Vec::new(); m];
902 for (i, wiring) in new_wiring.iter().enumerate() {
903 let mut producers: Vec<usize> = wiring
904 .iter()
905 .filter_map(|s| match s {
906 WireSource::NodeOutput(j, _) => Some(*j),
907 WireSource::Input(_) => None,
908 })
909 .collect();
910 producers.sort_unstable();
911 producers.dedup();
912 indegree[i] = producers.len();
913 for j in producers {
914 dependents[j].push(i);
915 }
916 }
917 let mut order: Vec<usize> = Vec::with_capacity(m);
918 let mut ready: std::collections::BinaryHeap<std::cmp::Reverse<usize>> = (0..m)
919 .filter(|&i| indegree[i] == 0)
920 .map(std::cmp::Reverse)
921 .collect();
922 while let Some(std::cmp::Reverse(i)) = ready.pop() {
923 order.push(i);
924 for &d in &dependents[i] {
925 indegree[d] -= 1;
926 if indegree[d] == 0 {
927 ready.push(std::cmp::Reverse(d));
928 }
929 }
930 }
931 assert_eq!(
932 order.len(),
933 m,
934 "cone splice must not introduce a cycle (old_n={old_n})"
935 );
936 let mut pos = vec![0usize; m];
937 for (new_idx, &i) in order.iter().enumerate() {
938 pos[i] = new_idx;
939 }
940
941 let mut sorted_nodes: Vec<Option<Box<dyn PolydatNode>>> =
942 new_nodes.into_iter().map(Some).collect();
943 dag.nodes = order
944 .iter()
945 .map(|&i| sorted_nodes[i].take().expect("each node placed once"))
946 .collect();
947 dag.wiring = order
948 .iter()
949 .map(|&i| {
950 new_wiring[i]
951 .iter()
952 .map(|s| match s {
953 WireSource::Input(k) => WireSource::Input(*k),
954 WireSource::NodeOutput(j, p) => WireSource::NodeOutput(pos[*j], *p),
955 })
956 .collect()
957 })
958 .collect();
959 dag.output_map = new_output_map
960 .into_iter()
961 .map(|(name, (j, p))| (name, (pos[j], p)))
962 .collect();
963 let _ = cone_plans;
964 }
965}