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#[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
38fn 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
78pub 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}