ternaria_arith/

balanced.rs

//! The balanced ternary numeric types: [`Tryte`], [`Word`], [`DoubleWord`].
//!
//! # Representation
//!
//! Each type stores its numeric value in a native host integer rather than a
//! packed trit string (D-03). Balanced ternary addition needs base-3 carry
//! propagation, which a host add does not perform, so the fast path is to
//! unpack, operate natively, and repack.
//!
//! BCT, two bits per trit, is the storage and interchange format, reached via
//! [`Word::to_bct`] and [`Word::from_bct`]. Guest memory holds BCT; registers
//! hold values.
//!
//! Widths and their host types (D-02):
//!
//! | Type | Trits | ~ bits | Range | Host |
//! |------|-------|--------|-------|------|
//! | [`Tryte`]      |  9 |  14.3 | +/-9,841 | `i32` |
//! | [`Word`]       | 27 |  42.8 | +/-3.81*10^12 | `i64` |
//! | [`DoubleWord`] | 54 |  85.6 | +/-2.91*10^25 | `i128` |

use crate::trit::Trit;
use core::fmt;
use core::str::FromStr;

/// 3^n as an `i128`. Exact for n <= 80.
///
/// Most calls are in const context ([`Word::MAX`] and friends) and cost
/// nothing at runtime. The exceptions are the trit shifts, which pass a runtime
/// `k` - hence the inline hint. If profiling ever shows this loop mattering,
/// the fix is a 81-entry lookup table, not a cleverer loop.
#[inline]
pub(crate) const fn pow3(n: u32) -> i128 {
    let mut r: i128 = 1;
    let mut i = 0;
    while i < n {
        r *= 3;
        i += 1;
    }
    r
}

/// Error from parsing a balanced ternary literal.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum ParseTritsError {
    /// A character outside the accepted digit set.
    ///
    /// Accepted are `T`, `t` and `-` for -1; `0` for zero; `1` and `+` for +1.
    /// Underscores are stripped before parsing and never reach here.
    BadDigit(char),
    /// No digits at all.
    Empty,
    /// More digits than the type has trits.
    TooLong {
        /// How many digits the literal actually had.
        got: usize,
        /// How many trits the target type provides.
        max: usize,
    },
}

impl fmt::Display for ParseTritsError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ParseTritsError::BadDigit(c) => write!(f, "invalid balanced ternary digit {c:?}"),
            ParseTritsError::Empty => f.write_str("empty balanced ternary literal"),
            ParseTritsError::TooLong { got, max } => {
                write!(f, "{got} digits exceeds the {max} trits available")
            }
        }
    }
}

/// Error from decoding a BCT-packed value.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct NatError {
    /// Index of the first trit position holding the `10` NaT pattern.
    pub position: u32,
}

impl fmt::Display for NatError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "NaT (not-a-trit) at trit position {}", self.position)
    }
}

macro_rules! balanced_type {
    ($name:ident, $trits:literal, $host:ty, $bct:ty, $what:literal) => {
        #[doc = concat!("A ", $what, ", ", stringify!($trits), " balanced trits.")]
        ///
        /// Stores the numeric value in the host integer; see the module docs.
        /// Every value of this type is in range by construction, so `value()`
        /// never returns something out of bounds.
        #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
        pub struct $name($host);

        impl $name {
            /// Number of trits in this type.
            pub const TRITS: u32 = $trits;

            /// 3^TRITS, the number of distinct values.
            pub const MODULUS: i128 = pow3($trits);

            /// The largest representable value, (3^n-1)/2.
            pub const MAX: Self = Self(((pow3($trits) - 1) / 2) as $host);

            /// The smallest representable value, -(3^n-1)/2.
            ///
            /// The range is symmetric, so there is no value whose negation
            /// overflows (D-01).
            pub const MIN: Self = Self(-(((pow3($trits) - 1) / 2) as $host));

            /// Zero.
            pub const ZERO: Self = Self(0);
            /// One.
            pub const ONE: Self = Self(1);
            /// Negative one.
            pub const NEG_ONE: Self = Self(-1);

            /// True when `v` is representable in this many trits.
            #[inline]
            pub const fn in_range(v: $host) -> bool {
                v >= Self::MIN.0 && v <= Self::MAX.0
            }

            /// Wraps a host value, or `None` if out of range.
            #[inline]
            pub const fn new(v: $host) -> Option<Self> {
                if Self::in_range(v) {
                    Some(Self(v))
                } else {
                    None
                }
            }

            /// Wraps a host value, panicking if out of range.
            ///
            /// For constants and tests where the range is known statically.
            #[inline]
            #[track_caller]
            pub const fn from_value(v: $host) -> Self {
                Self::new(v).expect(concat!(stringify!($name), ": value out of range"))
            }

            /// Reduces an `i128` into range modulo 3^n. Always succeeds.
            #[inline]
            pub const fn from_i128_wrapping(v: i128) -> Self {
                let m = Self::MODULUS;
                let mut r = v % m;
                let max = (m - 1) / 2;
                if r > max {
                    r -= m;
                } else if r < -max {
                    r += m;
                }
                Self(r as $host)
            }

            /// Reduces an `i128` into range, or `None` if it does not fit.
            #[inline]
            pub const fn from_i128(v: i128) -> Option<Self> {
                let max = (Self::MODULUS - 1) / 2;
                if v >= -max && v <= max {
                    Some(Self(v as $host))
                } else {
                    None
                }
            }

            /// The numeric value.
            #[inline]
            pub const fn value(self) -> $host {
                self.0
            }

            /// The numeric value widened to `i128`.
            #[inline]
            pub const fn to_i128(self) -> i128 {
                self.0 as i128
            }

            /// The trits, least significant first.
            pub const fn trits(self) -> [Trit; $trits] {
                let mut out = [Trit::Zero; $trits];
                let mut n = self.0;
                let mut i = 0;
                while i < $trits {
                    // Rust's `%` truncates toward zero, so for n < 0 the
                    // remainder lands in {-2,-1,0}. Both 2 and -2 are outside
                    // the digit set and borrow from the next position.
                    let mut r = n % 3;
                    n /= 3;
                    if r == 2 {
                        r = -1;
                        n += 1;
                    } else if r == -2 {
                        r = 1;
                        n -= 1;
                    }
                    out[i] = match r {
                        -1 => Trit::Neg,
                        1 => Trit::Pos,
                        _ => Trit::Zero,
                    };
                    i += 1;
                }
                out
            }

            /// Rebuilds a value from trits, least significant first.
            pub const fn from_trits(trits: [Trit; $trits]) -> Self {
                let mut acc: i128 = 0;
                let mut p: i128 = 1;
                let mut i = 0;
                while i < $trits {
                    acc += (trits[i] as i8 as i128) * p;
                    p *= 3;
                    i += 1;
                }
                Self(acc as $host)
            }

            /// The trit at `index`, counting from the least significant.
            ///
            /// # Panics
            /// If `index >= TRITS`.
            #[inline]
            #[track_caller]
            pub const fn trit(self, index: u32) -> Trit {
                assert!(index < $trits, "trit index out of range");
                self.trits()[index as usize]
            }

            // ---- BCT storage format (D-03) ----

            /// Packs into BCT: two bits per trit, least significant trit in the
            /// low bits.
            pub const fn to_bct(self) -> $bct {
                let trits = self.trits();
                let mut out: $bct = 0;
                let mut i = 0;
                while i < $trits {
                    out |= (trits[i].bct_code() as $bct) << (2 * i);
                    i += 1;
                }
                out
            }

            /// Unpacks from BCT, rejecting the `10` NaT pattern.
            ///
            /// NaT marks poisoned host memory and must never reach a guest, so
            /// decoding it is an error rather than a silent zero (D-03).
            pub const fn from_bct(bits: $bct) -> Result<Self, NatError> {
                let mut trits = [Trit::Zero; $trits];
                let mut i = 0;
                while i < $trits {
                    let code = ((bits >> (2 * i)) & 0b11) as u8;
                    match Trit::from_bct_code(code) {
                        Some(t) => trits[i] = t,
                        None => return Err(NatError { position: i as u32 }),
                    }
                    i += 1;
                }
                Ok(Self::from_trits(trits))
            }

            // ---- arithmetic ----
            //
            // Checked forms return None on overflow. Intermediates go through
            // i128. If an i128 operation overflows, the true result is outside
            // this type's range, so the i128 check is a sound overflow test.

            /// Addition, `None` on overflow.
            #[inline]
            pub const fn checked_add(self, rhs: Self) -> Option<Self> {
                match (self.0 as i128).checked_add(rhs.0 as i128) {
                    Some(v) => Self::from_i128(v),
                    None => None,
                }
            }

            /// Subtraction, `None` on overflow.
            #[inline]
            pub const fn checked_sub(self, rhs: Self) -> Option<Self> {
                match (self.0 as i128).checked_sub(rhs.0 as i128) {
                    Some(v) => Self::from_i128(v),
                    None => None,
                }
            }

            /// Multiplication, `None` on overflow.
            #[inline]
            pub const fn checked_mul(self, rhs: Self) -> Option<Self> {
                match (self.0 as i128).checked_mul(rhs.0 as i128) {
                    Some(v) => Self::from_i128(v),
                    None => None,
                }
            }

            /// Truncating division, `None` only when `rhs` is zero.
            ///
            /// Division cannot overflow: the range is symmetric, so there is no
            /// `MIN / -1` case (D-01).
            #[inline]
            pub const fn checked_div(self, rhs: Self) -> Option<Self> {
                if rhs.0 == 0 {
                    return None;
                }
                Self::from_i128((self.0 as i128) / (rhs.0 as i128))
            }

            /// Remainder, `None` only when `rhs` is zero.
            #[inline]
            pub const fn checked_rem(self, rhs: Self) -> Option<Self> {
                if rhs.0 == 0 {
                    return None;
                }
                Self::from_i128((self.0 as i128) % (rhs.0 as i128))
            }

            /// Negation. Cannot overflow; the range is symmetric (D-01).
            #[inline]
            pub const fn neg(self) -> Self {
                Self(-self.0)
            }

            /// Absolute value. Total; the range is symmetric (D-01).
            #[inline]
            pub const fn abs(self) -> Self {
                Self(if self.0 < 0 { -self.0 } else { self.0 })
            }

            /// Addition modulo 3^n.
            #[inline]
            pub const fn wrapping_add(self, rhs: Self) -> Self {
                Self::from_i128_wrapping(self.0 as i128 + rhs.0 as i128)
            }

            /// Subtraction modulo 3^n.
            #[inline]
            pub const fn wrapping_sub(self, rhs: Self) -> Self {
                Self::from_i128_wrapping(self.0 as i128 - rhs.0 as i128)
            }

            /// Shift left by `k` trits, multiplying by 3^k. `None` on
            /// overflow. Exact by construction.
            #[inline]
            pub const fn checked_shl_trits(self, k: u32) -> Option<Self> {
                if k >= 128 {
                    return if self.0 == 0 { Some(Self::ZERO) } else { None };
                }
                match (self.0 as i128).checked_mul(pow3(k)) {
                    Some(v) => Self::from_i128(v),
                    None => None,
                }
            }

            /// Shift right by `k` trits, dropping the low `k` digits. This is
            /// division by 3^k rounded to nearest: the dropped digits are worth
            /// at most (3^k-1)/2, half a unit in the last place.
            ///
            /// This is not the host's `/`, which truncates toward zero. For
            /// 1,598,046,941,971 / 27 truncation leaves a remainder of 19
            /// against a ULP of 27. Implementing this as `self.0 / pow3(k)` was
            /// a bug, caught by `truncation_is_round_to_nearest`.
            ///
            /// Ties cannot occur: 3^k is odd, so twice the remainder never
            /// equals it. No round-half-to-even rule is needed.
            #[inline]
            pub const fn shr_trits(self, k: u32) -> Self {
                if k >= 128 {
                    return Self::ZERO;
                }
                let d = pow3(k);
                let v = self.0 as i128;
                let q = v / d;
                let r = v - q * d;
                // Nudge toward the nearer multiple when the truncated-toward-zero
                // remainder exceeds half a ULP. Symmetric in sign.
                let adj = if 2 * r > d {
                    1
                } else if 2 * r < -d {
                    -1
                } else {
                    0
                };
                Self((q + adj) as $host)
            }

            /// Three-way comparison returning one trit (D-06).
            ///
            /// `Neg` for less, `Zero` for equal, `Pos` for greater.
            #[inline]
            pub const fn cmp3(self, rhs: Self) -> Trit {
                if self.0 < rhs.0 {
                    Trit::Neg
                } else if self.0 > rhs.0 {
                    Trit::Pos
                } else {
                    Trit::Zero
                }
            }

            /// The sign, as a trit.
            #[inline]
            pub const fn signum(self) -> Trit {
                self.cmp3(Self::ZERO)
            }

            /// Applies a per-trit unary operation.
            pub fn map_trits(self, f: impl Fn(Trit) -> Trit) -> Self {
                let mut t = self.trits();
                for slot in t.iter_mut() {
                    *slot = f(*slot);
                }
                Self::from_trits(t)
            }

            /// Applies a per-trit binary operation, used by `MIN` and `MAX`.
            pub fn zip_trits(self, rhs: Self, f: impl Fn(Trit, Trit) -> Trit) -> Self {
                let a = self.trits();
                let b = rhs.trits();
                let mut out = [Trit::Zero; $trits];
                for i in 0..$trits {
                    out[i] = f(a[i], b[i]);
                }
                Self::from_trits(out)
            }

            /// Per-trit logical AND (minimum).
            pub fn trit_and(self, rhs: Self) -> Self {
                self.zip_trits(rhs, Trit::and)
            }

            /// Per-trit logical OR (maximum).
            pub fn trit_or(self, rhs: Self) -> Self {
                self.zip_trits(rhs, Trit::or)
            }

            /// Per-trit logical NOT. Identical to arithmetic negation, since
            /// negating a balanced ternary number flips every digit (D-01).
            pub fn trit_not(self) -> Self {
                self.neg()
            }
        }

        impl fmt::Display for $name {
            /// Balanced ternary digits, most significant first, `T` for -1.
            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                let trits = self.trits();
                let mut hi = None;
                for i in (0..$trits).rev() {
                    if trits[i] != Trit::Zero {
                        hi = Some(i);
                        break;
                    }
                }
                match hi {
                    None => f.write_str("0"),
                    Some(h) => {
                        for i in (0..=h).rev() {
                            fmt::Display::fmt(&trits[i], f)?;
                        }
                        Ok(())
                    }
                }
            }
        }

        impl fmt::Debug for $name {
            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                write!(f, "{}({} = {})", stringify!($name), self, self.0)
            }
        }

        impl FromStr for $name {
            type Err = ParseTritsError;

            /// Parses balanced ternary digits, most significant first.
            /// Underscores are ignored as separators.
            fn from_str(s: &str) -> Result<Self, Self::Err> {
                let digits: Vec<char> = s.chars().filter(|c| *c != '_').collect();
                if digits.is_empty() {
                    return Err(ParseTritsError::Empty);
                }
                if digits.len() > $trits {
                    return Err(ParseTritsError::TooLong {
                        got: digits.len(),
                        max: $trits,
                    });
                }
                let mut trits = [Trit::Zero; $trits];
                // Input is most-significant-first; storage is the reverse.
                for (i, c) in digits.iter().rev().enumerate() {
                    trits[i] = Trit::from_char(*c).ok_or(ParseTritsError::BadDigit(*c))?;
                }
                Ok(Self::from_trits(trits))
            }
        }

        impl core::ops::Neg for $name {
            type Output = Self;
            #[inline]
            fn neg(self) -> Self {
                Self::neg(self)
            }
        }
    };
}

balanced_type!(Tryte, 9, i32, u32, "tryte, the addressable unit");
balanced_type!(Word, 27, i64, u64, "word");
balanced_type!(DoubleWord, 54, i128, u128, "double word");

/// Splits a value into trytes and rebuilds it from them.
///
/// A tryte is nine trits, so the tryte at index `k` carries weight 3^(9k) and
/// the decomposition is exact: `$trytes` trytes tile the range with nothing
/// left over. Memory and memory-mapped devices both store multi-tryte values
/// this way, least significant tryte first.
macro_rules! tryte_split {
    ($name:ident, $trytes:expr, $host:ty) => {
        impl $name {
            /// Number of trytes in the representation.
            pub const TRYTES: usize = $trytes;

            /// Tryte `index`, counting from the least significant.
            ///
            /// # Panics
            /// If `index >= TRYTES`.
            #[track_caller]
            pub const fn tryte(self, index: usize) -> Tryte {
                assert!(index < $trytes, "tryte index out of range");
                let trits = self.trits();
                let mut chunk = [Trit::Zero; 9];
                let mut i = 0;
                while i < 9 {
                    chunk[i] = trits[index * 9 + i];
                    i += 1;
                }
                Tryte::from_trits(chunk)
            }

            /// Rebuilds a value from its trytes, least significant first.
            pub const fn from_trytes(trytes: [Tryte; $trytes]) -> Self {
                let mut acc: i128 = 0;
                let mut weight: i128 = 1;
                let mut i = 0;
                while i < $trytes {
                    acc += (trytes[i].value() as i128) * weight;
                    weight *= 19683;
                    i += 1;
                }
                Self(acc as $host)
            }
        }
    };
}

tryte_split!(Word, 3, i64);
tryte_split!(DoubleWord, 6, i128);

/// Adds `wrapping_mul` where the exact product fits an `i128` intermediate.
///
/// Not provided for [`DoubleWord`]. Two 54-trit values multiply to at most
/// 3^108, about 4*10^51, which exceeds `i128`. `i128::wrapping_mul` would
/// reduce modulo 2^128 rather than 3^54 and give a wrong answer.
/// [`DoubleWord::checked_mul`] is exact and unaffected.
macro_rules! wrapping_mul_via_i128 {
    ($name:ident) => {
        impl $name {
            /// Multiplication modulo 3^n.
            #[inline]
            pub const fn wrapping_mul(self, rhs: Self) -> Self {
                Self::from_i128_wrapping(self.value() as i128 * rhs.value() as i128)
            }
        }
    };
}

wrapping_mul_via_i128!(Tryte);
wrapping_mul_via_i128!(Word);

impl Word {
    /// Widens to a [`DoubleWord`]. Always exact.
    #[inline]
    pub const fn widen(self) -> DoubleWord {
        DoubleWord(self.value() as i128)
    }

    /// The exact product of two words, as a double word.
    ///
    /// Cannot overflow: the product of two p-trit values is exactly 2p trits,
    /// and 27 + 27 is the double-word width. A `DoubleWord` accumulator
    /// therefore holds each product exactly and only the summation rounds
    /// (D-02).
    #[inline]
    pub const fn widening_mul(self, rhs: Self) -> DoubleWord {
        DoubleWord(self.value() as i128 * rhs.value() as i128)
    }
}

impl Tryte {
    /// Widens to a [`Word`]. Always exact.
    #[inline]
    pub const fn widen(self) -> Word {
        Word(self.value() as i64)
    }
}

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

    #[test]
    fn word_trytes_carry_powers_of_three_to_the_ninth() {
        // Tryte k has weight 3^(9k), so a value of exactly 3^(9k) puts a one
        // in tryte k and zero everywhere else.
        for k in 0..Word::TRYTES {
            let v = Word::from_value(19683i64.pow(k as u32));
            for j in 0..Word::TRYTES {
                let want = if j == k { 1 } else { 0 };
                assert_eq!(v.tryte(j).value(), want, "3^(9*{k}) tryte {j}");
            }
        }
    }

    #[test]
    fn word_round_trips_through_its_trytes() {
        let cases = [
            Word::MIN,
            Word::from_value(-1),
            Word::ZERO,
            Word::from_value(1),
            Word::from_value(-4_294_967_296),
            Word::from_value(1_598_046_941_971),
            Word::MAX,
        ];
        for v in cases {
            let trytes = [v.tryte(0), v.tryte(1), v.tryte(2)];
            assert_eq!(Word::from_trytes(trytes), v, "value {}", v.value());
        }
    }

    #[test]
    fn double_word_round_trips_through_its_trytes() {
        let cases = [
            DoubleWord::MIN,
            DoubleWord::ZERO,
            DoubleWord::from_value(-987_654_321_098_765_432),
            DoubleWord::MAX,
        ];
        for v in cases {
            let mut trytes = [Tryte::ZERO; DoubleWord::TRYTES];
            for (i, slot) in trytes.iter_mut().enumerate() {
                *slot = v.tryte(i);
            }
            assert_eq!(DoubleWord::from_trytes(trytes), v, "value {}", v.value());
        }
    }

    #[test]
    fn every_tryte_value_round_trips_in_the_low_position() {
        for n in Tryte::MIN.value()..=Tryte::MAX.value() {
            let t = Tryte::from_value(n);
            let w = Word::from_trytes([t, Tryte::ZERO, Tryte::ZERO]);
            assert_eq!(w.value(), n as i64);
            assert_eq!(w.tryte(0), t);
        }
    }
}