Skip to main content

polydat_core/compile/
marshal.rs

1// Copyright 2024-2026 Jonathan Shook
2// SPDX-License-Identifier: Apache-2.0
3
4//! Marshalling between typed `Value`s and compiled-tier slots. One
5//! place owns the rule for every port type, so embedded cones and
6//! whole compiled kernels decode alike:
7//!
8//! - scalars ride as their slot bits, one slot or two;
9//! - every `Ref2` kind rides as a `(ptr, len)` pair (jit_boundary.md,
10//!   axioms S1–S10): a typed vector, a string, or a byte string as a
11//!   slice of its elements, and a JSON, extension, or handle value as
12//!   a one-element slice holding the `Value`.
13//!
14//! Entering the compiled tier, a `Value` is borrowed: the pair points
15//! into the value the caller holds for the duration of the call.
16//! Leaving it, the pair is copied out to an owned `Value`: the
17//! interpreter and the host never hold a reference into a state's
18//! buffers.
19
20use crate::ast::SlotShape;
21use crate::ast::{PortType, Value};
22
23/// The `(ptr, len)` pair a `Ref2`-kind value is borrowed as: into the
24/// value's own bytes or element storage, or the value itself for a
25/// JSON, extension, or handle value. The pair is valid while `v` is.
26/// `None` for a scalar or for a variant the port type does not admit.
27pub(crate) fn borrow_pair(v: &Value) -> Option<(u64, u64)> {
28    Some(match v {
29        Value::Str(s) => (s.as_ptr() as usize as u64, s.len() as u64),
30        Value::Bytes(b) => (b.as_ptr() as usize as u64, b.len() as u64),
31        Value::Json(_) | Value::Ext(_) | Value::Handle(_) => (v as *const Value as usize as u64, 1),
32        Value::VecF32(v) => slice_pair(v.as_slice()),
33        Value::VecF64(v) => slice_pair(v.as_slice()),
34        Value::VecF16(v) => slice_pair(v.as_slice()),
35        Value::VecI8(v) => slice_pair(v.as_slice()),
36        Value::VecI16(v) => slice_pair(v.as_slice()),
37        Value::VecI32(v) => slice_pair(v.as_slice()),
38        Value::VecI64(v) => slice_pair(v.as_slice()),
39        _ => return None,
40    })
41}
42
43/// The `(ptr, len)` pair of a slice.
44#[inline]
45pub(crate) fn slice_pair<T>(s: &[T]) -> (u64, u64) {
46    (s.as_ptr() as usize as u64, s.len() as u64)
47}
48
49/// An empty pair for an unset `Ref2` slot: a dangling, non-null
50/// pointer with length zero, which every slice reader accepts.
51#[inline]
52pub(crate) fn empty_pair() -> (u64, u64) {
53    (
54        std::ptr::NonNull::<u8>::dangling().as_ptr() as usize as u64,
55        0,
56    )
57}
58
59/// A `Value` as the slot bits its port type `ty` rides as: the bits of a
60/// scalar, or the borrowed pair of a `Ref2` kind. A `Dyn` port names the
61/// value itself, whatever its variant, for its converter to read. `None`
62/// when the value's variant has no compiled representation.
63#[cfg(feature = "jit")]
64pub(crate) fn encode_slots(v: &Value, ty: PortType, out: &mut [u64]) -> Option<()> {
65    if ty == PortType::Dyn {
66        (out[0], out[1]) = match v {
67            Value::None => empty_pair(),
68            v => (v as *const Value as usize as u64, 1),
69        };
70        return Some(());
71    }
72    match v {
73        Value::U64(x) => out[0] = *x,
74        Value::I64(x) => out[0] = *x as u64,
75        Value::F64(x) => out[0] = x.to_bits(),
76        Value::Bool(b) => out[0] = *b as u64,
77        _ => {
78            let (p, l) = borrow_pair(v)?;
79            out[0] = p;
80            out[1] = l;
81        }
82    }
83    Some(())
84}
85
86/// The value a `Ref2` pair names, copied out as the `Value` its port
87/// type declares. Every read that leaves the compiled tier goes
88/// through here.
89///
90/// # Safety
91///
92/// The pair was published by a producer whose storage is alive: its
93/// own scratch, an extern's stored value, an interned constant, or a
94/// boundary value alive for the call (axioms S3, S4).
95pub(crate) unsafe fn decode_pair(ty: PortType, ptr: u64, len: u64) -> Value {
96    use crate::ast::SliceArc;
97    let (p, n) = (ptr as usize, len as usize);
98    // SAFETY: as documented on the function.
99    unsafe {
100        match ty {
101            PortType::Str => Value::Str(std::sync::Arc::from(std::str::from_utf8_unchecked(
102                std::slice::from_raw_parts(p as *const u8, n),
103            ))),
104            PortType::Bytes => Value::Bytes(std::sync::Arc::from(std::slice::from_raw_parts(
105                p as *const u8,
106                n,
107            ))),
108            PortType::Json | PortType::Ext | PortType::Handle | PortType::Dyn => {
109                if n == 0 {
110                    Value::None
111                } else {
112                    (*(p as *const Value)).clone()
113                }
114            }
115            PortType::VecF32 => Value::VecF32(SliceArc::from_vec(
116                std::slice::from_raw_parts(p as *const f32, n).to_vec(),
117            )),
118            PortType::VecF64 => Value::VecF64(SliceArc::from_vec(
119                std::slice::from_raw_parts(p as *const f64, n).to_vec(),
120            )),
121            PortType::VecF16 => Value::VecF16(SliceArc::from_vec(
122                std::slice::from_raw_parts(p as *const half::f16, n).to_vec(),
123            )),
124            PortType::VecI8 => Value::VecI8(SliceArc::from_vec(
125                std::slice::from_raw_parts(p as *const i8, n).to_vec(),
126            )),
127            PortType::VecI16 => Value::VecI16(SliceArc::from_vec(
128                std::slice::from_raw_parts(p as *const i16, n).to_vec(),
129            )),
130            PortType::VecI32 => Value::VecI32(SliceArc::from_vec(
131                std::slice::from_raw_parts(p as *const i32, n).to_vec(),
132            )),
133            PortType::VecI64 => Value::VecI64(SliceArc::from_vec(
134                std::slice::from_raw_parts(p as *const i64, n).to_vec(),
135            )),
136            other => panic!("{other:?} is not a Ref2-colored port type"),
137        }
138    }
139}
140
141/// The port type of the `Value` a port of type `ty` carries, which is
142/// what [`decode_slot`] reads back. A narrow integer rides the 64-bit
143/// carrier of its sign, and `f32` and `f16` ride their bit patterns in
144/// a `U64`, as their `Wire` impls inject them; every other type is its
145/// own carrier.
146pub(crate) fn carrier_port(ty: PortType) -> PortType {
147    match ty {
148        PortType::I8 | PortType::I16 | PortType::I32 => PortType::I64,
149        PortType::U32 | PortType::U16 | PortType::U8 | PortType::F32 | PortType::F16 => {
150            PortType::U64
151        }
152        other => other,
153    }
154}
155
156/// Slot bits as the `Value` their declared port type names, copied
157/// out where they are a pair.
158///
159/// The read is chosen by the type's slot color, as `write_poly`'s
160/// write is, so every compiled read of a value (an output a host
161/// pulls, and an input a polymorphic node takes) is the inverse of
162/// the one write. This used to match on the type with a `U64`
163/// fallback. The `Imm2` limb reassembly lived only in
164/// `decode_output`, so a register or 128-bit value reaching a
165/// polymorphic node's input fell through to the fallback and arrived
166/// as its low limb, typed `U64`. The type system knew better at every
167/// step; only the fallback arm did not.
168pub(crate) fn decode_slot(slots: &[u64], ty: PortType) -> Value {
169    use crate::ast::{Bits128, RegLanes, SlotColor};
170    match ty.slot_color() {
171        SlotColor::Imm1 => match ty {
172            PortType::F64 => Value::F64(f64::from_bits(slots[0])),
173            PortType::Bool => Value::Bool(slots[0] != 0),
174            // A signed narrow carrier rides sign-extended (alignment
175            // §2) and is the `I64` value its `Wire` impl injects.
176            PortType::I64 | PortType::I8 | PortType::I16 | PortType::I32 => {
177                Value::I64(slots[0] as i64)
178            }
179            // Unsigned narrow carriers ride zero-extended, and `f32`
180            // and `f16` ride their bit patterns, all in the `U64`
181            // carrier their `Wire` impls inject.
182            PortType::U64
183            | PortType::U32
184            | PortType::U16
185            | PortType::U8
186            | PortType::F32
187            | PortType::F16 => Value::U64(slots[0]),
188            other => unreachable!("{other:?} is colored Imm1 but has no one-slot carrier"),
189        },
190        // Two consecutive slots of limb data, low word first
191        // (alignment §6).
192        SlotColor::Imm2 => {
193            let limbs = Bits128([slots[0], slots[1]]);
194            match ty {
195                PortType::U128 => Value::U128(limbs),
196                PortType::I128 => Value::I128(limbs),
197                PortType::Reg128 => Value::Reg128(limbs, RegLanes::Raw),
198                PortType::RegI8x16 => Value::Reg128(limbs, RegLanes::I8x16),
199                PortType::RegI16x8 => Value::Reg128(limbs, RegLanes::I16x8),
200                PortType::RegI32x4 => Value::Reg128(limbs, RegLanes::I32x4),
201                PortType::RegI64x2 => Value::Reg128(limbs, RegLanes::I64x2),
202                PortType::RegF16x8 => Value::Reg128(limbs, RegLanes::F16x8),
203                PortType::RegF32x4 => Value::Reg128(limbs, RegLanes::F32x4),
204                PortType::RegF64x2 => Value::Reg128(limbs, RegLanes::F64x2),
205                other => unreachable!("{other:?} is colored Imm2 but is not a 128-bit value"),
206            }
207        }
208        // SAFETY: the pair in a kernel's buffer was published by a
209        // producer whose storage is alive (axioms S3, S4).
210        SlotColor::Ref2 => unsafe { decode_pair(ty, slots[0], slots[1]) },
211    }
212}
213
214/// An output at `slot` of `buffer` as the `Value` its port type names.
215/// This is the typed read every compiled kernel's `get_value` makes.
216pub fn decode_output(buffer: &[u64], slot: usize, ty: PortType) -> Value {
217    decode_slot(&buffer[slot..], ty)
218}
219
220/// An argument's slots as a borrowed view of the value, decoded by its
221/// port type: nothing is copied. A string, byte string, or value is
222/// borrowed through its pair, valid while its producer's storage is.
223///
224/// # Safety
225///
226/// As for `decode_pair`.
227pub unsafe fn arg_ref<'a>(ty: PortType, slots: &'a [u64]) -> crate::ast::ValueRef<'a> {
228    use crate::ast::ValueRef;
229    let (p, n) = (
230        slots.first().copied().unwrap_or(0) as usize,
231        slots.get(1).copied().unwrap_or(0) as usize,
232    );
233    // SAFETY: as documented on the function.
234    unsafe {
235        match ty {
236            PortType::U64 => ValueRef::U64(slots[0]),
237            PortType::I64 | PortType::I8 | PortType::I16 | PortType::I32 => {
238                ValueRef::I64(slots[0] as i64)
239            }
240            PortType::F64 => ValueRef::F64(f64::from_bits(slots[0])),
241            PortType::Bool => ValueRef::Bool(slots[0] != 0),
242            PortType::Str => ValueRef::Str(std::str::from_utf8_unchecked(
243                std::slice::from_raw_parts(p as *const u8, n),
244            )),
245            PortType::Bytes => ValueRef::Bytes(std::slice::from_raw_parts(p as *const u8, n)),
246            PortType::Json | PortType::Ext | PortType::Handle | PortType::Dyn => {
247                if n == 0 {
248                    ValueRef::None
249                } else {
250                    ValueRef::from(&*(p as *const Value))
251                }
252            }
253            _ => ValueRef::U64(slots[0]),
254        }
255    }
256}