xad-rs 0.2.0

//! `LabeledDual` — labeled wrapper over `Dual` (f64-only forward-mode).
//!
//! **Shape A (Phase 02.2):** does NOT carry an `Arc<VarRegistry>` field in
//! release builds. The struct layout is a single `inner: Dual` field plus,
//! under `#[cfg(debug_assertions)]` only, a `gen_id: u64` stamped by the
//! owning [`LabeledForwardTape`] scope for the cross-registry debug guard.
//! Release builds are bit-for-bit equivalent to a pure `Dual` wrapper and
//! carry zero atomic-refcount cost per operator.
//!
//! The only way to obtain a `LabeledDual` is via the `LabeledForwardTape`
//! constructor API (see `LabeledForwardTape::declare_dual` /
//! `LabeledForwardTape::freeze_dual` — the exact shape is resolved in
//! Plan 02.2-02 Task 3a and implemented in Task 3b).

use std::fmt;

use crate::dual::Dual;

/// Labeled wrapper around the positional [`Dual`] type.
///
/// # Example
///
/// The constructor API for `LabeledDual` is owned by `LabeledForwardTape`
/// via the `declare_dual` / `freeze_dual` / `scope.dual(handle)` pattern.
/// See the module docs in `src/labeled/forward_tape.rs` for the full
/// rationale.
///
/// ```
/// use xad_rs::labeled::{LabeledForwardTape, LabeledForwardScope};
///
/// let mut ft = LabeledForwardTape::new();
/// let x_h = ft.declare_dual("x", 2.0);
/// let y_h = ft.declare_dual("y", 3.0);
/// let scope: LabeledForwardScope = ft.freeze_dual();
///
/// let x = scope.dual(x_h);
/// let y = scope.dual(y_h);
/// let f = x * y + x - 2.0 * y;
///
/// assert_eq!(f.real(), 2.0);
/// assert_eq!(f.partial("x"), 4.0);
/// assert_eq!(f.partial("y"), 0.0);
/// ```
#[derive(Clone, Debug)]
pub struct LabeledDual {
    pub(super) inner: Dual,
    // NOTE: field name is `gen_id` — `gen` alone is a reserved keyword in
    // Rust 2024 edition. The D-01/D-02 CONTEXT blocks spell it as `gen`;
    // we carry the spelling adjustment forward to satisfy the compiler.
    #[cfg(debug_assertions)]
    pub(super) gen_id: u64,
}

impl LabeledDual {
    /// 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_inner(inner: Dual) -> Self {
        Self {
            inner,
            #[cfg(debug_assertions)]
            gen_id: crate::labeled::forward_tape::current_gen(),
        }
    }

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

    /// Partial derivative with respect to a named variable.
    ///
    /// Reads the active registry from the `LabeledForwardTape` thread-local
    /// slot (stamped at `freeze_dual()` time). Returns `0.0` if `name` is
    /// in the registry but not touched by this value (same as positional
    /// `Dual::partial` on an unused slot). Panics if `name` is not in the
    /// registry, or if called outside a frozen `LabeledForwardTape` scope.
    pub fn partial(&self, name: &str) -> f64 {
        let idx = crate::labeled::forward_tape::with_active_registry(|r| {
            let r =
                r.expect("LabeledDual::partial called outside a frozen LabeledForwardTape scope");
            r.index_of(name).unwrap_or_else(|| {
                panic!(
                    "LabeledDual::partial: name {:?} not present in registry",
                    name
                )
            })
        });
        self.inner.partial(idx)
    }

    /// Full gradient as a `Vec<(name, partial)>`.
    ///
    /// Iteration order matches the active registry's insertion order —
    /// deterministic. Reads the active registry from the thread-local slot.
    /// Panics if called outside a frozen `LabeledForwardTape` scope.
    pub fn gradient(&self) -> Vec<(String, f64)> {
        crate::labeled::forward_tape::with_active_registry(|r| {
            let r =
                r.expect("LabeledDual::gradient called outside a frozen LabeledForwardTape scope");
            let n = r.len();
            let mut out = Vec::with_capacity(n);
            for (i, name) in r.iter().enumerate() {
                out.push((name.to_string(), self.inner.partial(i)));
            }
            out
        })
    }

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

    // ============ Elementary math delegations ============
    // Each method forwards to the inherent `Dual` elementary and preserves
    // the parent's generation stamp explicitly (debug builds) — no TLS read
    // on the hot path.

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

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

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

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

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

    /// Tangent, preserving the parent scope's generation.
    #[inline]
    pub fn tan(&self) -> Self {
        Self {
            inner: self.inner.tan(),
            #[cfg(debug_assertions)]
            gen_id: self.gen_id,
        }
    }
}

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

// ============ Operator overloads — hand-written, Shape A ============
// No shared op-stamping macro is used: Shape A does not carry an
// `Arc` registry field on the per-value wrapper. The four reference
// variants (owned/owned, ref/ref, owned/ref, ref/owned) plus scalar-RHS
// variants are stamped explicitly via a local `__lbl_dual_binop!`
// macro, modelled on `__lbl_freal_binop!` in `src/labeled/freal.rs`
// but specialised for non-generic `LabeledDual` with `f64` as the
// scalar RHS. Each wrapper-vs-wrapper impl performs a debug-only
// `check_gen` between the two operands' generations, then constructs
// the result preserving the LHS's generation stamp.

macro_rules! __lbl_dual_binop {
    ($trait:ident, $method:ident, $op:tt) => {
        impl ::core::ops::$trait<LabeledDual> for LabeledDual {
            type Output = LabeledDual;
            #[inline]
            fn $method(self, rhs: LabeledDual) -> LabeledDual {
                #[cfg(debug_assertions)]
                crate::labeled::forward_tape::check_gen(self.gen_id, rhs.gen_id);
                LabeledDual {
                    inner: self.inner $op rhs.inner,
                    #[cfg(debug_assertions)]
                    gen_id: self.gen_id,
                }
            }
        }
        impl ::core::ops::$trait<&LabeledDual> for &LabeledDual {
            type Output = LabeledDual;
            #[inline]
            fn $method(self, rhs: &LabeledDual) -> LabeledDual {
                #[cfg(debug_assertions)]
                crate::labeled::forward_tape::check_gen(self.gen_id, rhs.gen_id);
                LabeledDual {
                    inner: &self.inner $op &rhs.inner,
                    #[cfg(debug_assertions)]
                    gen_id: self.gen_id,
                }
            }
        }
        impl ::core::ops::$trait<&LabeledDual> for LabeledDual {
            type Output = LabeledDual;
            #[inline]
            fn $method(self, rhs: &LabeledDual) -> LabeledDual {
                #[cfg(debug_assertions)]
                crate::labeled::forward_tape::check_gen(self.gen_id, rhs.gen_id);
                LabeledDual {
                    inner: self.inner $op &rhs.inner,
                    #[cfg(debug_assertions)]
                    gen_id: self.gen_id,
                }
            }
        }
        impl ::core::ops::$trait<LabeledDual> for &LabeledDual {
            type Output = LabeledDual;
            #[inline]
            fn $method(self, rhs: LabeledDual) -> LabeledDual {
                #[cfg(debug_assertions)]
                crate::labeled::forward_tape::check_gen(self.gen_id, rhs.gen_id);
                LabeledDual {
                    inner: &self.inner $op rhs.inner,
                    #[cfg(debug_assertions)]
                    gen_id: self.gen_id,
                }
            }
        }
        impl ::core::ops::$trait<f64> for LabeledDual {
            type Output = LabeledDual;
            #[inline]
            fn $method(self, rhs: f64) -> LabeledDual {
                LabeledDual {
                    inner: self.inner $op rhs,
                    #[cfg(debug_assertions)]
                    gen_id: self.gen_id,
                }
            }
        }
        impl ::core::ops::$trait<f64> for &LabeledDual {
            type Output = LabeledDual;
            #[inline]
            fn $method(self, rhs: f64) -> LabeledDual {
                LabeledDual {
                    inner: &self.inner $op rhs,
                    #[cfg(debug_assertions)]
                    gen_id: self.gen_id,
                }
            }
        }
    };
}

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

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

// ============ Scalar-on-LHS hand-written impls (orphan-rule escape) ============
// The inner `Dual` type supports the full `f64 op Dual` + `f64 op &Dual`
// surface (see src/dual.rs), so delegation is direct. Each result preserves
// the RHS's generation stamp (debug builds only).

// Add
impl ::core::ops::Add<LabeledDual> for f64 {
    type Output = LabeledDual;
    #[inline]
    fn add(self, rhs: LabeledDual) -> LabeledDual {
        LabeledDual {
            inner: self + rhs.inner,
            #[cfg(debug_assertions)]
            gen_id: rhs.gen_id,
        }
    }
}
impl ::core::ops::Add<&LabeledDual> for f64 {
    type Output = LabeledDual;
    #[inline]
    fn add(self, rhs: &LabeledDual) -> LabeledDual {
        LabeledDual {
            inner: self + &rhs.inner,
            #[cfg(debug_assertions)]
            gen_id: rhs.gen_id,
        }
    }
}

// Sub
impl ::core::ops::Sub<LabeledDual> for f64 {
    type Output = LabeledDual;
    #[inline]
    fn sub(self, rhs: LabeledDual) -> LabeledDual {
        LabeledDual {
            inner: self - rhs.inner,
            #[cfg(debug_assertions)]
            gen_id: rhs.gen_id,
        }
    }
}
impl ::core::ops::Sub<&LabeledDual> for f64 {
    type Output = LabeledDual;
    #[inline]
    fn sub(self, rhs: &LabeledDual) -> LabeledDual {
        LabeledDual {
            inner: self - &rhs.inner,
            #[cfg(debug_assertions)]
            gen_id: rhs.gen_id,
        }
    }
}

// Mul
impl ::core::ops::Mul<LabeledDual> for f64 {
    type Output = LabeledDual;
    #[inline]
    fn mul(self, rhs: LabeledDual) -> LabeledDual {
        LabeledDual {
            inner: self * rhs.inner,
            #[cfg(debug_assertions)]
            gen_id: rhs.gen_id,
        }
    }
}
impl ::core::ops::Mul<&LabeledDual> for f64 {
    type Output = LabeledDual;
    #[inline]
    fn mul(self, rhs: &LabeledDual) -> LabeledDual {
        LabeledDual {
            inner: self * &rhs.inner,
            #[cfg(debug_assertions)]
            gen_id: rhs.gen_id,
        }
    }
}

// Div
impl ::core::ops::Div<LabeledDual> for f64 {
    type Output = LabeledDual;
    #[inline]
    fn div(self, rhs: LabeledDual) -> LabeledDual {
        LabeledDual {
            inner: self / rhs.inner,
            #[cfg(debug_assertions)]
            gen_id: rhs.gen_id,
        }
    }
}
impl ::core::ops::Div<&LabeledDual> for f64 {
    type Output = LabeledDual;
    #[inline]
    fn div(self, rhs: &LabeledDual) -> LabeledDual {
        LabeledDual {
            inner: self / &rhs.inner,
            #[cfg(debug_assertions)]
            gen_id: rhs.gen_id,
        }
    }
}