outram_foam_basic_lib/polynomial/

roots.rs

// SPDX-License-Identifier: GPL-3.0-only
// Copyright (C) 2026 OUTRAM PARK contributors
// Derived from OpenFOAM (www.openfoam.com)
// Copyright (C) 2004-2023 OpenFOAM Foundation
// Copyright (C) 2016-2023 OpenCFD Ltd.
//
// This file is part of OUTRAM PARK.
//
// OUTRAM PARK is free software: you can redistribute it and/or modify it
// under the terms of the GNU General Public License as published by the
// Free Software Foundation, either version 3 of the License, or (at your
// option) any later version.
//
// OUTRAM PARK is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// General Public License for more details.
//
// You should have received a copy of the GNU General Public License along
// with OUTRAM PARK.  If not, see <https://www.gnu.org/licenses/>.

/// Root classification tag, matching `Foam::roots::type`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u64)]
pub enum RootType {
    Real = 0,
    Complex = 1,
    PosInf = 2,
    NegInf = 3,
    Nan = 4,
}

/// Tagged root container for N roots.
/// Types are packed 3 bits per root into a u64, identical to C++ `Roots<N>`.
/// For complex conjugate pairs, slot i holds the real part and slot i+1 holds
/// the imaginary part; both slots are tagged `Complex`.
#[derive(Debug, Clone, Copy)]
pub struct Roots<const N: usize> {
    pub(crate) values: [f64; N],
    types: u64,
}

impl<const N: usize> Roots<N> {
    /// Value stored at slot `i`.
    ///
    /// # Out of range reads `NaN`, and never panics
    ///
    /// `i >= N` is a caller bug either way. Returning `NaN` rather than
    /// panicking is what lets this type be lifted verbatim into `petir`, which
    /// targets bare metal -- there a panic is not a stack trace and a non-zero
    /// exit, it is the end of the program. `NaN` still surfaces the bug, just
    /// downstream of the read.
    #[inline]
    pub fn get(&self, i: usize) -> f64 {
        match self.values.get(i) {
            Some(&v) => v,
            None => f64::NAN,
        }
    }

    /// Root type at slot `i`.
    #[inline]
    pub fn root_type(&self, i: usize) -> RootType {
        match (self.types >> (3 * i)) & 7 {
            0 => RootType::Real,
            1 => RootType::Complex,
            2 => RootType::PosInf,
            3 => RootType::NegInf,
            _ => RootType::Nan,
        }
    }

    /// Overwrite the type at slot `i`.
    #[inline]
    pub fn set_type(&mut self, i: usize, t: RootType) {
        let shift = 3 * i;
        self.types = (self.types & !(7u64 << shift)) | ((t as u64) << shift);
    }
}

impl<const N: usize> std::ops::Index<usize> for Roots<N> {
    type Output = f64;
    #[inline]
    fn index(&self, i: usize) -> &f64 {
        // `Index` must hand back a reference, so an out-of-range slot borrows
        // a shared NaN rather than panicking -- matching `get` above. Using
        // the operator where a slot may be absent is still a caller bug; this
        // only changes how it is reported.
        static OUT_OF_RANGE: f64 = f64::NAN;
        match self.values.get(i) {
            Some(v) => v,
            None => &OUT_OF_RANGE,
        }
    }
}

// ---------------------------------------------------------------------------
// Constructors for each concrete arity.
// The C++ template constructors (concatenation by prepend/append) are exposed
// as named static methods on the concrete impl blocks.
// ---------------------------------------------------------------------------

impl Roots<1> {
    /// Single root with the given type and value.
    #[inline]
    pub fn new(t: RootType, x: f64) -> Self {
        Self {
            values: [x],
            types: t as u64,
        }
    }
}

impl Roots<2> {
    /// Concatenate two single roots.  C++ `Roots<2>(Roots<1>, Roots<1>)`.
    #[inline]
    pub fn from_pair(a: Roots<1>, b: Roots<1>) -> Self {
        // Destructured rather than subscripted. `Roots<1>::values` is
        // `[f64; 1]`, so rustc checks `[a0]` exhaustively; the pattern also
        // states the arity the concatenation depends on.
        let ([a0], [b0]) = (a.values, b.values);
        Self {
            values: [a0, b0],
            types: a.types | (b.types << 3),
        }
    }

    /// `Roots<1>` followed by one additional root.
    /// C++ `Roots<2>(Roots<1>, type, x)`.
    #[inline]
    pub fn with_tail(head: Roots<1>, t: RootType, x: f64) -> Self {
        let [h0] = head.values;
        Self {
            values: [h0, x],
            types: head.types | ((t as u64) << 3),
        }
    }

    /// Duplicate a single root into both slots.
    /// C++ `Roots<2>(r, r)`.
    #[inline]
    pub fn both(r: Roots<1>) -> Self {
        Self::from_pair(r, r)
    }
}

impl Roots<3> {
    /// All three slots get the same type and value.
    /// C++ `Roots<3>(type, x)`.
    #[inline]
    pub fn uniform(t: RootType, x: f64) -> Self {
        let ti = t as u64;
        Self {
            values: [x, x, x],
            types: ti | (ti << 3) | (ti << 6),
        }
    }

    /// Concatenate `Roots<1>` then `Roots<2>`.
    /// C++ `Roots<3>(Roots<1>, Roots<2>)`.
    #[inline]
    pub fn concat_1_2(a: Roots<1>, b: Roots<2>) -> Self {
        let ([a0], [b0, b1]) = (a.values, b.values);
        Self {
            values: [a0, b0, b1],
            types: a.types | (b.types << 3),
        }
    }

    /// Concatenate `Roots<2>` then `Roots<1>`.
    /// C++ `Roots<3>(Roots<2>, Roots<1>)`.
    #[inline]
    pub fn concat_2_1(a: Roots<2>, b: Roots<1>) -> Self {
        let ([a0, a1], [b0]) = (a.values, b.values);
        Self {
            values: [a0, a1, b0],
            types: a.types | (b.types << 6),
        }
    }

    /// `Roots<2>` followed by one additional root.
    /// C++ `Roots<3>(Roots<2>, type, x)`.
    #[inline]
    pub fn with_tail(head: Roots<2>, t: RootType, x: f64) -> Self {
        Self::concat_2_1(head, Roots::<1>::new(t, x))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn encoding_round_trip() {
        let r = Roots::<1>::new(RootType::Real, 3.14);
        assert_eq!(r.root_type(0), RootType::Real);
        assert_eq!(r.get(0), 3.14);
    }

    #[test]
    fn concat_2() {
        let a = Roots::<1>::new(RootType::Real, 1.0);
        let b = Roots::<1>::new(RootType::Nan, 0.0);
        let r = Roots::<2>::from_pair(a, b);
        assert_eq!(r.root_type(0), RootType::Real);
        assert_eq!(r.root_type(1), RootType::Nan);
        assert_eq!(r[0], 1.0);
        assert_eq!(r[1], 0.0);
    }

    #[test]
    fn concat_3() {
        let a = Roots::<1>::new(RootType::Real, 1.0);
        let b = Roots::<2>::from_pair(
            Roots::<1>::new(RootType::Real, 2.0),
            Roots::<1>::new(RootType::Nan, 0.0),
        );
        let r = Roots::<3>::concat_1_2(a, b);
        assert_eq!(r.root_type(0), RootType::Real);
        assert_eq!(r.root_type(1), RootType::Real);
        assert_eq!(r.root_type(2), RootType::Nan);
        assert_eq!(r[0], 1.0);
        assert_eq!(r[1], 2.0);
    }

    #[test]
    fn uniform_3() {
        let r = Roots::<3>::uniform(RootType::Complex, -5.0);
        for i in 0..3 {
            assert_eq!(r.root_type(i), RootType::Complex);
            assert_eq!(r[i], -5.0);
        }
    }
}