xad-rs 0.2.0

//! `LabeledDual2<T>` — labeled wrapper over seeded `Dual2<T>`.
//!
//! **Shape A (Phase 02.2):** does NOT carry an `Arc<VarRegistry>` field in
//! release builds. The struct layout is `{ inner: Dual2<T>, seeded: Option<usize> }`
//! plus, under `#[cfg(debug_assertions)]` only, a `gen_id: u64` stamped by the
//! owning [`LabeledForwardTape`] scope for the cross-registry debug guard.
//! Release builds carry zero atomic-refcount cost per operator.
//!
//! `Dual2<T>` tracks first AND second derivatives along ONE seed direction.
//! The labeled form gives that direction a name at construction time and
//! exposes [`first_derivative`](LabeledDual2::first_derivative) /
//! [`second_derivative`](LabeledDual2::second_derivative) accessors that
//! return the seeded values when the name matches and `T::zero()`
//! otherwise. The `seeded: Option<usize>` field is NOT the registry — it
//! is the positional index of the currently-seeded variable, a per-value
//! attribute that survives Shape A.
//!
//! Unlike `LabeledDual` and `LabeledFReal`, every binary op on
//! `LabeledDual2` must propagate the `seeded: Option<usize>` field through
//! a private [`merge_seeded`] combinator. Operations between two
//! `LabeledDual2<T>` values with different seeds panic in debug builds
//! (two-direction `Dual2` is a semantic violation).
//!
//! The only way to obtain a `LabeledDual2<T>` is via the
//! `LabeledForwardTape` constructor API (see module docs in
//! `src/labeled/forward_tape.rs`).

use std::fmt;

use crate::dual2::Dual2;
use crate::traits::Scalar;

/// Labeled wrapper around the positional [`Dual2<T>`] type.
///
/// # Example
///
/// The constructor API for `LabeledDual2<f64>` is owned by
/// `LabeledForwardTape` via the `declare_dual2_f64` / `freeze_dual` /
/// `scope.dual2(handle)` pattern. See the module docs in
/// `src/labeled/forward_tape.rs` for the full rationale.
///
/// ```
/// use xad_rs::labeled::{LabeledForwardScope, LabeledForwardTape};
///
/// let mut ft = LabeledForwardTape::new();
/// let x_h = ft.declare_dual2_f64("x", 3.0);
/// let scope: LabeledForwardScope = ft.freeze_dual();
///
/// let x = scope.dual2(x_h);
/// let f = x * x; // f(x) = x^2, f'(x) = 2x = 6, f''(x) = 2
///
/// assert_eq!(f.value(), 9.0);
/// assert_eq!(f.first_derivative("x"), 6.0);
/// assert_eq!(f.second_derivative("x"), 2.0);
/// ```
#[derive(Clone)]
pub struct LabeledDual2<T: Scalar> {
    pub(super) inner: Dual2<T>,
    pub(super) seeded: Option<usize>,
    // NOTE: field name is `gen_id` — `gen` alone is a reserved keyword in
    // Rust 2024 edition.
    #[cfg(debug_assertions)]
    pub(super) gen_id: u64,
}

impl<T: Scalar> LabeledDual2<T> {
    /// Internal constructor used by `LabeledForwardTape` input / freeze
    /// paths. Reads the TLS active generation (debug builds only) to stamp
    /// the `gen_id` field. Not part of the public API.
    #[inline]
    pub(crate) fn __from_parts(inner: Dual2<T>, seeded: Option<usize>) -> Self {
        Self {
            inner,
            seeded,
            #[cfg(debug_assertions)]
            gen_id: crate::labeled::forward_tape::current_gen(),
        }
    }

    /// Value part.
    #[inline]
    pub fn value(&self) -> T {
        self.inner.value()
    }

    /// First derivative with respect to `name`.
    ///
    /// Returns the seeded first derivative if `name` matches the seeded
    /// variable, else `T::zero()`. Reads the active registry from the
    /// `LabeledForwardTape` thread-local slot. Panics if `name` is not in
    /// the registry, or if called outside a frozen `LabeledForwardTape`
    /// scope.
    pub fn first_derivative(&self, name: &str) -> T {
        let idx = crate::labeled::forward_tape::with_active_registry(|r| {
            let r = r.expect(
                "LabeledDual2::first_derivative called outside a frozen LabeledForwardTape scope",
            );
            r.index_of(name).unwrap_or_else(|| {
                panic!(
                    "LabeledDual2::first_derivative: name {:?} not present in registry",
                    name
                )
            })
        });
        if self.seeded == Some(idx) {
            self.inner.first_derivative()
        } else {
            T::zero()
        }
    }

    /// Second derivative (diagonal Hessian entry) with respect to `name`.
    ///
    /// Returns the seeded second derivative if `name` matches the seeded
    /// variable, else `T::zero()`. Panics if `name` is not in the registry,
    /// or if called outside a frozen `LabeledForwardTape` scope.
    pub fn second_derivative(&self, name: &str) -> T {
        let idx = crate::labeled::forward_tape::with_active_registry(|r| {
            let r = r.expect(
                "LabeledDual2::second_derivative called outside a frozen LabeledForwardTape scope",
            );
            r.index_of(name).unwrap_or_else(|| {
                panic!(
                    "LabeledDual2::second_derivative: name {:?} not present in registry",
                    name
                )
            })
        });
        if self.seeded == Some(idx) {
            self.inner.second_derivative()
        } else {
            T::zero()
        }
    }

    /// Escape hatch: direct access to the inner positional `Dual2<T>`.
    #[inline]
    pub fn inner(&self) -> &Dual2<T> {
        &self.inner
    }

    // ============ Elementary math delegations ============
    // Each method forwards to the inherent `Dual2<T>` elementary (which
    // takes `self` by value — `Dual2<T>` is `Copy`), preserves the
    // `seeded` field because unary elementaries cannot change which
    // direction is active, and stamps the parent's generation explicitly.

    /// Natural exponential, preserving the seed and the parent's generation.
    #[inline]
    pub fn exp(&self) -> Self {
        Self {
            inner: self.inner.exp(),
            seeded: self.seeded,
            #[cfg(debug_assertions)]
            gen_id: self.gen_id,
        }
    }

    /// Natural logarithm, preserving the seed and the parent's generation.
    #[inline]
    pub fn ln(&self) -> Self {
        Self {
            inner: self.inner.ln(),
            seeded: self.seeded,
            #[cfg(debug_assertions)]
            gen_id: self.gen_id,
        }
    }

    /// Square root, preserving the seed and the parent's generation.
    #[inline]
    pub fn sqrt(&self) -> Self {
        Self {
            inner: self.inner.sqrt(),
            seeded: self.seeded,
            #[cfg(debug_assertions)]
            gen_id: self.gen_id,
        }
    }

    /// Sine, preserving the seed and the parent's generation.
    #[inline]
    pub fn sin(&self) -> Self {
        Self {
            inner: self.inner.sin(),
            seeded: self.seeded,
            #[cfg(debug_assertions)]
            gen_id: self.gen_id,
        }
    }

    /// Cosine, preserving the seed and the parent's generation.
    #[inline]
    pub fn cos(&self) -> Self {
        Self {
            inner: self.inner.cos(),
            seeded: self.seeded,
            #[cfg(debug_assertions)]
            gen_id: self.gen_id,
        }
    }
}

impl<T: Scalar> fmt::Debug for LabeledDual2<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("LabeledDual2")
            .field("value", &self.inner.value())
            .field("first", &self.inner.first_derivative())
            .field("second", &self.inner.second_derivative())
            .field("seeded", &self.seeded)
            .finish()
    }
}

impl<T: Scalar> fmt::Display for LabeledDual2<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "LabeledDual2({})", self.inner.value())
    }
}

/// Merge two seeded directions. Same seed or one constant is fine;
/// two different seeds is a `Dual2` semantic violation — debug-panic,
/// release-pick-LHS.
#[inline]
pub(super) fn merge_seeded(a: Option<usize>, b: Option<usize>) -> Option<usize> {
    match (a, b) {
        (None, None) => None,
        (Some(x), None) | (None, Some(x)) => Some(x),
        (Some(x), Some(y)) if x == y => Some(x),
        (Some(_), Some(_)) => {
            #[cfg(debug_assertions)]
            panic!(
                "LabeledDual2: operation between two differently-seeded variables; \
                 seeded Dual2 supports only one active direction"
            );
            #[cfg(not(debug_assertions))]
            a
        }
    }
}

// ============================ Operator impls ============================
// Shape A hand-written local macro (no shared stamping macro — the
// LBLF-07 scaffold was deleted in Plan 02.2-02 Task 4) because of the
// `seeded` field merge logic.
// `Dual2<T>` is `Copy`, so inner-value ops are by-value throughout.
// Each wrapper-vs-wrapper impl performs a debug-only `check_gen` between
// the two operands' generations, then merges `seeded` via `merge_seeded`,
// then constructs the result preserving the LHS's generation stamp.
// Pattern: 6 variants per binary op × 4 ops = 24 impls, plus 2 Neg impls.

macro_rules! __lbld2_binop {
    ($trait_:ident, $method:ident, $op:tt) => {
        // owned + owned
        impl<T: Scalar> ::core::ops::$trait_<LabeledDual2<T>> for LabeledDual2<T> {
            type Output = LabeledDual2<T>;
            #[inline]
            fn $method(self, rhs: LabeledDual2<T>) -> LabeledDual2<T> {
                #[cfg(debug_assertions)]
                crate::labeled::forward_tape::check_gen(self.gen_id, rhs.gen_id);
                LabeledDual2 {
                    inner: self.inner $op rhs.inner,
                    seeded: merge_seeded(self.seeded, rhs.seeded),
                    #[cfg(debug_assertions)]
                    gen_id: self.gen_id,
                }
            }
        }
        // ref + ref
        impl<T: Scalar> ::core::ops::$trait_<&LabeledDual2<T>> for &LabeledDual2<T> {
            type Output = LabeledDual2<T>;
            #[inline]
            fn $method(self, rhs: &LabeledDual2<T>) -> LabeledDual2<T> {
                #[cfg(debug_assertions)]
                crate::labeled::forward_tape::check_gen(self.gen_id, rhs.gen_id);
                LabeledDual2 {
                    inner: self.inner $op rhs.inner,
                    seeded: merge_seeded(self.seeded, rhs.seeded),
                    #[cfg(debug_assertions)]
                    gen_id: self.gen_id,
                }
            }
        }
        // owned + ref
        impl<T: Scalar> ::core::ops::$trait_<&LabeledDual2<T>> for LabeledDual2<T> {
            type Output = LabeledDual2<T>;
            #[inline]
            fn $method(self, rhs: &LabeledDual2<T>) -> LabeledDual2<T> {
                #[cfg(debug_assertions)]
                crate::labeled::forward_tape::check_gen(self.gen_id, rhs.gen_id);
                LabeledDual2 {
                    inner: self.inner $op rhs.inner,
                    seeded: merge_seeded(self.seeded, rhs.seeded),
                    #[cfg(debug_assertions)]
                    gen_id: self.gen_id,
                }
            }
        }
        // ref + owned
        impl<T: Scalar> ::core::ops::$trait_<LabeledDual2<T>> for &LabeledDual2<T> {
            type Output = LabeledDual2<T>;
            #[inline]
            fn $method(self, rhs: LabeledDual2<T>) -> LabeledDual2<T> {
                #[cfg(debug_assertions)]
                crate::labeled::forward_tape::check_gen(self.gen_id, rhs.gen_id);
                LabeledDual2 {
                    inner: self.inner $op rhs.inner,
                    seeded: merge_seeded(self.seeded, rhs.seeded),
                    #[cfg(debug_assertions)]
                    gen_id: self.gen_id,
                }
            }
        }
        // X op T (scalar, no cross-registry check)
        impl<T: Scalar> ::core::ops::$trait_<T> for LabeledDual2<T> {
            type Output = LabeledDual2<T>;
            #[inline]
            fn $method(self, rhs: T) -> LabeledDual2<T> {
                LabeledDual2 {
                    inner: self.inner $op rhs,
                    seeded: self.seeded,
                    #[cfg(debug_assertions)]
                    gen_id: self.gen_id,
                }
            }
        }
        impl<T: Scalar> ::core::ops::$trait_<T> for &LabeledDual2<T> {
            type Output = LabeledDual2<T>;
            #[inline]
            fn $method(self, rhs: T) -> LabeledDual2<T> {
                LabeledDual2 {
                    inner: self.inner $op rhs,
                    seeded: self.seeded,
                    #[cfg(debug_assertions)]
                    gen_id: self.gen_id,
                }
            }
        }
    };
}

__lbld2_binop!(Add, add, +);
__lbld2_binop!(Sub, sub, -);
__lbld2_binop!(Mul, mul, *);
__lbld2_binop!(Div, div, /);

impl<T: Scalar> ::core::ops::Neg for LabeledDual2<T> {
    type Output = LabeledDual2<T>;
    #[inline]
    fn neg(self) -> LabeledDual2<T> {
        LabeledDual2 {
            inner: -self.inner,
            seeded: self.seeded,
            #[cfg(debug_assertions)]
            gen_id: self.gen_id,
        }
    }
}
impl<T: Scalar> ::core::ops::Neg for &LabeledDual2<T> {
    type Output = LabeledDual2<T>;
    #[inline]
    fn neg(self) -> LabeledDual2<T> {
        LabeledDual2 {
            inner: -self.inner,
            seeded: self.seeded,
            #[cfg(debug_assertions)]
            gen_id: self.gen_id,
        }
    }
}

// ============ Scalar-on-LHS impls (f64 and f32) ============
// The inner `Dual2<T>` only provides owned `f64 op Dual2<f64>` / `f32 op
// Dual2<f32>` (no ref variants). The labeled ref variants dereference the
// `Copy` inner to call the owned op.

macro_rules! __lbld2_scalar_lhs {
    ($scalar:ty) => {
        impl ::core::ops::Add<LabeledDual2<$scalar>> for $scalar {
            type Output = LabeledDual2<$scalar>;
            #[inline]
            fn add(self, rhs: LabeledDual2<$scalar>) -> LabeledDual2<$scalar> {
                LabeledDual2 {
                    inner: self + rhs.inner,
                    seeded: rhs.seeded,
                    #[cfg(debug_assertions)]
                    gen_id: rhs.gen_id,
                }
            }
        }
        impl ::core::ops::Add<&LabeledDual2<$scalar>> for $scalar {
            type Output = LabeledDual2<$scalar>;
            #[inline]
            fn add(self, rhs: &LabeledDual2<$scalar>) -> LabeledDual2<$scalar> {
                LabeledDual2 {
                    inner: self + rhs.inner,
                    seeded: rhs.seeded,
                    #[cfg(debug_assertions)]
                    gen_id: rhs.gen_id,
                }
            }
        }
        impl ::core::ops::Sub<LabeledDual2<$scalar>> for $scalar {
            type Output = LabeledDual2<$scalar>;
            #[inline]
            fn sub(self, rhs: LabeledDual2<$scalar>) -> LabeledDual2<$scalar> {
                LabeledDual2 {
                    inner: self - rhs.inner,
                    seeded: rhs.seeded,
                    #[cfg(debug_assertions)]
                    gen_id: rhs.gen_id,
                }
            }
        }
        impl ::core::ops::Sub<&LabeledDual2<$scalar>> for $scalar {
            type Output = LabeledDual2<$scalar>;
            #[inline]
            fn sub(self, rhs: &LabeledDual2<$scalar>) -> LabeledDual2<$scalar> {
                LabeledDual2 {
                    inner: self - rhs.inner,
                    seeded: rhs.seeded,
                    #[cfg(debug_assertions)]
                    gen_id: rhs.gen_id,
                }
            }
        }
        impl ::core::ops::Mul<LabeledDual2<$scalar>> for $scalar {
            type Output = LabeledDual2<$scalar>;
            #[inline]
            fn mul(self, rhs: LabeledDual2<$scalar>) -> LabeledDual2<$scalar> {
                LabeledDual2 {
                    inner: self * rhs.inner,
                    seeded: rhs.seeded,
                    #[cfg(debug_assertions)]
                    gen_id: rhs.gen_id,
                }
            }
        }
        impl ::core::ops::Mul<&LabeledDual2<$scalar>> for $scalar {
            type Output = LabeledDual2<$scalar>;
            #[inline]
            fn mul(self, rhs: &LabeledDual2<$scalar>) -> LabeledDual2<$scalar> {
                LabeledDual2 {
                    inner: self * rhs.inner,
                    seeded: rhs.seeded,
                    #[cfg(debug_assertions)]
                    gen_id: rhs.gen_id,
                }
            }
        }
        impl ::core::ops::Div<LabeledDual2<$scalar>> for $scalar {
            type Output = LabeledDual2<$scalar>;
            #[inline]
            fn div(self, rhs: LabeledDual2<$scalar>) -> LabeledDual2<$scalar> {
                LabeledDual2 {
                    inner: self / rhs.inner,
                    seeded: rhs.seeded,
                    #[cfg(debug_assertions)]
                    gen_id: rhs.gen_id,
                }
            }
        }
        impl ::core::ops::Div<&LabeledDual2<$scalar>> for $scalar {
            type Output = LabeledDual2<$scalar>;
            #[inline]
            fn div(self, rhs: &LabeledDual2<$scalar>) -> LabeledDual2<$scalar> {
                LabeledDual2 {
                    inner: self / rhs.inner,
                    seeded: rhs.seeded,
                    #[cfg(debug_assertions)]
                    gen_id: rhs.gen_id,
                }
            }
        }
    };
}

__lbld2_scalar_lhs!(f64);
__lbld2_scalar_lhs!(f32);