Skip to main content

runmat_runtime/indexing/
end_expr.rs

1use crate::builtins::common::tensor::{
2    complex_tensor_element_len, complex_tensor_value_complex64, is_scalar_tensor, tensor_value_f64,
3};
4use runmat_types::{CallableFallbackPolicy, CallableIdentity};
5use runmat_value::{IntValue, Value};
6use serde::{Deserialize, Serialize};
7use std::{future::Future, pin::Pin};
8
9/// Executor-neutral representation of `end` arithmetic used by indexing
10/// plans, bytecode, native IR, and semantic call adapters.
11#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
12pub enum EndExpr {
13    End,
14    Const(f64),
15    Var(usize),
16    ResolvedCall {
17        identity: CallableIdentity,
18        fallback_policy: CallableFallbackPolicy,
19        args: Vec<EndExpr>,
20    },
21    Add(Box<EndExpr>, Box<EndExpr>),
22    Sub(Box<EndExpr>, Box<EndExpr>),
23    Mul(Box<EndExpr>, Box<EndExpr>),
24    Div(Box<EndExpr>, Box<EndExpr>),
25    LeftDiv(Box<EndExpr>, Box<EndExpr>),
26    Pow(Box<EndExpr>, Box<EndExpr>),
27    Neg(Box<EndExpr>),
28    Pos(Box<EndExpr>),
29    Floor(Box<EndExpr>),
30    Ceil(Box<EndExpr>),
31    Round(Box<EndExpr>),
32    Fix(Box<EndExpr>),
33}
34
35#[derive(Debug, Clone, Copy)]
36pub struct ValueToF64Error;
37
38/// Converts an integer only when the resulting double retains its exact value.
39/// `end` expressions are subsequently used as indices, so accepting a rounded
40/// integer here would select a different element than the source expression.
41fn exact_integer_to_f64(value: &IntValue) -> Result<f64, ValueToF64Error> {
42    let converted = value.to_f64();
43    let exact = match value {
44        IntValue::I8(value) => converted as i128 == i128::from(*value),
45        IntValue::I16(value) => converted as i128 == i128::from(*value),
46        IntValue::I32(value) => converted as i128 == i128::from(*value),
47        IntValue::I64(value) => converted as i128 == i128::from(*value),
48        IntValue::U8(value) => converted as u128 == u128::from(*value),
49        IntValue::U16(value) => converted as u128 == u128::from(*value),
50        IntValue::U32(value) => converted as u128 == u128::from(*value),
51        IntValue::U64(value) => converted as u128 == u128::from(*value),
52    };
53    exact.then_some(converted).ok_or(ValueToF64Error)
54}
55
56pub fn value_to_f64(v: &Value) -> Result<f64, ValueToF64Error> {
57    match v {
58        Value::Num(n) => Ok(*n),
59        Value::Int(i) => exact_integer_to_f64(i),
60        Value::Bool(b) => Ok(if *b { 1.0 } else { 0.0 }),
61        Value::Tensor(t) if is_scalar_tensor(t) => match t.integer_storage() {
62            Some(storage) => exact_integer_to_f64(&storage.value_at(0).ok_or(ValueToF64Error)?),
63            None => Ok(tensor_value_f64(t, 0)),
64        },
65        Value::Complex(re, im) if im.abs() < 1e-12 => Ok(*re),
66        Value::ComplexTensor(ct) if complex_tensor_element_len(ct) == 1 => {
67            let value = complex_tensor_value_complex64(ct, 0);
68            if value.im.abs() < 1e-12 {
69                Ok(value.re)
70            } else {
71                Err(ValueToF64Error)
72            }
73        }
74        _ => Err(ValueToF64Error),
75    }
76}
77
78/// Resolve one context-dependent indexing expression through caller-owned
79/// local storage and the canonical runtime callable path.
80pub async fn resolve_end_expr_value<F>(
81    dimension_length: usize,
82    expression: &EndExpr,
83    resolve_variable: F,
84) -> Result<f64, crate::RuntimeError>
85where
86    F: Fn(usize) -> Option<Value> + Copy,
87{
88    fn evaluate<'a, F>(
89        expression: &'a EndExpr,
90        end_value: f64,
91        resolve_variable: F,
92    ) -> Pin<Box<dyn Future<Output = Result<f64, crate::RuntimeError>> + 'a>>
93    where
94        F: Fn(usize) -> Option<Value> + Copy + 'a,
95    {
96        Box::pin(async move {
97            let invalid =
98                |identifier, message| crate::runtime_error::semantic_error(identifier, message);
99            match expression {
100                EndExpr::End => Ok(end_value),
101                EndExpr::Const(value) => Ok(*value),
102                EndExpr::Var(local) => {
103                    let mut value = resolve_variable(*local).ok_or_else(|| {
104                        invalid("MissingNumericIndex", "missing variable for end expression")
105                    })?;
106                    if matches!(value, Value::GpuTensor(_)) {
107                        value = crate::dispatcher::gather_if_needed_async(&value).await?;
108                    }
109                    value_to_f64(&value).map_err(|_| {
110                        invalid("UnsupportedIndexType", "end expression must be numeric")
111                    })
112                }
113                EndExpr::ResolvedCall {
114                    identity,
115                    fallback_policy,
116                    args,
117                } => {
118                    let mut arguments = Vec::with_capacity(args.len());
119                    for argument in args {
120                        arguments.push(Value::Num(
121                            evaluate(argument, end_value, resolve_variable).await?,
122                        ));
123                    }
124                    let descriptor = crate::call::descriptor::CallableDescriptor::resolved(
125                        identity.clone(),
126                        arguments,
127                        1,
128                        *fallback_policy,
129                        crate::call::descriptor::CallableCallKind::EndExpr,
130                    );
131                    let value =
132                        crate::call::descriptor::execute_callable_descriptor(descriptor).await?;
133                    let value = match value {
134                        Value::OutputList(mut values) if values.len() == 1 => values.remove(0),
135                        value => value,
136                    };
137                    value_to_f64(&value)
138                        .map_err(|_| invalid("UnsupportedIndexType", "end call must return scalar"))
139                }
140                EndExpr::Add(left, right) => Ok(evaluate(left, end_value, resolve_variable)
141                    .await?
142                    + evaluate(right, end_value, resolve_variable).await?),
143                EndExpr::Sub(left, right) => Ok(evaluate(left, end_value, resolve_variable)
144                    .await?
145                    - evaluate(right, end_value, resolve_variable).await?),
146                EndExpr::Mul(left, right) => Ok(evaluate(left, end_value, resolve_variable)
147                    .await?
148                    * evaluate(right, end_value, resolve_variable).await?),
149                EndExpr::Div(left, right) => {
150                    let denominator = evaluate(right, end_value, resolve_variable).await?;
151                    if denominator == 0.0 {
152                        return Err(invalid("IndexOutOfBounds", "Index out of bounds"));
153                    }
154                    Ok(evaluate(left, end_value, resolve_variable).await? / denominator)
155                }
156                EndExpr::LeftDiv(left, right) => {
157                    let denominator = evaluate(left, end_value, resolve_variable).await?;
158                    if denominator == 0.0 {
159                        return Err(invalid("IndexOutOfBounds", "Index out of bounds"));
160                    }
161                    Ok(evaluate(right, end_value, resolve_variable).await? / denominator)
162                }
163                EndExpr::Pow(left, right) => Ok(evaluate(left, end_value, resolve_variable)
164                    .await?
165                    .powf(evaluate(right, end_value, resolve_variable).await?)),
166                EndExpr::Neg(inner) => Ok(-evaluate(inner, end_value, resolve_variable).await?),
167                EndExpr::Pos(inner) => evaluate(inner, end_value, resolve_variable).await,
168                EndExpr::Floor(inner) => {
169                    Ok(evaluate(inner, end_value, resolve_variable).await?.floor())
170                }
171                EndExpr::Ceil(inner) => {
172                    Ok(evaluate(inner, end_value, resolve_variable).await?.ceil())
173                }
174                EndExpr::Round(inner) => {
175                    Ok(evaluate(inner, end_value, resolve_variable).await?.round())
176                }
177                EndExpr::Fix(inner) => {
178                    let value = evaluate(inner, end_value, resolve_variable).await?;
179                    Ok(if value >= 0.0 {
180                        value.floor()
181                    } else {
182                        value.ceil()
183                    })
184                }
185            }
186        })
187    }
188
189    evaluate(expression, dimension_length as f64, resolve_variable).await
190}
191
192#[cfg(test)]
193mod tests {
194    use super::{resolve_end_expr_value, value_to_f64, EndExpr};
195    use runmat_value::{IntValue, IntegerStorage, Tensor, Value};
196
197    #[test]
198    fn value_to_f64_reads_all_typed_integer_tensor_classes_without_f64_mirrors() {
199        macro_rules! assert_typed_scalar {
200            ($storage:expr, $expected:expr) => {{
201                let tensor = Tensor::new_integer($storage, vec![1, 1]).expect("scalar tensor");
202                assert_eq!(value_to_f64(&Value::Tensor(tensor)).unwrap(), $expected);
203            }};
204        }
205
206        assert_typed_scalar!(IntegerStorage::I8(vec![-8]), -8.0);
207        assert_typed_scalar!(IntegerStorage::I16(vec![-16]), -16.0);
208        assert_typed_scalar!(IntegerStorage::I32(vec![-32]), -32.0);
209        assert_typed_scalar!(IntegerStorage::I64(vec![-64]), -64.0);
210        assert_typed_scalar!(IntegerStorage::U8(vec![8]), 8.0);
211        assert_typed_scalar!(IntegerStorage::U16(vec![16]), 16.0);
212        assert_typed_scalar!(IntegerStorage::U32(vec![32]), 32.0);
213        assert_typed_scalar!(IntegerStorage::U64(vec![64]), 64.0);
214    }
215
216    #[test]
217    fn value_to_f64_rejects_wide_integer_values_that_would_round_indices() {
218        for value in [
219            IntValue::I64(i64::MAX),
220            IntValue::U64((1_u64 << 53) + 1),
221            IntValue::U64(u64::MAX),
222        ] {
223            assert!(value_to_f64(&Value::Int(value)).is_err());
224        }
225
226        let tensor = Tensor::new_integer(IntegerStorage::U64(vec![(1_u64 << 53) + 1]), vec![1, 1])
227            .expect("scalar tensor");
228        assert!(value_to_f64(&Value::Tensor(tensor)).is_err());
229    }
230
231    #[test]
232    fn shared_end_expression_evaluator_uses_base_extent_and_executor_locals() {
233        let expression = EndExpr::Sub(Box::new(EndExpr::End), Box::new(EndExpr::Var(2)));
234        let value = futures::executor::block_on(resolve_end_expr_value(8, &expression, |local| {
235            (local == 2).then_some(Value::Int(IntValue::U8(3)))
236        }))
237        .unwrap();
238        assert_eq!(value, 5.0);
239    }
240}