bb_ir/types/
storage.rs

1//! `Storage` — static link between a Rust storage type and its IR
2//! `TypeNode`. Library makers declare *where in the polymorphism tree*
3//! their concrete sits by picking the `Storage` impl for their
4//! associated type.
5
6use super::builtins::{
7    TYPE_SCALAR_F16, TYPE_SCALAR_F32, TYPE_SCALAR_F64, TYPE_SCALAR_I32, TYPE_SCALAR_U8,
8    TYPE_TENSOR, TYPE_TENSOR_BF16, TYPE_TENSOR_BOOL, TYPE_TENSOR_F16, TYPE_TENSOR_F32,
9    TYPE_TENSOR_F64, TYPE_TENSOR_I16, TYPE_TENSOR_I32, TYPE_TENSOR_I64, TYPE_TENSOR_I8,
10    TYPE_TENSOR_U16, TYPE_TENSOR_U32, TYPE_TENSOR_U64, TYPE_TENSOR_U8,
11};
12use super::TypeNode;
13
14/// Static link between a Rust storage type and its IR `TypeNode`.
15///
16/// Library makers declare what wire format a backend / index / model
17/// natively understands. The framework reads `Storage::TYPE` at
18/// recording time (and at slot binding time) to stamp every Contract
19/// port's `value_info` — the compiler's type solver then walks the
20/// graph and refuses any mismatch.
21pub trait Storage: Send + Sync + 'static {
22    /// Position-in-tree declaration. The `TypeNode` static this
23    /// constant points at decides what other storage types unify with
24    /// this one during the type-solver walk.
25    const TYPE: &'static TypeNode;
26}
27
28// --- Tensor (slice) leaf impls --------------------------------------
29
30impl Storage for [f32] {
31    const TYPE: &'static TypeNode = &TYPE_TENSOR_F32;
32}
33impl Storage for [f64] {
34    const TYPE: &'static TypeNode = &TYPE_TENSOR_F64;
35}
36// Half-precision float storage uses `half::f16` (canonical Rust newtype
37// around the binary16 bit pattern). The earlier polymorphic-storage
38// design conflated `[u16]` with packed half, but that prevented u16
39// from being a first-class unsigned-integer tensor type. Both are
40// distinct positions in the polymorphism tree now.
41impl Storage for [half::f16] {
42    const TYPE: &'static TypeNode = &TYPE_TENSOR_F16;
43}
44impl Storage for [half::bf16] {
45    const TYPE: &'static TypeNode = &TYPE_TENSOR_BF16;
46}
47impl Storage for [u8] {
48    const TYPE: &'static TypeNode = &TYPE_TENSOR_U8;
49}
50impl Storage for [u16] {
51    const TYPE: &'static TypeNode = &TYPE_TENSOR_U16;
52}
53impl Storage for [u32] {
54    const TYPE: &'static TypeNode = &TYPE_TENSOR_U32;
55}
56impl Storage for [u64] {
57    const TYPE: &'static TypeNode = &TYPE_TENSOR_U64;
58}
59impl Storage for [i8] {
60    const TYPE: &'static TypeNode = &TYPE_TENSOR_I8;
61}
62impl Storage for [i16] {
63    const TYPE: &'static TypeNode = &TYPE_TENSOR_I16;
64}
65impl Storage for [i32] {
66    const TYPE: &'static TypeNode = &TYPE_TENSOR_I32;
67}
68impl Storage for [i64] {
69    const TYPE: &'static TypeNode = &TYPE_TENSOR_I64;
70}
71impl Storage for [bool] {
72    const TYPE: &'static TypeNode = &TYPE_TENSOR_BOOL;
73}
74
75// --- Scalar leaf impls ----------------------------------------------
76
77impl Storage for f32 {
78    const TYPE: &'static TypeNode = &TYPE_SCALAR_F32;
79}
80impl Storage for f64 {
81    const TYPE: &'static TypeNode = &TYPE_SCALAR_F64;
82}
83// Scalar half stays as `u16` (single-element bit pattern) for now —
84// scalar variants don't yet have their own `half::f16` impl because
85// scalar polymorphism isn't as load-bearing as tensor polymorphism.
86impl Storage for u16 {
87    const TYPE: &'static TypeNode = &TYPE_SCALAR_F16;
88}
89impl Storage for u8 {
90    const TYPE: &'static TypeNode = &TYPE_SCALAR_U8;
91}
92impl Storage for i32 {
93    const TYPE: &'static TypeNode = &TYPE_SCALAR_I32;
94}
95
96// --- Generic-position storage ---------------------------------------
97
98/// Concrete-erased tensor. Stores raw bytes plus a runtime-known
99/// dtype + shape. Compute-outsourcing concretes (an index that
100/// delegates distance math to a bound Backend) declare
101/// `type Vector = AnyTensor` — `Storage::TYPE = &TYPE_TENSOR` puts
102/// the value at a non-leaf position in the tree, so any tensor
103/// subtype unifies into it.
104#[derive(Clone, Debug)]
105pub struct AnyTensor {
106    /// Raw little-endian bytes of the tensor payload, packed per `dtype`.
107    pub bytes: Vec<u8>,
108    /// Runtime dtype. Pair with `Self::shape` for full interpretation.
109    pub dtype: Dtype,
110    /// Per-axis shape. `shape.iter().product::<usize>()` × dtype-size
111    /// is expected to equal `bytes.len()`.
112    pub shape: Vec<usize>,
113}
114
115impl Storage for AnyTensor {
116    const TYPE: &'static TypeNode = &TYPE_TENSOR;
117}
118
119/// Runtime dtype tag for `AnyTensor` and any caller that needs to
120/// dispatch over the framework's known tensor dtypes.
121#[derive(Clone, Copy, Debug, PartialEq, Eq)]
122pub enum Dtype {
123    /// 32-bit IEEE 754 float — corresponds to `TYPE_TENSOR_F32`.
124    F32,
125    /// 64-bit IEEE 754 float — corresponds to `TYPE_TENSOR_F64`.
126    F64,
127    /// 16-bit IEEE 754 half — corresponds to `TYPE_TENSOR_F16`. Stored as `half::f16`.
128    F16,
129    /// 16-bit bfloat (Google's brain-float) — corresponds to `TYPE_TENSOR_BF16`. Stored as `half::bf16`.
130    BF16,
131    /// 8-bit unsigned — corresponds to `TYPE_TENSOR_U8`.
132    U8,
133    /// 16-bit unsigned — corresponds to `TYPE_TENSOR_U16`.
134    U16,
135    /// 32-bit unsigned — corresponds to `TYPE_TENSOR_U32`.
136    U32,
137    /// 64-bit unsigned — corresponds to `TYPE_TENSOR_U64`.
138    U64,
139    /// 8-bit signed — corresponds to `TYPE_TENSOR_I8` (edge quantization).
140    I8,
141    /// 16-bit signed — corresponds to `TYPE_TENSOR_I16`.
142    I16,
143    /// 32-bit signed — corresponds to `TYPE_TENSOR_I32`.
144    I32,
145    /// 64-bit signed — corresponds to `TYPE_TENSOR_I64` (indices, embeddings).
146    I64,
147    /// 1-bit boolean — corresponds to `TYPE_TENSOR_BOOL` (masks, predicates).
148    Bool,
149}
150
151impl Dtype {
152    /// Return the framework `TypeNode` static for this dtype.
153    pub fn type_node(self) -> &'static TypeNode {
154        match self {
155            Dtype::F32 => &TYPE_TENSOR_F32,
156            Dtype::F64 => &TYPE_TENSOR_F64,
157            Dtype::F16 => &TYPE_TENSOR_F16,
158            Dtype::BF16 => &TYPE_TENSOR_BF16,
159            Dtype::U8 => &TYPE_TENSOR_U8,
160            Dtype::U16 => &TYPE_TENSOR_U16,
161            Dtype::U32 => &TYPE_TENSOR_U32,
162            Dtype::U64 => &TYPE_TENSOR_U64,
163            Dtype::I8 => &TYPE_TENSOR_I8,
164            Dtype::I16 => &TYPE_TENSOR_I16,
165            Dtype::I32 => &TYPE_TENSOR_I32,
166            Dtype::I64 => &TYPE_TENSOR_I64,
167            Dtype::Bool => &TYPE_TENSOR_BOOL,
168        }
169    }
170}
bb_ir/types/storage.rs

bb_ir/types/
storage.rs
