runmat_runtime/builtins/math/poly/
roots.rs

1//! MATLAB-compatible `roots` builtin with GPU-aware semantics for RunMat.
2//!
3//! This implementation mirrors MATLAB behaviour, including handling for leading
4//! zeros, constant polynomials, and complex-valued coefficients. GPU inputs are
5//! gathered to the host because companion matrix eigenvalue computations are
6//! currently performed on the CPU.
7
8use nalgebra::DMatrix;
9use num_complex::Complex64;
10use runmat_builtins::{
11    BuiltinCompletionPolicy, BuiltinDescriptor, BuiltinErrorDescriptor, BuiltinExtensionDescriptor,
12    BuiltinExtensionMode, BuiltinIntegerBackendRule, BuiltinIntegerCapabilityDescriptor,
13    BuiltinIntegerComputationDomain, BuiltinIntegerInputAvailability,
14    BuiltinIntegerInputCapability, BuiltinIntegerOutputClassRule, BuiltinIntegerOverflowRule,
15    BuiltinIntegerOverloadKind, BuiltinIntegerScalarDoubleRule, BuiltinOutputMode,
16    BuiltinParamArity, BuiltinParamDescriptor, BuiltinParamType, BuiltinSignatureDescriptor,
17};
18use runmat_macros::runtime_builtin;
19use runmat_value::{ComplexTensor, Tensor, Value};
20
21use crate::builtins::common::spec::{
22    BroadcastSemantics, BuiltinFusionSpec, BuiltinGpuSpec, ConstantStrategy, GpuOpKind,
23    ReductionNaN, ResidencyPolicy, ShapeRequirements,
24};
25use crate::builtins::common::{gpu_helpers, tensor};
26use crate::builtins::math::poly::type_resolvers::roots_type;
27use crate::{build_runtime_error, BuiltinResult, RuntimeError};
28
29const LEADING_ZERO_TOL: f64 = 1.0e-12;
30const RESULT_ZERO_TOL: f64 = 1.0e-10;
31const BUILTIN_NAME: &str = "roots";
32
33const ROOTS_OUTPUT: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
34    name: "r",
35    ty: BuiltinParamType::Any,
36    arity: BuiltinParamArity::Required,
37    default: None,
38    description: "Roots of the polynomial as a column vector.",
39}];
40
41const ROOTS_INPUTS: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
42    name: "c",
43    ty: BuiltinParamType::Any,
44    arity: BuiltinParamArity::Required,
45    default: None,
46    description: "Polynomial coefficient vector in descending power order.",
47}];
48
49const ROOTS_SIGNATURES: [BuiltinSignatureDescriptor; 1] = [BuiltinSignatureDescriptor {
50    label: "r = roots(c)",
51    inputs: &ROOTS_INPUTS,
52    outputs: &ROOTS_OUTPUT,
53}];
54
55const ROOTS_ERROR_INVALID_INPUT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
56    code: "RM.ROOTS.INVALID_INPUT",
57    identifier: Some("RunMat:roots:InvalidInput"),
58    when: "Input cannot be interpreted as a numeric coefficient vector.",
59    message: "roots: invalid input",
60};
61
62const ROOTS_ERROR_INTERNAL: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
63    code: "RM.ROOTS.INTERNAL",
64    identifier: Some("RunMat:roots:Internal"),
65    when: "Runtime fails while building companion matrix outputs or solving eigenvalues.",
66    message: "roots: internal runtime failure",
67};
68
69const ROOTS_ERRORS: [BuiltinErrorDescriptor; 2] = [ROOTS_ERROR_INVALID_INPUT, ROOTS_ERROR_INTERNAL];
70
71pub const ROOTS_DESCRIPTOR: BuiltinDescriptor = BuiltinDescriptor {
72    signatures: &ROOTS_SIGNATURES,
73    output_mode: BuiltinOutputMode::Fixed,
74    completion_policy: BuiltinCompletionPolicy::Public,
75    errors: &ROOTS_ERRORS,
76};
77
78const ROOTS_INTEGER_COEFFICIENTS_EXTENSION: BuiltinExtensionDescriptor =
79    BuiltinExtensionDescriptor {
80        id: "roots-integer-coefficients",
81        mode: BuiltinExtensionMode::RunMatOnly,
82        description: "roots accepts typed-integer polynomial coefficients as a RunMat extension",
83        error_identifier: Some("RunMat:compatibility:RootsIntegerCoefficientsExtension"),
84    };
85pub const ROOTS_EXTENSIONS: [BuiltinExtensionDescriptor; 1] =
86    [ROOTS_INTEGER_COEFFICIENTS_EXTENSION];
87const ROOTS_INTEGER_INPUTS: [BuiltinIntegerInputCapability; 1] =
88    [BuiltinIntegerInputCapability {
89        name: "c",
90        classes: &crate::builtins::common::integer_capability::ALL_INTEGER_CLASSES,
91        availability: BuiltinIntegerInputAvailability::RunMatOnly,
92        scalar_double: BuiltinIntegerScalarDoubleRule::NotApplicable,
93        notes: "The compatibility target documents single and double coefficients; RunMat admits typed integers only when every coefficient is exactly representable in binary64.",
94    }];
95pub const ROOTS_INTEGER_CAPABILITIES: [BuiltinIntegerCapabilityDescriptor; 1] =
96    [BuiltinIntegerCapabilityDescriptor {
97        form: "r = roots(integer_c)",
98        inputs: &ROOTS_INTEGER_INPUTS,
99        computation_domain: BuiltinIntegerComputationDomain::FloatingPoint,
100        output_class: BuiltinIntegerOutputClassRule::Double,
101        overflow: BuiltinIntegerOverflowRule::Error,
102        backend: BuiltinIntegerBackendRule::GatherFallback,
103        overload: BuiltinIntegerOverloadKind::Multiple,
104        notes: "The checked extension crosses once into the double companion-matrix/eigenvalue algorithm. Automatic residency gathers through the owning provider; explicit typed-integer input is gated before provider access.",
105    }];
106
107#[runmat_macros::register_gpu_spec(builtin_path = "crate::builtins::math::poly::roots")]
108pub const GPU_SPEC: BuiltinGpuSpec = BuiltinGpuSpec {
109    name: "roots",
110    op_kind: GpuOpKind::Custom("polynomial-roots"),
111    supported_precisions: &[],
112    broadcast: BroadcastSemantics::None,
113    provider_hooks: &[],
114    constant_strategy: ConstantStrategy::InlineLiteral,
115    residency: ResidencyPolicy::GatherImmediately,
116    nan_mode: ReductionNaN::Include,
117    two_pass_threshold: None,
118    workgroup_size: None,
119    accepts_nan_mode: false,
120    notes: "Companion matrix eigenvalue solve executes on the host; providers currently fall back to the CPU implementation.",
121};
122
123fn roots_error(message: impl Into<String>) -> RuntimeError {
124    roots_error_with(message, &ROOTS_ERROR_INVALID_INPUT)
125}
126
127fn roots_error_with(
128    message: impl Into<String>,
129    error: &'static BuiltinErrorDescriptor,
130) -> RuntimeError {
131    let mut builder = build_runtime_error(message).with_builtin(BUILTIN_NAME);
132    if let Some(identifier) = error.identifier {
133        builder = builder.with_identifier(identifier);
134    }
135    builder.build()
136}
137
138#[runmat_macros::register_fusion_spec(builtin_path = "crate::builtins::math::poly::roots")]
139pub const FUSION_SPEC: BuiltinFusionSpec = BuiltinFusionSpec {
140    name: "roots",
141    shape: ShapeRequirements::Any,
142    constant_strategy: ConstantStrategy::InlineLiteral,
143    elementwise: None,
144    reduction: None,
145    emits_nan: true,
146    notes: "Non-elementwise builtin that terminates fusion and gathers inputs to the host.",
147};
148
149#[runtime_builtin(
150    name = "roots",
151    category = "math/poly",
152    summary = "Compute polynomial roots from a coefficient vector.",
153    keywords = "roots,polynomial,eigenvalues,companion",
154    accel = "sink",
155    type_resolver(roots_type),
156    descriptor(crate::builtins::math::poly::roots::ROOTS_DESCRIPTOR),
157    extensions(crate::builtins::math::poly::roots::ROOTS_EXTENSIONS),
158    integer_capabilities(crate::builtins::math::poly::roots::ROOTS_INTEGER_CAPABILITIES),
159    builtin_path = "crate::builtins::math::poly::roots"
160)]
161async fn roots_builtin(coefficients: Value) -> crate::BuiltinResult<Value> {
162    roots_value(coefficients).await
163}
164
165pub(crate) async fn roots_value(coefficients: Value) -> crate::BuiltinResult<Value> {
166    crate::builtins::common::validation::reject_typed_complex_integer(&coefficients, BUILTIN_NAME)?;
167    crate::builtins::common::validation::ensure_runmat_integer_f64_boundary(
168        &coefficients,
169        &ROOTS_INTEGER_COEFFICIENTS_EXTENSION,
170        BUILTIN_NAME,
171        "coefficient",
172    )
173    .await?;
174    let coeffs = coefficients_to_complex(coefficients).await?;
175    let trimmed = trim_leading_zeros(coeffs);
176    if trimmed.is_empty() || trimmed.len() == 1 {
177        return empty_column();
178    }
179    let roots = solve_roots(&trimmed)?;
180    roots_to_value(&roots)
181}
182
183async fn coefficients_to_complex(value: Value) -> BuiltinResult<Vec<Complex64>> {
184    match value {
185        Value::GpuTensor(handle) => {
186            let tensor = gpu_helpers::gather_tensor_async(&handle).await?;
187            tensor_to_complex(tensor)
188        }
189        Value::Tensor(tensor) => tensor_to_complex(tensor),
190        Value::ComplexTensor(tensor) => complex_tensor_to_vec(tensor),
191        Value::LogicalArray(logical) => {
192            let tensor = tensor::logical_to_tensor(&logical).map_err(roots_error)?;
193            tensor_to_complex(tensor)
194        }
195        Value::Num(n) => {
196            let tensor =
197                Tensor::new(vec![n], vec![1, 1]).map_err(|e| roots_error(format!("roots: {e}")))?;
198            tensor_to_complex(tensor)
199        }
200        Value::Int(i) => {
201            let tensor = Tensor::new(vec![i.to_f64()], vec![1, 1])
202                .map_err(|e| roots_error(format!("roots: {e}")))?;
203            tensor_to_complex(tensor)
204        }
205        Value::Bool(b) => {
206            let tensor = Tensor::new(vec![if b { 1.0 } else { 0.0 }], vec![1, 1])
207                .map_err(|e| roots_error(format!("roots: {e}")))?;
208            tensor_to_complex(tensor)
209        }
210        other => Err(roots_error(format!(
211            "roots: expected a numeric vector of polynomial coefficients, got {other:?}"
212        ))),
213    }
214}
215
216fn tensor_to_complex(tensor: Tensor) -> BuiltinResult<Vec<Complex64>> {
217    ensure_vector_shape("roots", &tensor.shape)?;
218    Ok(tensor::tensor_values_f64(&tensor)
219        .into_iter()
220        .map(|value| Complex64::new(value, 0.0))
221        .collect())
222}
223
224fn complex_tensor_to_vec(tensor: ComplexTensor) -> BuiltinResult<Vec<Complex64>> {
225    ensure_vector_shape("roots", &tensor.shape)?;
226    Ok(tensor
227        .materialize_f64()
228        .into_iter()
229        .map(|(re, im)| Complex64::new(re, im))
230        .collect())
231}
232
233fn ensure_vector_shape(name: &str, shape: &[usize]) -> BuiltinResult<()> {
234    let is_vector = match shape.len() {
235        0 => true,
236        1 => true,
237        2 => shape[0] == 1 || shape[1] == 1 || shape.iter().product::<usize>() == 0,
238        _ => shape.iter().filter(|&&dim| dim > 1).count() <= 1,
239    };
240    if !is_vector {
241        return Err(roots_error(format!(
242            "{name}: coefficients must be a vector (row or column), got shape {:?}",
243            shape
244        )));
245    }
246    Ok(())
247}
248
249fn trim_leading_zeros(mut coeffs: Vec<Complex64>) -> Vec<Complex64> {
250    if coeffs.is_empty() {
251        return coeffs;
252    }
253    let scale = coeffs.iter().map(|c| c.norm()).fold(0.0_f64, f64::max);
254    let tol = if scale == 0.0 {
255        LEADING_ZERO_TOL
256    } else {
257        LEADING_ZERO_TOL * scale
258    };
259    let first_nonzero = coeffs
260        .iter()
261        .position(|c| c.norm() > tol)
262        .unwrap_or(coeffs.len());
263    coeffs.split_off(first_nonzero)
264}
265
266fn solve_roots(coeffs: &[Complex64]) -> BuiltinResult<Vec<Complex64>> {
267    if coeffs.len() <= 1 {
268        return Ok(Vec::new());
269    }
270    if coeffs.len() == 2 {
271        let a = coeffs[0];
272        let b = coeffs[1];
273        if a.norm() <= LEADING_ZERO_TOL {
274            return Err(roots_error(
275                "roots: leading coefficient must be non-zero after trimming",
276            ));
277        }
278        return Ok(vec![-b / a]);
279    }
280
281    let degree = coeffs.len() - 1;
282    if degree == 3 {
283        return Ok(cubic_roots(coeffs[0], coeffs[1], coeffs[2], coeffs[3]));
284    }
285    let leading = coeffs[0];
286    if leading.norm() <= LEADING_ZERO_TOL {
287        return Err(roots_error(
288            "roots: leading coefficient must be non-zero after trimming",
289        ));
290    }
291
292    let mut companion = DMatrix::<Complex64>::zeros(degree, degree);
293    for row in 1..degree {
294        companion[(row, row - 1)] = Complex64::new(1.0, 0.0);
295    }
296
297    for (idx, coeff) in coeffs.iter().enumerate().skip(1) {
298        let value = -(*coeff) / leading;
299        let column = idx - 1;
300        if column < degree {
301            companion[(0, column)] = value;
302        }
303    }
304
305    let eigenvalues = companion.clone().eigenvalues().ok_or_else(|| {
306        roots_error_with(
307            "roots: failed to compute eigenvalues of the companion matrix",
308            &ROOTS_ERROR_INTERNAL,
309        )
310    })?;
311    Ok(eigenvalues.iter().map(|&z| canonicalize_root(z)).collect())
312}
313
314fn cubic_roots(a: Complex64, b: Complex64, c: Complex64, d: Complex64) -> Vec<Complex64> {
315    // Depressed cubic via Cardano: x = y - b/(3a), y^3 + p y + q = 0
316    let three = 3.0;
317    let nine = 9.0;
318    let twenty_seven = 27.0;
319    let a2 = a * a;
320    let a3 = a2 * a;
321    let p = (three * a * c - b * b) / (three * a2);
322    let q = (twenty_seven * a2 * d - nine * a * b * c + Complex64::new(2.0, 0.0) * b * b * b)
323        / (twenty_seven * a3);
324    let half = Complex64::new(0.5, 0.0);
325    let disc = (q * q) * half * half + (p * p * p) / Complex64::new(27.0, 0.0);
326    let sqrt_disc = disc.sqrt();
327    let u = (-q * half + sqrt_disc).powf(1.0 / 3.0);
328    let v = (-q * half - sqrt_disc).powf(1.0 / 3.0);
329    let omega = Complex64::new(-0.5, (3.0f64).sqrt() * 0.5);
330    let omega2 = omega * omega;
331    let shift = b / (three * a);
332    let y0 = u + v;
333    let y1 = u * omega + v * omega.conj();
334    let y2 = u * omega2 + v * omega;
335    vec![y0 - shift, y1 - shift, y2 - shift]
336}
337
338fn canonicalize_root(z: Complex64) -> Complex64 {
339    if !z.re.is_finite() || !z.im.is_finite() {
340        return z;
341    }
342    let mut real = z.re;
343    let mut imag = z.im;
344    let scale = 1.0 + real.abs();
345    if imag.abs() <= RESULT_ZERO_TOL * scale {
346        imag = 0.0;
347    }
348    if real.abs() <= RESULT_ZERO_TOL {
349        real = 0.0;
350    }
351    Complex64::new(real, imag)
352}
353
354fn roots_to_value(roots: &[Complex64]) -> BuiltinResult<Value> {
355    if roots.is_empty() {
356        return empty_column();
357    }
358    let all_real = roots
359        .iter()
360        .all(|z| z.im.abs() <= RESULT_ZERO_TOL * (1.0 + z.re.abs()));
361    if all_real {
362        let mut data: Vec<f64> = Vec::with_capacity(roots.len());
363        for &root in roots {
364            data.push(root.re);
365        }
366        let tensor = Tensor::new(data, vec![roots.len(), 1])
367            .map_err(|e| roots_error_with(format!("roots: {e}"), &ROOTS_ERROR_INTERNAL))?;
368        Ok(Value::Tensor(tensor))
369    } else {
370        let data: Vec<(f64, f64)> = roots.iter().map(|z| (z.re, z.im)).collect();
371        let tensor = ComplexTensor::new(data, vec![roots.len(), 1])
372            .map_err(|e| roots_error_with(format!("roots: {e}"), &ROOTS_ERROR_INTERNAL))?;
373        Ok(Value::ComplexTensor(tensor))
374    }
375}
376
377fn empty_column() -> BuiltinResult<Value> {
378    let tensor = Tensor::new(Vec::new(), vec![0, 1])
379        .map_err(|e| roots_error_with(format!("roots: {e}"), &ROOTS_ERROR_INTERNAL))?;
380    Ok(Value::Tensor(tensor))
381}
382
383#[cfg(test)]
384pub(crate) mod tests {
385    use super::*;
386    use crate::builtins::common::test_support;
387    use futures::executor::block_on;
388    use runmat_accelerate_api::HostTensorView;
389    use runmat_value::{ComplexTensor, IntegerStorage, LogicalArray, Tensor};
390
391    fn assert_error_contains(err: crate::RuntimeError, needle: &str) {
392        assert!(
393            err.message().contains(needle),
394            "expected error containing '{needle}', got '{}'",
395            err.message()
396        );
397    }
398
399    #[test]
400    fn roots_descriptor_signatures_cover_core_forms() {
401        let labels: Vec<&str> = ROOTS_DESCRIPTOR
402            .signatures
403            .iter()
404            .map(|signature| signature.label)
405            .collect();
406        assert!(labels.contains(&"r = roots(c)"));
407    }
408
409    #[test]
410    fn roots_descriptor_errors_have_stable_codes() {
411        let codes: Vec<&str> = ROOTS_DESCRIPTOR
412            .errors
413            .iter()
414            .map(|error| error.code)
415            .collect();
416        assert!(codes.contains(&"RM.ROOTS.INVALID_INPUT"));
417        assert!(codes.contains(&"RM.ROOTS.INTERNAL"));
418    }
419
420    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
421    #[test]
422    fn roots_quadratic_real() {
423        let coeffs = Tensor::new(vec![1.0, -3.0, 2.0], vec![3, 1]).unwrap();
424        let result = roots_builtin(Value::Tensor(coeffs)).expect("roots");
425        match result {
426            Value::Tensor(t) => {
427                assert_eq!(t.shape, vec![2, 1]);
428                let mut roots = t.materialize_f64();
429                roots.sort_by(|a, b| a.partial_cmp(b).unwrap());
430                assert!((roots[0] - 1.0).abs() < 1e-10);
431                assert!((roots[1] - 2.0).abs() < 1e-10);
432            }
433            other => panic!("expected real tensor, got {other:?}"),
434        }
435    }
436
437    #[test]
438    fn roots_typed_integer_coefficients_cross_double_boundary_exactly() {
439        let _compat = crate::compatibility::push_runmat_extensions_enabled(true);
440        let coeffs = Tensor::new_integer(IntegerStorage::I16(vec![1, -3, 2]), vec![3, 1]).unwrap();
441        let result = roots_builtin(Value::Tensor(coeffs)).expect("roots");
442        match result {
443            Value::Tensor(t) => {
444                assert_eq!(t.shape, vec![2, 1]);
445                let mut roots = t.materialize_f64();
446                roots.sort_by(|a, b| a.partial_cmp(b).unwrap());
447                assert!((roots[0] - 1.0).abs() < 1e-10);
448                assert!((roots[1] - 2.0).abs() < 1e-10);
449            }
450            other => panic!("expected real tensor, got {other:?}"),
451        }
452    }
453
454    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
455    #[test]
456    fn roots_leading_zeros_trimmed() {
457        let coeffs = Tensor::new(vec![0.0, 0.0, 1.0, -4.0], vec![4, 1]).unwrap();
458        let result = roots_builtin(Value::Tensor(coeffs)).expect("roots");
459        match result {
460            Value::Tensor(t) => {
461                assert_eq!(t.shape, vec![1, 1]);
462                assert!((t.materialize_f64()[0] - 4.0).abs() < 1e-10);
463            }
464            other => panic!("expected tensor, got {other:?}"),
465        }
466    }
467
468    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
469    #[test]
470    fn roots_complex_pair() {
471        let coeffs = Tensor::new(vec![1.0, 0.0, 1.0], vec![3, 1]).unwrap();
472        let result = roots_builtin(Value::Tensor(coeffs)).expect("roots");
473        match result {
474            Value::ComplexTensor(t) => {
475                assert_eq!(t.shape, vec![2, 1]);
476                let mut roots = t.materialize_f64();
477                roots.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
478                assert!((roots[0].0).abs() < 1e-10);
479                assert!((roots[0].1 + 1.0).abs() < 1e-10);
480                assert!((roots[1].0).abs() < 1e-10);
481                assert!((roots[1].1 - 1.0).abs() < 1e-10);
482            }
483            other => panic!("expected complex tensor, got {other:?}"),
484        }
485    }
486
487    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
488    #[test]
489    fn roots_quartic_all_zero_roots() {
490        // p(x) = x^4 => 4 roots at 0
491        let coeffs = Tensor::new(vec![1.0, 0.0, 0.0, 0.0, 0.0], vec![5, 1]).unwrap();
492        let result = roots_builtin(Value::Tensor(coeffs)).expect("roots quartic");
493        match result {
494            Value::Tensor(t) => {
495                assert_eq!(t.shape, vec![4, 1]);
496                for &r in &t.materialize_f64() {
497                    assert!(r.abs() < 1e-8);
498                }
499            }
500            Value::ComplexTensor(t) => {
501                assert_eq!(t.shape, vec![4, 1]);
502                for &(re, im) in &t.materialize_f64() {
503                    assert!(re.abs() < 1e-7 && im.abs() < 1e-7);
504                }
505            }
506            other => panic!("unexpected output {other:?}"),
507        }
508    }
509
510    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
511    #[test]
512    fn roots_accepts_complex_coefficients_input() {
513        // p(x) = x^2 + 1 with complex coefficients path
514        let coeffs =
515            ComplexTensor::new(vec![(1.0, 0.0), (0.0, 0.0), (1.0, 0.0)], vec![3, 1]).unwrap();
516        let result = roots_builtin(Value::ComplexTensor(coeffs)).expect("roots complex input");
517        match result {
518            Value::ComplexTensor(t) => {
519                assert_eq!(t.shape, vec![2, 1]);
520                // roots at i and -i
521                let mut roots = t.materialize_f64();
522                roots.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
523                assert!(roots[0].0.abs() < 1e-10 && (roots[0].1 + 1.0).abs() < 1e-6);
524                assert!(roots[1].0.abs() < 1e-10 && (roots[1].1 - 1.0).abs() < 1e-6);
525            }
526            other => panic!("expected complex tensor, got {other:?}"),
527        }
528    }
529
530    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
531    #[test]
532    fn roots_accepts_logical_coefficients() {
533        // p(x) = x with logical coefficients [1 0]
534        let la = LogicalArray::new(vec![1, 0], vec![1, 2]).unwrap();
535        let result = roots_builtin(Value::LogicalArray(la)).expect("roots logical");
536        match result {
537            Value::Tensor(t) => {
538                assert_eq!(t.shape, vec![1, 1]);
539                assert!(t.materialize_f64()[0].abs() < 1e-12);
540            }
541            other => panic!("expected real tensor, got {other:?}"),
542        }
543    }
544
545    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
546    #[test]
547    fn roots_scalar_num_returns_empty() {
548        let result = roots_builtin(Value::Num(5.0)).expect("roots scalar num");
549        match result {
550            Value::Tensor(t) => {
551                assert_eq!(t.shape, vec![0, 1]);
552                assert!(t.materialize_f64().is_empty());
553            }
554            other => panic!("expected empty tensor, got {other:?}"),
555        }
556    }
557
558    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
559    #[test]
560    fn roots_rejects_non_vector_input() {
561        let coeffs = Tensor::new(vec![1.0, 0.0, 0.0, 1.0], vec![2, 2]).unwrap();
562        let err = roots_builtin(Value::Tensor(coeffs)).expect_err("expected vector-shape error");
563        assert_eq!(err.identifier(), ROOTS_ERROR_INVALID_INPUT.identifier);
564        assert_error_contains(err, "vector");
565    }
566
567    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
568    #[test]
569    fn roots_all_zero_coefficients_returns_empty() {
570        let coeffs = Tensor::new(vec![0.0, 0.0, 0.0], vec![3, 1]).unwrap();
571        let result = roots_builtin(Value::Tensor(coeffs)).expect("roots");
572        match result {
573            Value::Tensor(t) => {
574                assert_eq!(t.shape, vec![0, 1]);
575                assert!(t.materialize_f64().is_empty());
576            }
577            other => panic!("expected empty tensor, got {other:?}"),
578        }
579    }
580
581    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
582    #[test]
583    fn roots_gpu_input_gathers_to_host() {
584        test_support::with_test_provider(|provider| {
585            let coeffs = Tensor::new(vec![1.0, 0.0, -9.0, 0.0], vec![4, 1]).unwrap();
586            let view = HostTensorView {
587                data: &coeffs.materialize_f64(),
588                shape: &coeffs.shape,
589            };
590            let handle = provider.upload(&view).expect("upload");
591            let result = roots_builtin(Value::GpuTensor(handle)).expect("roots");
592            let gathered = test_support::gather(result).expect("gather");
593            assert_eq!(gathered.shape, vec![3, 1]);
594            let mut roots = gathered.materialize_f64();
595            roots.sort_by(|a, b| a.partial_cmp(b).unwrap());
596            assert!((roots[0] + 3.0).abs() < 1e-9);
597            assert!((roots[1]).abs() < 1e-9);
598            assert!((roots[2] - 3.0).abs() < 1e-9);
599        });
600    }
601
602    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
603    #[test]
604    fn roots_constant_polynomial_returns_empty() {
605        let coeffs = Tensor::new(vec![5.0], vec![1, 1]).unwrap();
606        let result = roots_builtin(Value::Tensor(coeffs)).expect("roots");
607        match result {
608            Value::Tensor(t) => {
609                assert_eq!(t.shape, vec![0, 1]);
610            }
611            other => panic!("expected empty tensor, got {other:?}"),
612        }
613    }
614
615    fn roots_builtin(coefficients: Value) -> BuiltinResult<Value> {
616        block_on(super::roots_builtin(coefficients))
617    }
618}
runmat_runtime/builtins/math/poly/roots.rs

runmat_runtime/builtins/math/poly/
roots.rs
