tenferro_tensor_core/

scalar_set.rs

//! Closed scalar sets.
//!
//! A scalar set is the closed list of scalar types one tensor value type can
//! carry. tenferro declares its own set once and gets the value enum, the tag,
//! and the membership query from that single declaration. A crate that needs a
//! different set declares it with [`define_scalar_set!`] in its own crate, and
//! its values implement this trait without touching tenferro's set.

/// Kind of arithmetic a set member belongs to.
///
/// # Examples
///
/// ```rust
/// use tenferro_tensor_core::MemberKind;
///
/// assert_ne!(MemberKind::Integer, MemberKind::Float);
/// ```
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum MemberKind {
    /// Boolean.
    Boolean,
    /// Signed integer.
    Integer,
    /// Real floating point.
    Float,
    /// Complex floating point.
    Complex,
    /// A scalar tenferro does not declare, so its facts are unknown.
    External,
}

/// Promotion-relevant facts about one set member.
///
/// `level` orders members within a kind and `width` is the component width in
/// bits. `level` is what orders two members of the same kind, so a set can rank
/// an extended-precision member above a standard one without claiming a wider
/// exponent range.
///
/// # Examples
///
/// ```rust
/// use tenferro_tensor_core::{MemberKind, MemberSpec};
///
/// let widened = tenferro_tensor_core::promote_specs(
///     MemberSpec::new(MemberKind::Float, 0, 32),
///     MemberSpec::new(MemberKind::Float, 1, 64),
/// );
/// assert_eq!(widened, MemberSpec::new(MemberKind::Float, 1, 64));
/// ```
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct MemberSpec {
    /// Arithmetic kind of the member.
    pub kind: MemberKind,
    /// Rank within the kind.
    pub level: u32,
    /// Component width in bits.
    pub width: u32,
}

impl MemberSpec {
    /// Build a member fact.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use tenferro_tensor_core::{MemberKind, MemberSpec};
    ///
    /// let spec = MemberSpec::new(MemberKind::Float, 1, 64);
    /// assert_eq!(spec.level, 1);
    /// ```
    #[must_use]
    pub const fn new(kind: MemberKind, level: u32, width: u32) -> Self {
        Self { kind, level, width }
    }
}

/// Combine two member facts under the ordinary numeric promotion rules.
///
/// A boolean yields to anything. Two members of one kind keep the higher level.
/// An integer with a float or complex yields to that kind's widest member,
/// because an integer is widened rather than mixed. A float with a complex takes
/// the narrowest complex that can still hold both, which keeps `f32 + c32` in
/// `c32` while `f64 + c32` becomes `c64`.
///
/// # Examples
///
/// ```rust
/// use tenferro_tensor_core::{promote_specs, MemberKind, MemberSpec};
///
/// let f32_ = MemberSpec::new(MemberKind::Float, 0, 32);
/// let c32 = MemberSpec::new(MemberKind::Complex, 0, 32);
/// assert_eq!(promote_specs(f32_, c32), c32);
/// ```
#[must_use]
pub const fn promote_specs(lhs: MemberSpec, rhs: MemberSpec) -> MemberSpec {
    use MemberKind::{Boolean, Complex, External, Float, Integer};
    match (lhs.kind, rhs.kind) {
        (Boolean, _) => rhs,
        (_, Boolean) => lhs,
        (External, _) => lhs,
        (_, External) => rhs,
        (Integer, Float) | (Float, Integer) => MemberSpec::new(Float, u32::MAX, u32::MAX),
        (Integer, Complex) | (Complex, Integer) => MemberSpec::new(Complex, u32::MAX, u32::MAX),
        (Integer, Integer) | (Float, Float) | (Complex, Complex) => {
            if lhs.level >= rhs.level {
                lhs
            } else {
                rhs
            }
        }
        (Float, Complex) | (Complex, Float) => {
            let width = if lhs.width >= rhs.width {
                lhs.width
            } else {
                rhs.width
            };
            MemberSpec::new(Complex, 0, width)
        }
    }
}

/// Promote two members of one set from its parallel tag and fact tables.
///
/// `tags` and `specs` are emitted together in the same order by
/// [`define_scalar_tag!`], so a tag that is not listed in `tags` (an externally
/// defined member) promotes to itself.
///
/// # Examples
///
/// ```rust
/// use tenferro_tensor_core::{promote_in_set, DType};
///
/// assert_eq!(
///     promote_in_set(DType::TAGS, DType::SPECS, DType::I32, DType::F32),
///     DType::F64
/// );
/// ```
#[must_use]
pub fn promote_in_set<Tag: Copy + PartialEq>(
    tags: &[Tag],
    specs: &[MemberSpec],
    lhs: Tag,
    rhs: Tag,
) -> Tag {
    // INVARIANT: `define_scalar_tag!` emits `TAGS` and `SPECS` from one member
    // list, so the two tables always hold the same entries in the same order.
    let (Some(lhs_index), Some(rhs_index)) = (
        tags.iter().position(|tag| *tag == lhs),
        tags.iter().position(|tag| *tag == rhs),
    ) else {
        return lhs;
    };
    let (Some(lhs_spec), Some(rhs_spec)) = (specs.get(lhs_index), specs.get(rhs_index)) else {
        return lhs;
    };
    let target = promote_specs(*lhs_spec, *rhs_spec);
    for (index, spec) in specs.iter().enumerate() {
        if *spec == target {
            if let Some(tag) = tags.get(index) {
                return *tag;
            }
        }
    }
    let mut chosen: Option<(Tag, MemberSpec)> = None;
    for (index, spec) in specs.iter().enumerate() {
        let Some(tag) = tags.get(index).copied() else {
            break;
        };
        if spec.kind != target.kind {
            continue;
        }
        let better = match chosen {
            None => true,
            Some((_, current)) => {
                let current_wide_enough = current.width >= target.width;
                let candidate_wide_enough = spec.width >= target.width;
                match (current_wide_enough, candidate_wide_enough) {
                    (false, true) => true,
                    (true, false) => false,
                    (true, true) => {
                        spec.width < current.width
                            || (spec.width == current.width && spec.level < current.level)
                    }
                    (false, false) => {
                        spec.width > current.width
                            || (spec.width == current.width && spec.level < current.level)
                    }
                }
            }
        };
        if better {
            chosen = Some((tag, *spec));
        }
    }
    match chosen {
        Some((tag, _)) => tag,
        None => lhs,
    }
}

/// Define a closed scalar set's tag type and its promotion facts.
///
/// The declaration lists each member once. The macro emits the tag enum and the
/// member facts (`spec`, `TAGS`, `SPECS`), plus the `external` variant when the
/// declaration names one. Use this form when the value enum lives in another
/// crate: tenferro's own preset set declares its tag here and its value enum
/// beside the tensor family that carries it.
///
/// # Examples
///
/// ```rust
/// use tenferro_tensor_core::{define_scalar_tag, MemberKind};
///
/// define_scalar_tag! {
///     /// Tag for a two-member set.
///     pub enum PairTag {
///         /// Double precision.
///         F64 => f64 : Float 1 64,
///         /// Single precision.
///         F32 => f32 : Float 0 32,
///     }
/// }
///
/// assert_eq!(PairTag::F32.spec().kind, MemberKind::Float);
/// assert_eq!(PairTag::TAGS, &[PairTag::F64, PairTag::F32]);
/// ```
#[macro_export]
macro_rules! define_scalar_tag {
    (
        $(#[$tag_meta:meta])*
        $tag_vis:vis enum $tag:ident {
            $(
                $(#[$variant_meta:meta])*
                $variant:ident => $ty:ty : $kind:ident $level:literal $width:literal
            ),+ $(,)?
        }
        $( external $ext_variant:ident($ext_ty:ty); )?
    ) => {
        $(#[$tag_meta])*
        #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
        $tag_vis enum $tag {
            $(
                $(#[$variant_meta])*
                $variant,
            )+
            $( $ext_variant($ext_ty), )?
        }

        impl $tag {
            /// Promotion facts of this member.
            ///
            /// # Examples
            ///
            /// ```rust
            /// use tenferro_tensor_core::{MemberKind, DType};
            ///
            /// assert_eq!(DType::F64.spec().kind, MemberKind::Float);
            /// ```
            #[must_use]
            pub const fn spec(self) -> $crate::MemberSpec {
                match self {
                    $(
                        $tag::$variant => $crate::MemberSpec::new(
                            $crate::MemberKind::$kind,
                            $level,
                            $width,
                        ),
                    )+
                    $( $tag::$ext_variant(_) => {
                        $crate::MemberSpec::new($crate::MemberKind::External, 0, 0)
                    } )?
                }
            }

            /// Every member tag, in declaration order.
            pub const TAGS: &'static [Self] = &[
                $(
                    $tag::$variant,
                )+
            ];

            /// Promotion facts of every member, in declaration order.
            pub const SPECS: &'static [$crate::MemberSpec] = &[
                $(
                    $crate::MemberSpec::new($crate::MemberKind::$kind, $level, $width),
                )+
            ];
        }
    };
}