warp-types 0.3.2

//! Inactive Lane Handling: Preventing Reads from Masked-Off Lanes
//!
//! **STATUS: Validated** — Research exploration complete. See conclusions below.
//!
//! THE PROBLEM:
//! When a warp diverges, inactive lanes don't execute. Their register values
//! are stale/garbage. Reading from inactive lanes (via shuffle) is undefined.
//!
//! ```ignore
//! if lane_id % 2 == 0 {
//!     x = compute();  // Only even lanes execute
//! }
//! // x in odd lanes is GARBAGE
//!
//! y = shuffle_xor(x, 1);  // BUG! Even lanes read garbage from odd lanes
//! ```
//!
//! This module explores type-safe solutions to prevent inactive lane reads.

use std::marker::PhantomData;

// ============================================================================
// APPROACH 1: DIVERGENT VALUES (Track which lanes have valid data)
// ============================================================================

/// A value that exists only in lanes within active set S.
///
/// Key insight: After divergence, values have a "validity mask".
/// The type tracks this mask. Operations that would read invalid lanes fail.
pub mod divergent_values {
    use super::*;

    pub trait ActiveSet {
        const MASK: u32;
        fn name() -> &'static str;
    }

    pub struct All;
    impl ActiveSet for All {
        const MASK: u32 = 0xFFFFFFFF;
        fn name() -> &'static str {
            "All"
        }
    }

    pub struct Even;
    impl ActiveSet for Even {
        const MASK: u32 = 0x55555555;
        fn name() -> &'static str {
            "Even"
        }
    }

    pub struct Odd;
    impl ActiveSet for Odd {
        const MASK: u32 = 0xAAAAAAAA;
        fn name() -> &'static str {
            "Odd"
        }
    }

    /// A value valid only in lanes within S.
    ///
    /// - `Divergent<T, All>` = valid everywhere (equivalent to PerLane<T>)
    /// - `Divergent<T, Even>` = valid only in even lanes
    /// - Reading from odd lanes is UNDEFINED
    #[derive(Clone, Copy)]
    pub struct Divergent<T, S: ActiveSet> {
        value: T,
        _marker: PhantomData<S>,
    }

    impl<T, S: ActiveSet> Divergent<T, S> {
        /// Create a divergent value (only valid in active lanes!)
        pub fn new(value: T) -> Self {
            Divergent {
                value,
                _marker: PhantomData,
            }
        }

        /// Read the value (caller must be in an active lane!)
        ///
        /// SAFETY: This is always "safe" in Rust terms, but semantically
        /// the value is only meaningful in lanes where S is active.
        pub fn get(&self) -> &T {
            &self.value
        }

        /// Get the validity mask
        pub fn valid_mask() -> u32 {
            S::MASK
        }
    }

    /// Shuffle that respects validity masks.
    ///
    /// KEY RULE: Can only shuffle if ALL source lanes have valid data.
    ///
    /// shuffle_xor with mask 1 reads from lane (i ^ 1).
    /// For even lanes, source is odd. For odd lanes, source is even.
    /// So we need BOTH even and odd lanes to have valid data.
    ///
    /// This means: shuffle_xor on Divergent<T, Even> is FORBIDDEN
    /// because odd lanes (sources for even lanes) have garbage.
    pub trait CanShuffle<Perm> {
        type Output;
    }

    /// XOR shuffle permutation
    pub struct XorPerm<const MASK: u32>;

    /// shuffle_xor(1) on Divergent<T, All> is OK - all source lanes valid
    impl<T> CanShuffle<XorPerm<1>> for Divergent<T, All> {
        type Output = Divergent<T, All>;
    }

    // shuffle_xor(1) on Divergent<T, Even> is FORBIDDEN
    // (would read from odd lanes which have garbage)
    // NO IMPL - this is a compile error!

    // shuffle_xor(2) on Divergent<T, Even> is... complicated
    // Lane 0 reads from lane 2 (even, valid)
    // Lane 2 reads from lane 0 (even, valid)
    // But this only works because XOR(2) stays within even lanes!

    /// Safe shuffle: only when permutation stays within active set
    pub fn shuffle_within<T: Copy, S: ActiveSet, const MASK: u32>(
        data: Divergent<T, S>,
    ) -> Divergent<T, S>
where
        // Compile-time check: XOR(MASK) maps S to S
        // i.e., for all i in S, (i ^ MASK) is also in S
    {
        // This would need const evaluation or a trait bound
        // For now, placeholder
        data
    }

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

        #[test]
        fn test_divergent_values() {
            let _all_valid: Divergent<i32, All> = Divergent::new(42);
            let _even_valid: Divergent<i32, Even> = Divergent::new(42);

            assert_eq!(Divergent::<i32, All>::valid_mask(), 0xFFFFFFFF);
            assert_eq!(Divergent::<i32, Even>::valid_mask(), 0x55555555);
        }
    }
}

// ============================================================================
// APPROACH 2: SHUFFLE WITH EXPLICIT MASK (Runtime check)
// ============================================================================

/// Instead of type-level tracking, use explicit masks at shuffle time.
///
/// Pro: More flexible, handles dynamic patterns
/// Con: Runtime overhead, can't catch all bugs statically
pub mod explicit_mask {

    /// Per-lane value (no validity tracking)
    #[derive(Clone, Copy)]
    pub struct PerLane<T>(pub T);

    /// Masked shuffle: only read from lanes in `valid_mask`
    ///
    /// If source lane not in mask, returns `default` instead of garbage.
    pub fn shuffle_xor_masked<T: Copy>(
        data: PerLane<T>,
        xor_mask: u32,
        valid_mask: u32,
        default: T,
        lane_id: u32,
    ) -> T {
        let source_lane = lane_id ^ xor_mask;
        let source_valid = (valid_mask >> source_lane) & 1 != 0;

        if source_valid {
            // In real GPU: __shfl_xor_sync(valid_mask, data, xor_mask)
            data.0 // Placeholder
        } else {
            default
        }
    }

    /// Variant: return Option<T> instead of default
    pub fn shuffle_xor_checked<T: Copy>(
        data: PerLane<T>,
        xor_mask: u32,
        valid_mask: u32,
        lane_id: u32,
    ) -> Option<T> {
        let source_lane = lane_id ^ xor_mask;
        let source_valid = (valid_mask >> source_lane) & 1 != 0;

        if source_valid {
            Some(data.0)
        } else {
            None
        }
    }

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

        #[test]
        fn test_masked_shuffle() {
            let data = PerLane(42);
            let even_mask = 0x55555555u32;

            // Lane 0 (even) reading from lane 1 (odd) - invalid!
            let result = shuffle_xor_checked(data, 1, even_mask, 0);
            assert!(result.is_none());

            // Lane 0 (even) reading from lane 2 (even) - valid!
            let result = shuffle_xor_checked(data, 2, even_mask, 0);
            assert!(result.is_some());
        }
    }
}

// ============================================================================
// APPROACH 3: SENTINEL VALUES (Propagating invalidity)
// ============================================================================

/// Use a sentinel value (like NaN or None) for inactive lanes.
/// Operations propagate the sentinel.
///
/// Pro: No type complexity, works with existing code patterns
/// Con: Runtime overhead, silent propagation
pub mod sentinel {

    /// A value that might be invalid (from inactive lane)
    #[derive(Clone, Copy, Debug, PartialEq)]
    pub enum MaybeValid<T> {
        Valid(T),
        Invalid, // From inactive lane
    }

    impl<T> MaybeValid<T> {
        pub fn map<U>(self, f: impl FnOnce(T) -> U) -> MaybeValid<U> {
            match self {
                MaybeValid::Valid(x) => MaybeValid::Valid(f(x)),
                MaybeValid::Invalid => MaybeValid::Invalid,
            }
        }

        pub fn unwrap_or(self, default: T) -> T {
            match self {
                MaybeValid::Valid(x) => x,
                MaybeValid::Invalid => default,
            }
        }
    }

    /// Initialize inactive lanes with Invalid sentinel
    pub fn diverge_with_sentinel<T>(value: T, lane_id: u32, active_mask: u32) -> MaybeValid<T> {
        if (active_mask >> lane_id) & 1 != 0 {
            MaybeValid::Valid(value)
        } else {
            MaybeValid::Invalid
        }
    }

    /// Shuffle that propagates Invalid
    pub fn shuffle_xor_sentinel<T: Copy>(
        data: MaybeValid<T>,
        _xor_mask: u32,
        _lane_id: u32,
    ) -> MaybeValid<T> {
        // In real GPU: shuffle, then check if source was valid
        // If source had Invalid, result is Invalid
        data // Placeholder - would need actual shuffle
    }

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

        #[test]
        fn test_sentinel_propagation() {
            let valid = MaybeValid::Valid(42);
            let invalid: MaybeValid<i32> = MaybeValid::Invalid;

            // Operations on valid values work
            assert_eq!(valid.map(|x| x * 2), MaybeValid::Valid(84));

            // Operations on invalid propagate
            assert_eq!(invalid.map(|x| x * 2), MaybeValid::Invalid);

            // Can provide defaults
            assert_eq!(valid.unwrap_or(0), 42);
            assert_eq!(invalid.unwrap_or(0), 0);
        }
    }
}

// ============================================================================
// APPROACH 4: WARP-RESTRICTED OPERATIONS (Type-level active set)
// ============================================================================

/// The cleanest approach: Warp<S> only provides operations valid for S.
///
/// - Warp<All> has shuffle_xor (all permutations safe)
/// - Warp<Even> has shuffle_within_even (only safe permutations)
/// - Cross-set shuffles require explicit merge first
pub mod warp_restricted {
    use super::*;

    pub trait ActiveSet {
        const MASK: u32;
    }

    pub struct All;
    impl ActiveSet for All {
        const MASK: u32 = 0xFFFFFFFF;
    }

    pub struct Even;
    impl ActiveSet for Even {
        const MASK: u32 = 0x55555555;
    }

    pub struct Odd;
    impl ActiveSet for Odd {
        const MASK: u32 = 0xAAAAAAAA;
    }

    /// Warp with typed active set
    pub struct Warp<S: ActiveSet> {
        _marker: PhantomData<S>,
    }

    impl<S: ActiveSet> Warp<S> {
        pub fn new() -> Self {
            Warp {
                _marker: PhantomData,
            }
        }
    }

    /// Full shuffles only on Warp<All>
    impl Warp<All> {
        /// Any XOR shuffle is safe - all lanes valid
        pub fn shuffle_xor<T: Copy>(&self, data: T, _mask: u32) -> T {
            data // Placeholder
        }

        /// Any permutation shuffle is safe
        pub fn shuffle_idx<T: Copy>(&self, data: T, _source: u32) -> T {
            data // Placeholder
        }
    }

    /// Restricted shuffles on Warp<Even>
    impl Warp<Even> {
        /// XOR shuffles that stay within even lanes
        ///
        /// Safe masks: 2, 4, 6, 8, 10, 12, 14, ... (even numbers)
        /// Unsafe masks: 1, 3, 5, 7, ... (would read from odd lanes)
        pub fn shuffle_xor_within<T: Copy>(&self, data: T, mask: u32) -> Option<T> {
            // Check at runtime that mask keeps us within even lanes
            // For XOR shuffle, this means mask must be even
            if mask % 2 == 0 {
                Some(data) // Safe
            } else {
                None // Would read from odd lanes
            }
        }

        /// Broadcast from lane 0 (always in Even)
        pub fn broadcast_from_0<T: Copy>(&self, data: T) -> T {
            data
        }
    }

    /// Restricted shuffles on Warp<Odd>
    impl Warp<Odd> {
        /// XOR shuffles that stay within odd lanes
        pub fn shuffle_xor_within<T: Copy>(&self, data: T, mask: u32) -> Option<T> {
            if mask % 2 == 0 {
                Some(data) // Safe - stays within odd lanes
            } else {
                None // Would read from even lanes
            }
        }

        /// Broadcast from lane 1 (first odd lane)
        pub fn broadcast_from_1<T: Copy>(&self, data: T) -> T {
            data
        }
    }

    /// KEY INSIGHT: To do a full shuffle after divergence, you MUST merge first.
    ///
    /// ```ignore
    /// let (evens, odds) = warp.diverge();
    /// // evens.shuffle_xor(1) - COMPILE ERROR, method doesn't exist!
    ///
    /// let merged = merge(evens, odds);  // Back to Warp<All>
    /// merged.shuffle_xor(1)  // OK, all lanes valid
    /// ```
    pub fn merge<S1: ActiveSet, S2: ActiveSet>(_left: Warp<S1>, _right: Warp<S2>) -> Warp<All> {
        // In reality, would need ComplementOf bound
        Warp::new()
    }

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

        #[test]
        fn test_warp_all_has_full_shuffle() {
            let warp: Warp<All> = Warp::new();
            let data = 42;

            // All XOR masks work
            let _ = warp.shuffle_xor(data, 1);
            let _ = warp.shuffle_xor(data, 5);
            let _ = warp.shuffle_xor(data, 31);
        }

        #[test]
        fn test_warp_even_restricted() {
            let warp: Warp<Even> = Warp::new();
            let data = 42;

            // Even masks work (stay within even lanes)
            assert!(warp.shuffle_xor_within(data, 2).is_some());
            assert!(warp.shuffle_xor_within(data, 4).is_some());

            // Odd masks rejected (would read from odd lanes)
            assert!(warp.shuffle_xor_within(data, 1).is_none());
            assert!(warp.shuffle_xor_within(data, 3).is_none());
        }

        #[test]
        fn test_must_merge_for_full_shuffle() {
            let warp_even: Warp<Even> = Warp::new();
            let warp_odd: Warp<Odd> = Warp::new();

            // warp_even.shuffle_xor(1) - doesn't exist!

            // Merge first
            let warp_all = merge(warp_even, warp_odd);

            // Now full shuffle works
            let _ = warp_all.shuffle_xor(42, 1);
        }
    }
}

// ============================================================================
// APPROACH 5: HYBRID - Type tracking + Runtime validation
// ============================================================================

/// Combine static and dynamic checking for defense in depth.
pub mod hybrid {
    use super::*;

    pub trait ActiveSet: Copy {
        const MASK: u32;
    }

    #[derive(Copy, Clone)]
    pub struct All;
    impl ActiveSet for All {
        const MASK: u32 = 0xFFFFFFFF;
    }

    #[derive(Copy, Clone)]
    pub struct Even;
    impl ActiveSet for Even {
        const MASK: u32 = 0x55555555;
    }

    /// Value with both static type AND runtime mask
    #[derive(Clone, Copy)]
    pub struct TrackedValue<T, S: ActiveSet> {
        value: T,
        /// Runtime mask for defense in depth
        /// Should always equal S::MASK, but checked at runtime too
        runtime_mask: u32,
        _marker: PhantomData<S>,
    }

    impl<T, S: ActiveSet> TrackedValue<T, S> {
        pub fn new(value: T) -> Self {
            TrackedValue {
                value,
                runtime_mask: S::MASK,
                _marker: PhantomData,
            }
        }

        /// Verify runtime mask matches static type
        pub fn verify(&self) -> bool {
            self.runtime_mask == S::MASK
        }
    }

    /// Shuffle with both static and runtime checks
    pub fn checked_shuffle_xor<T: Copy, S: ActiveSet>(
        data: TrackedValue<T, S>,
        xor_mask: u32,
        lane_id: u32,
    ) -> Result<T, &'static str> {
        // Runtime check: source lane must be in valid mask
        let source_lane = lane_id ^ xor_mask;
        if (data.runtime_mask >> source_lane) & 1 == 0 {
            return Err("Source lane is inactive");
        }

        Ok(data.value)
    }

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

        #[test]
        fn test_hybrid_checking() {
            let data: TrackedValue<i32, Even> = TrackedValue::new(42);

            // Lane 0 reading from lane 1 (odd) - rejected
            assert!(checked_shuffle_xor(data, 1, 0).is_err());

            // Lane 0 reading from lane 2 (even) - accepted
            assert!(checked_shuffle_xor(data, 2, 0).is_ok());
        }
    }
}

// ============================================================================
// SUMMARY: Recommended Approach
// ============================================================================
//
// For warp-types, recommend APPROACH 4 (Warp-Restricted Operations):
//
// 1. Warp<All> has all shuffle methods
// 2. Warp<S> for S ≠ All has only "within-set" shuffles
// 3. To do cross-set operations, must merge first
//
// This is:
// - Statically enforced (compile-time errors)
// - Zero runtime overhead in happy path
// - Matches mental model (diverge restricts, merge restores)
// - Already implemented in our existing prototype!
//
// The key insight: We don't need to track VALUE validity separately.
// We track WARP validity. If you have Warp<Even>, the warp itself
// only provides operations safe for even lanes. Values inherit this.
//
// Invalid patterns become compile errors:
// ```compile_fail
// let (evens, odds) = warp.diverge();
// evens.shuffle_xor(1);  // ERROR: method doesn't exist on Warp<Even>
// ```

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

    #[test]
    fn test_recommended_pattern() {
        // Start with all lanes
        let warp: Warp<All> = Warp::new();

        // Full shuffle works
        let data = 42;
        let _ = warp.shuffle_xor(data, 1);

        // After divergence, only restricted shuffles
        let warp_even: Warp<Even> = Warp::new(); // Simulating diverge
        let warp_odd: Warp<Odd> = Warp::new();

        // warp_even.shuffle_xor(1) - doesn't compile!
        // Must use within-set shuffles or merge first

        // Merge restores full capability
        let warp_merged = merge(warp_even, warp_odd);
        let _ = warp_merged.shuffle_xor(data, 1); // Works again
    }
}