xad-rs 0.8.0

//! `NamedForwardTape<T>` — declare-and-scope lifecycle for named
//! forward-mode values.
//!
//! Build phase: declare named inputs against a mutable
//! [`NamedForwardTape<T>`]. Each `declare_*` returns an opaque handle.
//! Lock-and-use phase: consume the tape via
//! [`NamedForwardTape::into_scope`] to obtain a
//! [`NamedForwardScope<T>`]; the scope owns the constructed values and
//! exposes them via the handles.
//!
//! Both types are generic over a single `T: Passive` (defaulting to
//! `f64`). The TLS slot used for the cross-registry debug guard is
//! scalar-type-independent and is shared across all `T`.

use std::cell::Cell;
use std::fmt;
use std::marker::PhantomData;
use std::sync::Arc;
#[cfg(debug_assertions)]
use std::sync::atomic::{AtomicU64, Ordering};

use indexmap::IndexSet;

use crate::passive::Passive;
use crate::registry::VarRegistry;

// === TLS generation counter (debug-only) ===

#[cfg(debug_assertions)]
static NEXT_GEN: AtomicU64 = AtomicU64::new(1);

thread_local! {
    #[cfg(debug_assertions)]
    static ACTIVE_GEN: Cell<u64> = const { Cell::new(0) };

    static ACTIVE_REGISTRY: Cell<*const VarRegistry> = const { Cell::new(std::ptr::null()) };
}

pub(crate) fn with_active_registry<R>(f: impl FnOnce(Option<&VarRegistry>) -> R) -> R {
    ACTIVE_REGISTRY.with(|c| {
        let ptr = c.get();
        let reg_ref: Option<&VarRegistry> = if ptr.is_null() {
            None
        } else {
            Some(unsafe { &*ptr })
        };
        f(reg_ref)
    })
}

#[cfg(debug_assertions)]
#[inline(always)]
pub(crate) fn current_gen() -> u64 {
    ACTIVE_GEN.with(|c| c.get())
}

#[cfg(debug_assertions)]
#[inline(always)]
pub(crate) fn check_gen(lhs: u64, rhs: u64) {
    assert_eq!(
        lhs, rhs,
        "xad_rs::named: cross-registry forward-mode op detected (lhs tape generation = {lhs}, rhs tape generation = {rhs}). \
         Both operands must come from the same NamedForwardTape scope."
    );
}

#[cfg(not(debug_assertions))]
#[inline(always)]
#[allow(dead_code)]
pub(crate) fn check_gen(_lhs: (), _rhs: ()) {}

// === Handles ===
//
// Handles carry the scalar type `T` so the wrong scope cannot be
// indexed with a foreign-T handle. `PhantomData<fn() -> T>` is the
// covariant, drop-check-free encoding — `T: Passive` is `'static` so
// either form is sound, but `fn() -> T` is the conventional choice for
// "this handle stands in for a value of type T".

/// Handle to a `NamedJet1<T>` declared on a `NamedForwardTape<T>`.
pub struct Jet1Handle<T: Passive = f64> {
    idx: usize,
    _t: PhantomData<fn() -> T>,
}

impl<T: Passive> Clone for Jet1Handle<T> {
    fn clone(&self) -> Self {
        *self
    }
}
impl<T: Passive> Copy for Jet1Handle<T> {}
impl<T: Passive> fmt::Debug for Jet1Handle<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Jet1Handle").field("idx", &self.idx).finish()
    }
}
impl<T: Passive> PartialEq for Jet1Handle<T> {
    fn eq(&self, other: &Self) -> bool {
        self.idx == other.idx
    }
}
impl<T: Passive> Eq for Jet1Handle<T> {}

/// Handle to a `NamedJet1Vec<T>` declared on a `NamedForwardTape<T>`.
pub struct Jet1VecHandle<T: Passive = f64> {
    idx: usize,
    _t: PhantomData<fn() -> T>,
}

impl<T: Passive> Clone for Jet1VecHandle<T> {
    fn clone(&self) -> Self {
        *self
    }
}
impl<T: Passive> Copy for Jet1VecHandle<T> {}
impl<T: Passive> fmt::Debug for Jet1VecHandle<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Jet1VecHandle")
            .field("idx", &self.idx)
            .finish()
    }
}
impl<T: Passive> PartialEq for Jet1VecHandle<T> {
    fn eq(&self, other: &Self) -> bool {
        self.idx == other.idx
    }
}
impl<T: Passive> Eq for Jet1VecHandle<T> {}

/// Handle to a `NamedJet2<T>` declared on a `NamedForwardTape<T>`.
pub struct Jet2Handle<T: Passive = f64> {
    idx: usize,
    _t: PhantomData<fn() -> T>,
}

impl<T: Passive> Clone for Jet2Handle<T> {
    fn clone(&self) -> Self {
        *self
    }
}
impl<T: Passive> Copy for Jet2Handle<T> {}
impl<T: Passive> fmt::Debug for Jet2Handle<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Jet2Handle").field("idx", &self.idx).finish()
    }
}
impl<T: Passive> PartialEq for Jet2Handle<T> {
    fn eq(&self, other: &Self) -> bool {
        self.idx == other.idx
    }
}
impl<T: Passive> Eq for Jet2Handle<T> {}

// === NamedForwardTape ===

/// Builder used to declare named inputs of scalar type `T`. Convert to
/// a [`NamedForwardScope<T>`] via [`NamedForwardTape::into_scope`] to
/// access the constructed values.
pub struct NamedForwardTape<T: Passive = f64> {
    builder: IndexSet<String>,
    pending_jet1: Vec<(String, T)>,
    pending_jet1vec: Vec<(String, T)>,
    pending_jet2: Vec<(String, T)>,
    #[cfg(debug_assertions)]
    #[allow(dead_code)]
    gen_id: u64,
    #[cfg(debug_assertions)]
    prev_gen: u64,
    /// `true` iff `into_scope` consumed this tape — suppresses the
    /// generation-restore in `Drop` (the scope owns that responsibility).
    consumed: bool,
    _not_send: PhantomData<*const ()>,
}

impl<T: Passive> NamedForwardTape<T> {
    /// Create a new, empty tape. In debug builds this also bumps the
    /// global TLS generation so cross-tape ops can be detected.
    pub fn new() -> Self {
        #[cfg(debug_assertions)]
        let new_gen = NEXT_GEN.fetch_add(1, Ordering::Relaxed);
        #[cfg(debug_assertions)]
        let prev_gen = ACTIVE_GEN.with(|c| {
            let p = c.get();
            c.set(new_gen);
            p
        });
        Self {
            builder: IndexSet::new(),
            pending_jet1: Vec::new(),
            pending_jet1vec: Vec::new(),
            pending_jet2: Vec::new(),
            #[cfg(debug_assertions)]
            gen_id: new_gen,
            #[cfg(debug_assertions)]
            prev_gen,
            consumed: false,
            _not_send: PhantomData,
        }
    }

    /// Declare a named [`crate::forward::jet1::NamedJet1`] input.
    pub fn declare_jet1(&mut self, name: &str, value: T) -> Jet1Handle<T> {
        assert!(
            !self.consumed,
            "NamedForwardTape::declare_jet1({:?}) called after into_scope",
            name
        );
        if !self.builder.contains(name) {
            self.builder.insert(name.to_string());
        }
        let idx = self.pending_jet1.len();
        self.pending_jet1.push((name.to_string(), value));
        Jet1Handle {
            idx,
            _t: PhantomData,
        }
    }

    /// Declare a named [`crate::forward::jet1vec::NamedJet1Vec`] input.
    pub fn declare_jet1vec(&mut self, name: &str, value: T) -> Jet1VecHandle<T> {
        assert!(
            !self.consumed,
            "NamedForwardTape::declare_jet1vec({:?}) called after into_scope",
            name
        );
        if !self.builder.contains(name) {
            self.builder.insert(name.to_string());
        }
        let idx = self.pending_jet1vec.len();
        self.pending_jet1vec.push((name.to_string(), value));
        Jet1VecHandle {
            idx,
            _t: PhantomData,
        }
    }

    /// Declare a named [`crate::forward::jet2::NamedJet2`] input.
    pub fn declare_jet2(&mut self, name: &str, value: T) -> Jet2Handle<T> {
        assert!(
            !self.consumed,
            "NamedForwardTape::declare_jet2({:?}) called after into_scope",
            name
        );
        if !self.builder.contains(name) {
            self.builder.insert(name.to_string());
        }
        let idx = self.pending_jet2.len();
        self.pending_jet2.push((name.to_string(), value));
        Jet2Handle {
            idx,
            _t: PhantomData,
        }
    }

    /// Consume the tape and produce a [`NamedForwardScope<T>`] that
    /// owns all declared values. Activates the registry on the TLS slot
    /// so named-derivative accessors find it.
    pub fn into_scope(mut self) -> NamedForwardScope<T> {
        assert!(
            !self.consumed,
            "NamedForwardTape::into_scope called twice"
        );

        let reg = Arc::new(VarRegistry::from_names(self.builder.iter().cloned()));

        // Activate the registry on the TLS slot, saving the previous pointer
        // so the scope's Drop can restore it.
        let new_ptr: *const VarRegistry = Arc::as_ptr(&reg);
        let prev_registry = ACTIVE_REGISTRY.with(|c| {
            let prev = c.get();
            c.set(new_ptr);
            prev
        });

        // Mark consumed so the tape's Drop is a no-op for the gen restore
        // (the scope now owns that responsibility).
        self.consumed = true;
        #[cfg(debug_assertions)]
        let prev_gen = self.prev_gen;

        // Drain pending lists into owned scope storage.
        let pending_jet1 = std::mem::take(&mut self.pending_jet1);
        let n_jet1 = pending_jet1.len();
        let mut jet1s: Vec<crate::forward::jet1::NamedJet1<T>> = Vec::with_capacity(n_jet1);
        for (_name, value) in pending_jet1.into_iter() {
            // `Jet1<T>` carries one tangent direction; seeding with `1.0`
            // matches the original `input_jet1` semantics.
            let inner = crate::forward::jet1::Jet1::<T>::new(value, T::one());
            jet1s.push(crate::forward::jet1::NamedJet1::<T>::__from_inner(inner));
        }

        let pending_jet1vec = std::mem::take(&mut self.pending_jet1vec);
        let n_jet1vec = pending_jet1vec.len();
        let mut jet1vecs: Vec<crate::forward::jet1vec::NamedJet1Vec<T>> =
            Vec::with_capacity(n_jet1vec);
        for (name, value) in pending_jet1vec.into_iter() {
            let reg_idx = reg
                .index_of(&name)
                .expect("declared name missing from frozen registry");
            let inner = crate::forward::jet1vec::Jet1Vec::<T>::variable(value, reg_idx, n_jet1vec);
            jet1vecs.push(crate::forward::jet1vec::NamedJet1Vec::<T>::__from_inner(
                inner,
            ));
        }

        let pending_jet2 = std::mem::take(&mut self.pending_jet2);
        let mut jet2s: Vec<crate::forward::jet2::NamedJet2<T>> =
            Vec::with_capacity(pending_jet2.len());
        for (name, value) in pending_jet2.into_iter() {
            let reg_idx = reg
                .index_of(&name)
                .expect("declared name missing from frozen registry");
            let inner = crate::forward::jet2::Jet2::<T>::variable(value);
            jet2s.push(crate::forward::jet2::NamedJet2::<T>::__from_parts(
                inner,
                Some(reg_idx),
            ));
        }

        NamedForwardScope {
            registry: reg,
            jet1s,
            jet1vecs,
            jet2s,
            prev_registry,
            #[cfg(debug_assertions)]
            prev_gen,
            _not_send: PhantomData,
        }
    }

    /// Recovery hook for the `std::mem::forget` path. Resets the active
    /// registry pointer and (in debug) the active generation. Static so
    /// it can be called without a live tape instance.
    pub fn deactivate_all() {
        ACTIVE_REGISTRY.with(|c| c.set(std::ptr::null()));
        #[cfg(debug_assertions)]
        ACTIVE_GEN.with(|c| c.set(0));
    }
}

impl<T: Passive> Default for NamedForwardTape<T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T: Passive> fmt::Debug for NamedForwardTape<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("NamedForwardTape")
            .field("inputs", &self.builder.len())
            .field("pending_jet1", &self.pending_jet1.len())
            .field("pending_jet1vec", &self.pending_jet1vec.len())
            .field("pending_jet2", &self.pending_jet2.len())
            .field("consumed", &self.consumed)
            .finish()
    }
}

impl<T: Passive> Drop for NamedForwardTape<T> {
    fn drop(&mut self) {
        // If `into_scope` ran, the scope owns the TLS-restore. If the
        // tape was dropped without calling `into_scope` (no scope was
        // constructed), restore the previous generation here so nested
        // tapes don't leak their generation upward.
        if !self.consumed {
            #[cfg(debug_assertions)]
            ACTIVE_GEN.with(|c| c.set(self.prev_gen));
        }
    }
}

// === NamedForwardScope ===

/// Owns the values declared on a [`NamedForwardTape<T>`] and the active
/// registry. The scope is `!Send` and restores the previous TLS state
/// on drop.
pub struct NamedForwardScope<T: Passive = f64> {
    registry: Arc<VarRegistry>,
    jet1s: Vec<crate::forward::jet1::NamedJet1<T>>,
    jet1vecs: Vec<crate::forward::jet1vec::NamedJet1Vec<T>>,
    jet2s: Vec<crate::forward::jet2::NamedJet2<T>>,
    prev_registry: *const VarRegistry,
    #[cfg(debug_assertions)]
    prev_gen: u64,
    _not_send: PhantomData<*const ()>,
}

impl<T: Passive> NamedForwardScope<T> {
    /// Look up a declared `NamedJet1<T>` by handle.
    #[inline]
    pub fn jet1(&self, handle: Jet1Handle<T>) -> &crate::forward::jet1::NamedJet1<T> {
        &self.jet1s[handle.idx]
    }

    /// Look up a declared `NamedJet1Vec<T>` by handle.
    #[inline]
    pub fn jet1vec(&self, handle: Jet1VecHandle<T>) -> &crate::forward::jet1vec::NamedJet1Vec<T> {
        &self.jet1vecs[handle.idx]
    }

    /// Look up a declared `NamedJet2<T>` by handle.
    #[inline]
    pub fn jet2(&self, handle: Jet2Handle<T>) -> &crate::forward::jet2::NamedJet2<T> {
        &self.jet2s[handle.idx]
    }

    /// Construct a scope-stamped `NamedJet1<T>` constant (zero tangent).
    #[inline]
    pub fn constant_jet1(&self, value: T) -> crate::forward::jet1::NamedJet1<T> {
        let inner = crate::forward::jet1::Jet1::<T>::constant(value);
        crate::forward::jet1::NamedJet1::<T>::__from_inner(inner)
    }

    /// Construct a scope-stamped `NamedJet1Vec<T>` constant (zero gradient,
    /// length matching the active registry).
    #[inline]
    pub fn constant_jet1vec(&self, value: T) -> crate::forward::jet1vec::NamedJet1Vec<T> {
        let inner = crate::forward::jet1vec::Jet1Vec::<T>::constant(value, self.registry.len());
        crate::forward::jet1vec::NamedJet1Vec::<T>::__from_inner(inner)
    }

    /// Construct a scope-stamped `NamedJet2<T>` constant (no seeded direction).
    #[inline]
    pub fn constant_jet2(&self, value: T) -> crate::forward::jet2::NamedJet2<T> {
        let inner = crate::forward::jet2::Jet2::<T>::constant(value);
        crate::forward::jet2::NamedJet2::<T>::__from_parts(inner, None)
    }

    /// Borrow the scope's frozen registry.
    #[inline]
    pub fn registry(&self) -> &Arc<VarRegistry> {
        &self.registry
    }
}

impl<T: Passive> fmt::Debug for NamedForwardScope<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("NamedForwardScope")
            .field("registry_len", &self.registry.len())
            .field("jet1s", &self.jet1s.len())
            .field("jet1vecs", &self.jet1vecs.len())
            .field("jet2s", &self.jet2s.len())
            .finish()
    }
}

impl<T: Passive> Drop for NamedForwardScope<T> {
    fn drop(&mut self) {
        ACTIVE_REGISTRY.with(|c| c.set(self.prev_registry));
        #[cfg(debug_assertions)]
        ACTIVE_GEN.with(|c| c.set(self.prev_gen));
    }
}