sim_lib_numbers_tensor/implementation/
value.rs

1//! The uniform `Tensor` value type: its shape, dtype, and cell storage, with
2//! indexing, construction, and number-value behavior backing the tensor domain.
3
4use std::cmp::Ordering;
5use std::collections::{BTreeMap, BinaryHeap};
6use std::sync::Arc;
7
8use sim_kernel::{
9    ClassRef, Cx, DefaultFactory, Error, Expr, Factory, NoopEvalPolicy, NumberValue, Object,
10    ObjectCompat, ObjectEncode, ObjectEncoding, Result, Symbol, Value,
11};
12
13use super::citizen::tensor_value_class_symbol;
14use super::domain::number_domain;
15
16/// The uniform tensor value: an n-dimensional array of scalar number cells.
17///
18/// A tensor is row-major (last axis varies fastest) and homogeneous: every cell
19/// shares the [`dtype`](Tensor::dtype) number domain. An empty [`shape`](Tensor::shape)
20/// denotes a rank-0 scalar holding a single cell. Tensors are the value backing
21/// the `numbers/tensor` domain and are constructed through
22/// [`build_tensor_value`] rather than parsed from literals.
23#[derive(Clone)]
24pub struct Tensor {
25    /// Length of each axis, outermost first. Empty for a rank-0 scalar.
26    shape: Vec<usize>,
27    /// The shared scalar number domain of every cell (for example
28    /// `numbers/i64` or `numbers/f64`).
29    dtype: Symbol,
30    /// Row-major cell storage; its length equals the product of `shape`
31    /// (one for a scalar).
32    data: Vec<Value>,
33}
34
35impl Tensor {
36    /// Builds a tensor after checking shape, scalar-cell, and promotion
37    /// invariants against the loaded number-domain registry.
38    pub fn new_checked(
39        cx: &mut Cx,
40        shape: Vec<usize>,
41        dtype: Symbol,
42        data: Vec<Value>,
43    ) -> Result<Self> {
44        validate_shape_and_data_len(&shape, data.len())?;
45        validate_cells(cx, &data)?;
46        validate_dtype_accepts_cells(cx, &dtype, &data)?;
47        Ok(Self { shape, dtype, data })
48    }
49
50    /// Builds a tensor whose cells already exactly match `dtype`.
51    ///
52    /// Specialized tensor backends use this when converting packed storage into
53    /// uniform tensor cells without needing a loaded registry. It is stricter
54    /// than [`Tensor::new_checked`]: every cell must report exactly `dtype`.
55    pub fn new_exact(shape: Vec<usize>, dtype: Symbol, data: Vec<Value>) -> Result<Self> {
56        validate_shape_and_data_len(&shape, data.len())?;
57        validate_exact_cell_dtype(&dtype, &data)?;
58        Ok(Self { shape, dtype, data })
59    }
60
61    /// The tensor shape, outermost axis first. Empty means a rank-0 scalar.
62    pub fn shape(&self) -> &[usize] {
63        &self.shape
64    }
65
66    /// The shared scalar number domain accepted by every tensor cell.
67    pub fn dtype(&self) -> &Symbol {
68        &self.dtype
69    }
70
71    /// Row-major scalar cell storage.
72    pub fn data(&self) -> &[Value] {
73        &self.data
74    }
75
76    /// The number of axes, i.e. the length of [`shape`](Tensor::shape). Zero
77    /// for a scalar.
78    pub fn rank(&self) -> usize {
79        self.shape.len()
80    }
81
82    /// Computes the row-major flat offset into [`data`](Tensor::data) for a
83    /// multi-dimensional `indices` coordinate against `shape`.
84    ///
85    /// Returns an error if the index rank does not match `shape` or any
86    /// component is out of bounds.
87    ///
88    /// # Examples
89    ///
90    /// ```
91    /// use sim_lib_numbers_tensor::Tensor;
92    ///
93    /// // Row-major 2x3 tensor: element (1, 2) is at flat offset 5.
94    /// assert_eq!(Tensor::flat_offset(&[2, 3], &[1, 2]).unwrap(), 5);
95    /// assert_eq!(Tensor::flat_offset(&[2, 3], &[0, 0]).unwrap(), 0);
96    /// // Out-of-bounds and rank-mismatched indices are rejected.
97    /// assert!(Tensor::flat_offset(&[2, 3], &[2, 0]).is_err());
98    /// assert!(Tensor::flat_offset(&[2, 3], &[0]).is_err());
99    /// ```
100    pub fn flat_offset(shape: &[usize], indices: &[usize]) -> Result<usize> {
101        if shape.len() != indices.len() {
102            return Err(Error::Eval("tensor index rank mismatch".to_owned()));
103        }
104        let mut stride = 1usize;
105        let mut offset = 0usize;
106        for (dim, index) in shape.iter().rev().zip(indices.iter().rev()) {
107            if *index >= *dim {
108                return Err(Error::Eval("tensor index was out of bounds".to_owned()));
109            }
110            offset += index * stride;
111            stride = stride.saturating_mul(*dim);
112        }
113        Ok(offset)
114    }
115
116    /// Enumerates every multi-dimensional coordinate of `shape` in row-major
117    /// order. An empty shape yields a single empty coordinate (the scalar cell).
118    pub fn coordinates(shape: &[usize]) -> Vec<Vec<usize>> {
119        if shape.is_empty() {
120            return vec![Vec::new()];
121        }
122        let mut out = Vec::new();
123        let mut coord = vec![0usize; shape.len()];
124        loop {
125            out.push(coord.clone());
126            let mut axis = shape.len();
127            while axis > 0 {
128                axis -= 1;
129                coord[axis] += 1;
130                if coord[axis] < shape[axis] {
131                    break;
132                }
133                coord[axis] = 0;
134                if axis == 0 {
135                    return out;
136                }
137            }
138        }
139    }
140}
141
142impl Object for Tensor {
143    fn display(&self, cx: &mut Cx) -> Result<String> {
144        match self.as_expr(cx)? {
145            Expr::Call { .. } => Ok(format!("{}<{:?}>", tensor_display_name(), self.shape)),
146            expr => Ok(format!("{expr:?}")),
147        }
148    }
149
150    fn as_any(&self) -> &dyn std::any::Any {
151        self
152    }
153}
154
155impl sim_kernel::ObjectCompat for Tensor {
156    fn class(&self, cx: &mut Cx) -> Result<ClassRef> {
157        if let Some(value) = cx.registry().class_by_symbol(&tensor_value_class_symbol()) {
158            return Ok(value.clone());
159        }
160        if let Some(value) = cx
161            .registry()
162            .class_by_symbol(&Symbol::qualified("core", "Number"))
163        {
164            return Ok(value.clone());
165        }
166        DefaultFactory.class_stub(
167            sim_kernel::CORE_NUMBER_CLASS_ID,
168            Symbol::qualified("core", "Number"),
169        )
170    }
171    fn as_expr(&self, cx: &mut Cx) -> Result<Expr> {
172        match self.rank() {
173            0 => Ok(Expr::Call {
174                operator: Box::new(Expr::Symbol(Symbol::new("scalar"))),
175                args: vec![
176                    self.data
177                        .first()
178                        .ok_or_else(|| Error::Eval("scalar tensor is missing its cell".to_owned()))?
179                        .object()
180                        .as_expr(cx)?,
181                ],
182            }),
183            1 => Ok(Expr::Vector(exprs(cx, &self.data)?)),
184            2 => {
185                let width = self.shape[1];
186                let rows = self
187                    .data
188                    .chunks(width)
189                    .map(|row| exprs(cx, row).map(Expr::Vector))
190                    .collect::<Result<Vec<_>>>()?;
191                Ok(Expr::Vector(rows))
192            }
193            _ => Ok(Expr::Call {
194                operator: Box::new(Expr::Symbol(Symbol::new("tensor"))),
195                args: vec![
196                    Expr::Vector(
197                        self.shape
198                            .iter()
199                            .map(|dim| Expr::String(dim.to_string()))
200                            .collect(),
201                    ),
202                    Expr::Symbol(self.dtype.clone()),
203                    Expr::Vector(exprs(cx, &self.data)?),
204                ],
205            }),
206        }
207    }
208    fn as_table(&self, cx: &mut Cx) -> Result<Value> {
209        let shape = cx.factory().list(
210            self.shape
211                .iter()
212                .map(|dim| cx.factory().string(dim.to_string()))
213                .collect::<Result<Vec<_>>>()?,
214        )?;
215        let data = cx.factory().list(self.data.clone())?;
216        cx.factory().table(vec![
217            (
218                Symbol::new("kind"),
219                cx.factory().string("tensor".to_owned())?,
220            ),
221            (Symbol::new("shape"), shape),
222            (
223                Symbol::new("dtype"),
224                cx.factory().symbol(self.dtype.clone())?,
225            ),
226            (Symbol::new("data"), data),
227        ])
228    }
229    fn as_number_value(&self) -> Option<&dyn NumberValue> {
230        Some(self)
231    }
232
233    fn as_object_encoder(&self) -> Option<&dyn ObjectEncode> {
234        Some(self)
235    }
236}
237
238impl NumberValue for Tensor {
239    fn number_domain(&self, _cx: &mut Cx) -> Result<Symbol> {
240        Ok(number_domain())
241    }
242}
243
244impl ObjectEncode for Tensor {
245    fn object_encoding(&self, cx: &mut Cx) -> Result<ObjectEncoding> {
246        Ok(ObjectEncoding::Constructor {
247            class: tensor_value_class_symbol(),
248            args: vec![
249                Expr::Symbol(Symbol::new("v1")),
250                Expr::List(
251                    self.shape
252                        .iter()
253                        .map(|dim| {
254                            Expr::Number(sim_kernel::NumberLiteral {
255                                domain: Symbol::qualified("citizen", "int"),
256                                canonical: dim.to_string(),
257                            })
258                        })
259                        .collect(),
260                ),
261                Expr::List(exprs(cx, &self.data)?),
262                Expr::Symbol(self.dtype.clone()),
263            ],
264        })
265    }
266}
267
268impl sim_citizen::Citizen for Tensor {
269    fn citizen_symbol() -> Symbol {
270        tensor_value_class_symbol()
271    }
272
273    fn citizen_version() -> u32 {
274        1
275    }
276
277    fn citizen_arity() -> usize {
278        3
279    }
280
281    fn citizen_fields() -> &'static [&'static str] {
282        &["shape", "data", "domain"]
283    }
284}
285
286/// Builds a tensor [`Value`] of the given `shape` from row-major `data` cells.
287///
288/// The cell count must equal the product of `shape` (one for an empty, scalar
289/// shape). Every cell must be a scalar number value (not a nested tensor). When
290/// `dtype_hint` is `Some`, all cells must promote to that domain; otherwise the
291/// element domain is chosen as the cheapest join of the cell domains. Returns an
292/// error on a cell-count mismatch, a non-scalar cell, or an impossible dtype.
293pub fn build_tensor_value(
294    cx: &mut Cx,
295    shape: Vec<usize>,
296    dtype_hint: Option<Symbol>,
297    data: Vec<Value>,
298) -> Result<Value> {
299    validate_shape_and_data_len(&shape, data.len())?;
300    let dtype = choose_dtype(cx, dtype_hint, &data)?;
301    let tensor = Tensor::new_checked(cx, shape, dtype, data)?;
302    cx.factory().opaque(Arc::new(tensor))
303}
304
305/// Builds a rank-0 scalar tensor wrapping a single scalar number `value`.
306pub fn build_scalar_tensor_value(cx: &mut Cx, value: Value) -> Result<Value> {
307    build_tensor_value(cx, Vec::new(), None, vec![value])
308}
309
310/// Borrows the [`Tensor`] backing a value, or `None` if it is not a tensor.
311pub fn tensor_value_ref(value: &Value) -> Option<&Tensor> {
312    value.object().downcast_ref::<Tensor>()
313}
314
315/// The shared element number domain (dtype) of a tensor's cells.
316pub fn tensor_dtype(tensor: &Tensor) -> &Symbol {
317    tensor.dtype()
318}
319
320/// Clones a tensor's row-major cell values as a flat vector.
321pub fn flatten_tensor_scalar_cells(tensor: &Tensor) -> Vec<Value> {
322    tensor.data().to_vec()
323}
324
325pub fn tensor_display_name() -> &'static str {
326    "tensor"
327}
328
329fn exprs(cx: &mut Cx, data: &[Value]) -> Result<Vec<Expr>> {
330    data.iter()
331        .map(|value| value.object().as_expr(cx))
332        .collect()
333}
334
335use crate::spec::checked_element_count;
336
337fn validate_shape_and_data_len(shape: &[usize], data_len: usize) -> Result<()> {
338    let expected = checked_element_count(shape)?;
339    if data_len != expected {
340        return Err(Error::Eval(format!(
341            "tensor shape {:?} expects {expected} cells, found {data_len}",
342            shape
343        )));
344    }
345    if data_len == 0 {
346        return Err(Error::Eval("tensor requires at least one cell".to_owned()));
347    }
348    Ok(())
349}
350
351fn validate_cells(cx: &mut Cx, data: &[Value]) -> Result<()> {
352    for cell in data {
353        let Some(number) = cx.number_value_ref(cell.clone())? else {
354            return Err(Error::Eval(
355                "tensor cells must all be scalar number values".to_owned(),
356            ));
357        };
358        if number.domain == number_domain() {
359            return Err(Error::Eval(
360                "tensor cells must be scalar numbers, not nested tensors".to_owned(),
361            ));
362        }
363    }
364    Ok(())
365}
366
367fn validate_dtype_accepts_cells(cx: &mut Cx, dtype: &Symbol, data: &[Value]) -> Result<()> {
368    let domains = cell_domains(cx, data)?;
369    if domains
370        .iter()
371        .all(|domain| promotion_cost(cx, domain, dtype).is_some())
372    {
373        return Ok(());
374    }
375    Err(Error::Eval(format!(
376        "tensor dtype {dtype} is not a valid join for cell domains {domains:?}"
377    )))
378}
379
380fn validate_exact_cell_dtype(dtype: &Symbol, data: &[Value]) -> Result<()> {
381    let mut cx = Cx::new(Arc::new(NoopEvalPolicy), Arc::new(DefaultFactory));
382    for cell in data {
383        let Some(number) = cell.object().as_number_value() else {
384            return Err(Error::Eval(
385                "tensor cells must all be scalar number values".to_owned(),
386            ));
387        };
388        let domain = number.number_domain(&mut cx)?;
389        if domain == number_domain() {
390            return Err(Error::Eval(
391                "tensor cells must be scalar numbers, not nested tensors".to_owned(),
392            ));
393        }
394        if &domain != dtype {
395            return Err(Error::Eval(format!(
396                "tensor dtype {dtype} does not match cell domain {domain}"
397            )));
398        }
399    }
400    Ok(())
401}
402
403fn choose_dtype(cx: &mut Cx, dtype_hint: Option<Symbol>, data: &[Value]) -> Result<Symbol> {
404    if data.is_empty() {
405        return Err(Error::Eval("tensor requires at least one cell".to_owned()));
406    }
407    let domains = cell_domains(cx, data)?;
408    if domains.is_empty() {
409        return Err(Error::Eval("tensor requires at least one cell".to_owned()));
410    }
411    if let Some(dtype) = dtype_hint {
412        if domains
413            .iter()
414            .all(|domain| promotion_cost(cx, domain, &dtype).is_some())
415        {
416            return Ok(dtype);
417        }
418        return Err(Error::Eval(format!(
419            "tensor dtype {dtype} is not a valid join for cell domains {domains:?}"
420        )));
421    }
422    let candidates = cx
423        .registry()
424        .number_domains()
425        .keys()
426        .filter(|symbol| **symbol != number_domain())
427        .cloned()
428        .collect::<Vec<_>>();
429    let mut best = None::<(u32, Symbol)>;
430    for candidate in candidates {
431        let mut total = 0u32;
432        let mut valid = true;
433        for domain in &domains {
434            let Some(cost) = promotion_cost(cx, domain, &candidate) else {
435                valid = false;
436                break;
437            };
438            total += cost;
439        }
440        if !valid {
441            continue;
442        }
443        match &best {
444            Some((best_cost, best_symbol))
445                if total > *best_cost || (total == *best_cost && candidate >= *best_symbol) => {}
446            _ => best = Some((total, candidate)),
447        }
448    }
449    best.map(|(_, symbol)| symbol).ok_or_else(|| {
450        Error::Eval(format!(
451            "no join domain exists for tensor cells {domains:?}"
452        ))
453    })
454}
455
456fn cell_domains(cx: &mut Cx, data: &[Value]) -> Result<Vec<Symbol>> {
457    data.iter()
458        .map(|value| {
459            cx.number_value_ref(value.clone())?
460                .map(|number| number.domain)
461                .ok_or_else(|| {
462                    Error::Eval("tensor cells must all be scalar number values".to_owned())
463                })
464        })
465        .collect()
466}
467
468fn promotion_cost(cx: &Cx, from: &Symbol, to: &Symbol) -> Option<u32> {
469    if from == to {
470        return Some(0);
471    }
472
473    #[derive(Clone, Eq, PartialEq)]
474    struct State {
475        cost: u32,
476        symbol: Symbol,
477    }
478
479    impl Ord for State {
480        fn cmp(&self, other: &Self) -> Ordering {
481            other
482                .cost
483                .cmp(&self.cost)
484                .then_with(|| other.symbol.cmp(&self.symbol))
485        }
486    }
487
488    impl PartialOrd for State {
489        fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
490            Some(self.cmp(other))
491        }
492    }
493
494    let mut best = BTreeMap::<Symbol, u32>::new();
495    let mut heap = BinaryHeap::new();
496    best.insert(from.clone(), 0);
497    heap.push(State {
498        cost: 0,
499        symbol: from.clone(),
500    });
501
502    while let Some(State { cost, symbol }) = heap.pop() {
503        if &symbol == to {
504            return Some(cost);
505        }
506        if best.get(&symbol).copied().unwrap_or(u32::MAX) < cost {
507            continue;
508        }
509        for rule in cx
510            .registry()
511            .value_promotion_rules()
512            .iter()
513            .filter(|rule| rule.from_domain == symbol)
514        {
515            let next = cost + rule.cost as u32;
516            let entry = best.entry(rule.to_domain.clone()).or_insert(u32::MAX);
517            if next < *entry {
518                *entry = next;
519                heap.push(State {
520                    cost: next,
521                    symbol: rule.to_domain.clone(),
522                });
523            }
524        }
525    }
526    None
527}
sim_lib_numbers_tensor/implementation/value.rs

sim_lib_numbers_tensor/implementation/
value.rs
