Skip to main content

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    fn classify_lifecycles(dag: &ResolvedDag) -> Vec<crate::kernel::EvalLifecycle> {
269        crate::kernel::PolydatProgram::classify_lifecycle(
270            &dag.nodes,
271            &dag.wiring,
272            &dag.input_defs,
273            &dag.output_map,
274            &dag.output_modifiers,
275        )
276        .lifecycle
277    }
278
279    /// Dedup/lookup key for a boundary wire source.
280    fn src_key(src: &WireSource) -> (u8, usize, usize) {
281        match src {
282            WireSource::Input(i) => (0, *i, 0),
283            WireSource::NodeOutput(j, p) => (1, *j, *p),
284        }
285    }
286
287    struct ConePlan {
288        /// Member node indices, ascending (inherits topo order).
289        members: Vec<usize>,
290        /// Boundary input sources, deduped, in first-use order.
291        boundary_in: Vec<WireSource>,
292        in_types: Vec<PortType>,
293        /// Boundary output ports `(member_idx, port)`, first-use order.
294        boundary_out: Vec<(usize, usize)>,
295        out_types: Vec<PortType>,
296    }
297
298    /// Replace eligible cones in `dag` with compiled cone nodes.
299    /// On any per-cone failure the cone's members stay interpreter
300    /// nodes; the DAG is always left valid and topologically sorted.
301    pub(crate) fn extract_jit_cones(dag: &mut ResolvedDag, mode: JitMode) {
302        let min_members = match mode {
303            JitMode::Off => return,
304            JitMode::Auto => 2,
305            JitMode::Force => 1,
306        };
307        let n = dag.nodes.len();
308        if n == 0 {
309            return;
310        }
311
312        let lifecycles = classify_lifecycles(dag);
313        // Eligibility in topological order, because the SRD-74 None
314        // rule for a None-tolerant node depends on its sources: the
315        // kernel guard makes a fused cone None whenever a boundary
316        // input is None, so a node that would have seen the None and
317        // produced a value (`tile_encode` writes `null`, `to_json`
318        // keeps going) may join only when every input is an intra-cone
319        // wire from an eligible node, where no None can arrive. Every
320        // other node is guarded the same way fused or not.
321        let mut eligible: Vec<bool> = vec![false; n];
322        for i in 0..n {
323            if lifecycles[i] != crate::kernel::EvalLifecycle::Dynamic {
324                continue;
325            }
326            let nd = dag.nodes[i].as_ref();
327            if !node_eligible(nd, &crate::compile::assembly::wire_types_of(dag, i)) {
328                continue;
329            }
330            if !crate::compile::none_rule_admits(
331                nd.accepts_none_inputs(),
332                &dag.wiring[i],
333                &eligible,
334            ) {
335                continue;
336            }
337            eligible[i] = true;
338        }
339
340        // Connected components over eligible-to-eligible wires, by the
341        // rule every fusing engine shares (compile::fusion_units).
342        let preds: Vec<Vec<usize>> = dag
343            .wiring
344            .iter()
345            .map(|w| {
346                w.iter()
347                    .filter_map(|src| match src {
348                        WireSource::NodeOutput(j, _) => Some(*j),
349                        WireSource::Input(_) => None,
350                    })
351                    .collect()
352            })
353            .collect();
354        let components = crate::compile::fusion_units::components(&preds, &eligible, &vec![0; n]);
355
356        // Consumer adjacency over the ORIGINAL node graph — the
357        // convexity walk below routes through it.
358        let mut consumers: Vec<Vec<usize>> = vec![Vec::new(); n];
359        for (i, ps) in preds.iter().enumerate() {
360            for &j in ps {
361                consumers[j].push(i);
362            }
363        }
364
365        let mut nodes_opt: Vec<Option<Box<dyn PolydatNode>>> = std::mem::take(&mut dag.nodes)
366            .into_iter()
367            .map(Some)
368            .collect();
369        let mut cones: Vec<(ConePlan, JitConeNode)> = Vec::new();
370
371        for members in &components {
372            if members.len() < min_members {
373                continue;
374            }
375            // Connected components are not necessarily CONVEX: an
376            // eligible→ineligible→eligible sandwich whose ends
377            // connect through some other eligible path lands both
378            // ends in one component while the middle stays kept.
379            // Fusing that component makes the kept middle both a
380            // consumer of the cone and one of its producers — a
381            // cycle in the spliced graph (the rebuild topo-sort
382            // assert). Detection: walk the consumer graph from the
383            // members' external consumers, only through
384            // non-members (other cones' members are ordinary route
385            // nodes here, which also covers cross-cone quotient
386            // cycles); reaching a member proves an external path
387            // re-enters this cone. Per the module's fallback rule,
388            // such a component stays on the interpreter.
389            if !crate::compile::fusion_units::is_convex(members, &consumers) {
390                audit_skip(
391                    members.len(),
392                    "non-convex component (an external path re-enters the cone)",
393                );
394                continue;
395            }
396            let Some(plan) = plan_cone(dag, members, &nodes_opt) else {
397                // plan_cone audit-logs its own rejection reason;
398                // the component stays on the interpreter.
399                continue;
400            };
401            match build_cone(dag, &plan, &mut nodes_opt) {
402                Ok(cone) => {
403                    // Formation is diagnosable state too — the B2
404                    // sweep and cone-aware bench reporting key on
405                    // this line to verify extraction actually ran.
406                    crate::library::support::audit::debug(&format!(
407                        "jit cone: fused {} members ({} boundary in, {} out): {}",
408                        plan.members.len(),
409                        plan.boundary_in.len(),
410                        plan.boundary_out.len(),
411                        cone.meta().name,
412                    ));
413                    cones.push((plan, cone));
414                }
415                // Members were restored by build_cone; the cone
416                // stays on the interpreter (SRD-105 fallback rule:
417                // a JIT failure never fails a compile). Eligibility
418                // prescreens classification, so a codegen error
419                // here is unexpected — surface it.
420                Err(e) => {
421                    crate::library::support::audit::warn(&format!(
422                        "jit cone: codegen failed for a {}-member                          cone — staying on the interpreter: {e}",
423                        plan.members.len(),
424                    ));
425                }
426            }
427        }
428
429        if cones.is_empty() {
430            dag.nodes = nodes_opt.into_iter().map(Option::unwrap).collect();
431            return;
432        }
433        rebuild(dag, nodes_opt, cones);
434    }
435
436    /// Compute the cone's boundaries; `None` rejects the component
437    /// (dead outputs, oversized boundary, unmarshalable edge type).
438    fn plan_cone(
439        dag: &ResolvedDag,
440        members: &[usize],
441        nodes: &[Option<Box<dyn PolydatNode>>],
442    ) -> Option<ConePlan> {
443        let is_member = |j: usize| members.binary_search(&j).is_ok();
444
445        let mut boundary_in: Vec<WireSource> = Vec::new();
446        let mut in_types: Vec<PortType> = Vec::new();
447        let mut seen_in: HashMap<(u8, usize, usize), usize> = HashMap::new();
448        for &m in members {
449            let member = nodes[m].as_ref()?;
450            let member_ports: Vec<PortType> =
451                member.meta().wire_inputs().iter().map(|p| p.typ).collect();
452            let wire_types: Vec<PortType> = dag.wiring[m]
453                .iter()
454                .map(|src| match src {
455                    WireSource::Input(i) => Some(dag.input_defs[*i].port_type),
456                    WireSource::NodeOutput(j, p) => Some(nodes[*j].as_ref()?.meta().outs[*p].typ),
457                })
458                .collect::<Option<_>>()?;
459            // A node that lowers as a slot call runs the kit built for
460            // its wire types (compiled_handles.md §6), so its advertised
461            // port types do not bind its wires: a variadic that inspects
462            // `Value`s at P1 reads each wire as the wire is. A named
463            // native lowering takes its ports as declared.
464            let typed_by_wires = matches!(
465                classify_node_typed(member.as_ref(), &wire_types),
466                JitOp::SlotCall { .. }
467            );
468            for (k, src) in dag.wiring[m].iter().enumerate() {
469                let ty = wire_types[k];
470                // Inside a cone every wire is exactly its port's type.
471                if !typed_by_wires
472                    && let Some(expected) = member_ports.get(k)
473                    && *expected != ty
474                {
475                    audit_skip(
476                        members.len(),
477                        &format!(
478                            "input [{k}] of `{}` is a {ty:?} wire on a {expected:?} port",
479                            member.meta().name
480                        ),
481                    );
482                    return None;
483                }
484                let intra = matches!(src, WireSource::NodeOutput(j, _) if is_member(*j));
485                // SRD-74: a None-tolerant member must not sit on the
486                // boundary, where a None could reach it (see the
487                // eligibility pass); a component split can put it there.
488                if !intra && member.accepts_none_inputs() {
489                    audit_skip(
490                        members.len(),
491                        &format!(
492                            "`{}` tolerates None inputs and input [{k}] is a boundary wire",
493                            member.meta().name
494                        ),
495                    );
496                    return None;
497                }
498                if intra {
499                    continue;
500                }
501                let key = src_key(src);
502                if seen_in.contains_key(&key) {
503                    continue;
504                }
505                if !scalar_ok(ty) {
506                    audit_skip(
507                        members.len(),
508                        &format!("boundary input of type {ty:?} is not marshalable"),
509                    );
510                    return None;
511                }
512                seen_in.insert(key, boundary_in.len());
513                boundary_in.push(src.clone());
514                in_types.push(ty);
515            }
516        }
517        // SRD-105: cones are bounded at 64 boundary inputs. The bound
518        // is a size cap on a cone's boundary, kept from when a
519        // provenance mask was one word (`ProvMask` is now multi-word);
520        // a cone over it is skipped rather than re-split.
521        if boundary_in.len() > 64 {
522            audit_skip(
523                members.len(),
524                &format!(
525                    "{} boundary inputs exceeds the 64-input bound (no                  re-split implemented — catchup item B2)",
526                    boundary_in.len()
527                ),
528            );
529            return None;
530        }
531        // A cone with no boundary inputs is a compile-time
532        // constant: it would evaluate exactly once (node_clean)
533        // and belongs to const folding, not per-cycle fusion.
534        // It also breaks lifecycle analysis (a no-input node
535        // claiming per-cycle outputs). Leave it interpreted.
536        if boundary_in.is_empty() {
537            // Normal outcome for const subgraphs — the fold passes
538            // own them; not worth an audit line.
539            return None;
540        }
541
542        let mut boundary_out: Vec<(usize, usize)> = Vec::new();
543        let mut seen_out: HashMap<(usize, usize), usize> = HashMap::new();
544        let mut note_out = |j: usize, p: usize| {
545            if let std::collections::hash_map::Entry::Vacant(e) = seen_out.entry((j, p)) {
546                e.insert(boundary_out.len());
547                boundary_out.push((j, p));
548            }
549        };
550        for (i, wiring) in dag.wiring.iter().enumerate() {
551            if is_member(i) {
552                continue;
553            }
554            for src in wiring {
555                if let WireSource::NodeOutput(j, p) = src
556                    && is_member(*j)
557                {
558                    note_out(*j, *p);
559                }
560            }
561        }
562        for (j, p) in dag.output_map.values() {
563            if is_member(*j) {
564                note_out(*j, *p);
565            }
566        }
567        if boundary_out.is_empty() {
568            // Dead subgraph (no observable outputs) — DCE
569            // territory, not worth an audit line.
570            return None;
571        }
572        let out_types: Vec<PortType> = boundary_out
573            .iter()
574            .map(|(j, p)| nodes[*j].as_ref().map(|nd| nd.meta().outs[*p].typ))
575            .collect::<Option<_>>()?;
576        if out_types.iter().any(|t| !scalar_ok(*t)) {
577            audit_skip(members.len(), "a boundary output type is not marshalable");
578            return None;
579        }
580
581        Some(ConePlan {
582            members: members.to_vec(),
583            boundary_in,
584            in_types,
585            boundary_out,
586            out_types,
587        })
588    }
589
590    /// A boundary input's declared default, of its own type; the cone
591    /// is always evaluated with its inputs bound, so the default is
592    /// never read, but the definition is typed like any input's.
593    fn default_for(ty: PortType) -> Value {
594        match ty {
595            PortType::F64 => Value::F64(0.0),
596            PortType::Bool => Value::Bool(false),
597            PortType::Str => Value::Str("".into()),
598            PortType::Bytes => Value::Bytes(Vec::new().into()),
599            PortType::Json => Value::Json(std::sync::Arc::new(serde_json::Value::Null)),
600            PortType::U64 => Value::U64(0),
601            _ => Value::None,
602        }
603    }
604
605    /// Attempt native compilation of the planned cone. Codegen runs
606    /// before the members leave the graph permanently: on any error
607    /// they are restored and the caller keeps the interpreter form.
608    fn build_cone(
609        dag: &ResolvedDag,
610        plan: &ConePlan,
611        nodes: &mut [Option<Box<dyn PolydatNode>>],
612    ) -> Result<JitConeNode, String> {
613        let local: HashMap<usize, usize> = plan
614            .members
615            .iter()
616            .enumerate()
617            .map(|(l, &g)| (g, l))
618            .collect();
619        let in_pos: HashMap<(u8, usize, usize), usize> = plan
620            .boundary_in
621            .iter()
622            .enumerate()
623            .map(|(i, s)| (src_key(s), i))
624            .collect();
625
626        let sub_wiring: Vec<Vec<WireSource>> = plan
627            .members
628            .iter()
629            .map(|&m| {
630                dag.wiring[m]
631                    .iter()
632                    .map(|src| match src {
633                        WireSource::NodeOutput(j, p) if local.contains_key(j) => {
634                            WireSource::NodeOutput(local[j], *p)
635                        }
636                        other => WireSource::Input(in_pos[&src_key(other)]),
637                    })
638                    .collect()
639            })
640            .collect();
641        let sub_input_defs: Vec<InputDef> = plan
642            .in_types
643            .iter()
644            .enumerate()
645            .map(|(i, ty)| InputDef {
646                name: format!("c{i}"),
647                default: default_for(*ty),
648                port_type: *ty,
649                kind: InputKind::Coordinate,
650                type_origin: crate::kernel::TypeOrigin::Declared,
651                converts_to: None,
652            })
653            .collect();
654        let mut sub_output_map: HashMap<String, (usize, usize)> = HashMap::new();
655        let mut sub_output_order: Vec<String> = Vec::new();
656        for (k, (j, p)) in plan.boundary_out.iter().enumerate() {
657            let name = format!("o{k}");
658            sub_output_map.insert(name.clone(), (local[j], *p));
659            sub_output_order.push(name);
660        }
661
662        let taken: Vec<Box<dyn PolydatNode>> = plan
663            .members
664            .iter()
665            .map(|&m| nodes[m].take().expect("cone member present"))
666            .collect();
667        let member_label = cone_label(&taken);
668
669        let mut sub = ResolvedDag {
670            nodes: taken,
671            wiring: sub_wiring,
672            input_defs: sub_input_defs,
673            coord_count: plan.boundary_in.len(),
674            output_map: sub_output_map,
675            output_order: sub_output_order,
676            cursor_schemas: Vec::new(),
677            source: String::new(),
678            // A member's failure is reported against the program the
679            // cone stands in, as the same node's failure is reported on
680            // every other engine (A7); the cone is not a frame of its own.
681            context: dag.context.clone(),
682            output_modifiers: HashMap::new(),
683            const_outputs: std::collections::HashSet::new(),
684            // A cone is a fragment of the program that stands in the
685            // tree's ledger already, not a program of its own: its
686            // kernel is recorded nowhere.
687            ledger: crate::kernel::CompileLedger::new(),
688        };
689
690        let restore = |sub_nodes: Vec<Box<dyn PolydatNode>>,
691                       nodes: &mut [Option<Box<dyn PolydatNode>>]| {
692            for (&m, nd) in plan.members.iter().zip(sub_nodes) {
693                nodes[m] = Some(nd);
694            }
695        };
696
697        let layout = match PolydatAssembler::build_jit_layout(&sub) {
698            Ok(l) => l,
699            Err(e) => {
700                restore(sub.nodes, nodes);
701                return Err(e);
702            }
703        };
704        let (coord_slots, total_slots, jit_steps, jit_outputs, scratch, _volatile) = layout;
705        // Boundary inputs occupy the first slots, each as wide as its
706        // type.
707        let mut in_slots = Vec::with_capacity(plan.in_types.len());
708        let mut next = 0usize;
709        for ty in &plan.in_types {
710            in_slots.push(next);
711            next += ty.slot_width();
712        }
713        debug_assert_eq!(coord_slots, next);
714        let compiled = crate::compile::jit::compile_jit_entry(&jit_steps, Some(total_slots));
715        let (code_fn, code) = match compiled {
716            Ok(parts) => parts,
717            Err(e) => {
718                restore(sub.nodes, nodes);
719                return Err(e);
720            }
721        };
722
723        let out_slots: Vec<usize> = (0..plan.boundary_out.len())
724            .map(|k| jit_outputs[&format!("o{k}")])
725            .collect();
726        // Port metadata mirrors the fused subgraph rather than
727        // being synthesized: outputs clone the member's original
728        // port (lifecycle analysis and downstream diagnostics see
729        // what the interpreter form would have declared); inputs
730        // clone the source port where one exists (graph inputs are
731        // per-cycle by definition).
732        let meta = NodeMeta {
733            name: member_label,
734            ins: plan
735                .boundary_in
736                .iter()
737                .zip(&plan.in_types)
738                .enumerate()
739                .map(|(i, (src, ty))| {
740                    // Boundary producers are ineligible nodes by
741                    // definition, so they are never cone members
742                    // and always present in the slot vec.
743                    let mut port = match src {
744                        WireSource::NodeOutput(j, p) => nodes[*j]
745                            .as_ref()
746                            .map(|nd| nd.meta().outs[*p].clone())
747                            .unwrap_or_else(|| Port::new("", *ty)),
748                        WireSource::Input(_) => Port::new("", *ty),
749                    };
750                    port.name = format!("c{i}");
751                    port.constraint = None;
752                    Slot::Wire(port)
753                })
754                .collect(),
755            outs: plan
756                .boundary_out
757                .iter()
758                .enumerate()
759                .map(|(k, (j, p))| {
760                    let mut port = sub.nodes[local[j]].meta().outs[*p].clone();
761                    port.name = format!("o{k}");
762                    port.constraint = None;
763                    port
764                })
765                .collect(),
766        };
767        let out_ports: Vec<(usize, usize)> = plan
768            .boundary_out
769            .iter()
770            .map(|(j, p)| (local[j], *p))
771            .collect();
772        // A member's failure names the member's outputs as the program
773        // names them (A7), not as the cone numbers them: the boundary
774        // outputs take the program's names for the attribution.
775        let mut named = sub.output_map.clone();
776        for (k, (j, p)) in plan.boundary_out.iter().enumerate() {
777            let names: Vec<String> = dag
778                .output_map
779                .iter()
780                .filter(|(_, v)| **v == (*j, *p))
781                .map(|(n, _)| n.clone())
782                .collect();
783            if !names.is_empty()
784                && let Some(target) = named.remove(&format!("o{k}"))
785            {
786                for n in names {
787                    named.insert(n, target);
788                }
789            }
790        }
791        let numbered = std::mem::replace(&mut sub.output_map, named);
792        let attribution = std::sync::Arc::new(PolydatAssembler::attribution_of(&sub));
793        sub.output_map = numbered;
794        Ok(JitConeNode {
795            attribution,
796            in_slots,
797            meta,
798            code_fn,
799            total_slots,
800            out_slots,
801            in_types: plan.in_types.clone(),
802            out_types: plan.out_types.clone(),
803            members: sub.nodes,
804            sub_wiring: sub.wiring,
805            out_ports,
806            scratch,
807            fallible: code.fallible(),
808            _module: code,
809        })
810    }
811
812    /// Diagnostic name carrying the fused members, so an enriched
813    /// eval panic attributes the interior functions.
814    fn cone_label(members: &[Box<dyn PolydatNode>]) -> String {
815        const SHOWN: usize = 6;
816        let names: Vec<&str> = members
817            .iter()
818            .take(SHOWN)
819            .map(|n| n.meta().name.as_str())
820            .collect();
821        let suffix = if members.len() > SHOWN {
822            format!("+{} more", members.len() - SHOWN)
823        } else {
824            String::new()
825        };
826        format!("jit_cone[{}{}]", names.join("+"), suffix)
827    }
828
829    /// Splice the compiled cones into the DAG and restore
830    /// topological order.
831    fn rebuild(
832        dag: &mut ResolvedDag,
833        nodes_opt: Vec<Option<Box<dyn PolydatNode>>>,
834        cones: Vec<(ConePlan, JitConeNode)>,
835    ) {
836        let old_n = nodes_opt.len();
837        // (old_idx, port) → (cone_ordinal, cone_out_port)
838        let mut cone_port: HashMap<(usize, usize), (usize, usize)> = HashMap::new();
839        for (ci, (plan, _)) in cones.iter().enumerate() {
840            for (k, (j, p)) in plan.boundary_out.iter().enumerate() {
841                cone_port.insert((*j, *p), (ci, k));
842            }
843        }
844
845        let mut kept_map: HashMap<usize, usize> = HashMap::new();
846        let mut new_nodes: Vec<Box<dyn PolydatNode>> = Vec::new();
847        let mut new_wiring: Vec<Vec<WireSource>> = Vec::new();
848        for (old, slot) in nodes_opt.into_iter().enumerate() {
849            if let Some(node) = slot {
850                kept_map.insert(old, new_nodes.len());
851                new_nodes.push(node);
852                new_wiring.push(dag.wiring[old].clone());
853            }
854        }
855        let cone_base = new_nodes.len();
856        let mut cone_plans: Vec<ConePlan> = Vec::with_capacity(cones.len());
857        for (plan, cone) in cones {
858            new_nodes.push(Box::new(cone));
859            new_wiring.push(plan.boundary_in.clone());
860            cone_plans.push(plan);
861        }
862
863        let remap = |src: &WireSource| -> WireSource {
864            match src {
865                WireSource::Input(i) => WireSource::Input(*i),
866                WireSource::NodeOutput(j, p) => {
867                    if let Some(&nj) = kept_map.get(j) {
868                        WireSource::NodeOutput(nj, *p)
869                    } else {
870                        let (ci, k) = cone_port[&(*j, *p)];
871                        WireSource::NodeOutput(cone_base + ci, k)
872                    }
873                }
874            }
875        };
876        for wiring in new_wiring.iter_mut() {
877            for src in wiring.iter_mut() {
878                *src = remap(src);
879            }
880        }
881        let mut new_output_map: HashMap<String, (usize, usize)> = HashMap::new();
882        for (name, (j, p)) in dag.output_map.iter() {
883            let (nj, np) = match remap(&WireSource::NodeOutput(*j, *p)) {
884                WireSource::NodeOutput(a, b) => (a, b),
885                WireSource::Input(_) => unreachable!("outputs map to nodes"),
886            };
887            new_output_map.insert(name.clone(), (nj, np));
888        }
889
890        // Kahn topo sort — consumers of cone interiors may sit at
891        // indices below the spliced cone node.
892        let m = new_nodes.len();
893        let mut indegree = vec![0usize; m];
894        let mut dependents: Vec<Vec<usize>> = vec![Vec::new(); m];
895        for (i, wiring) in new_wiring.iter().enumerate() {
896            let mut producers: Vec<usize> = wiring
897                .iter()
898                .filter_map(|s| match s {
899                    WireSource::NodeOutput(j, _) => Some(*j),
900                    WireSource::Input(_) => None,
901                })
902                .collect();
903            producers.sort_unstable();
904            producers.dedup();
905            indegree[i] = producers.len();
906            for j in producers {
907                dependents[j].push(i);
908            }
909        }
910        let mut order: Vec<usize> = Vec::with_capacity(m);
911        let mut ready: std::collections::BinaryHeap<std::cmp::Reverse<usize>> = (0..m)
912            .filter(|&i| indegree[i] == 0)
913            .map(std::cmp::Reverse)
914            .collect();
915        while let Some(std::cmp::Reverse(i)) = ready.pop() {
916            order.push(i);
917            for &d in &dependents[i] {
918                indegree[d] -= 1;
919                if indegree[d] == 0 {
920                    ready.push(std::cmp::Reverse(d));
921                }
922            }
923        }
924        assert_eq!(
925            order.len(),
926            m,
927            "cone splice must not introduce a cycle (old_n={old_n})"
928        );
929        let mut pos = vec![0usize; m];
930        for (new_idx, &i) in order.iter().enumerate() {
931            pos[i] = new_idx;
932        }
933
934        let mut sorted_nodes: Vec<Option<Box<dyn PolydatNode>>> =
935            new_nodes.into_iter().map(Some).collect();
936        dag.nodes = order
937            .iter()
938            .map(|&i| sorted_nodes[i].take().expect("each node placed once"))
939            .collect();
940        dag.wiring = order
941            .iter()
942            .map(|&i| {
943                new_wiring[i]
944                    .iter()
945                    .map(|s| match s {
946                        WireSource::Input(k) => WireSource::Input(*k),
947                        WireSource::NodeOutput(j, p) => WireSource::NodeOutput(pos[*j], *p),
948                    })
949                    .collect()
950            })
951            .collect();
952        dag.output_map = new_output_map
953            .into_iter()
954            .map(|(name, (j, p))| (name, (pos[j], p)))
955            .collect();
956        let _ = cone_plans;
957    }
958}