warp-types 0.3.2

//! Prototype: Can we verify lane masks statically?
//!
//! **STATUS: Superseded** — Promoted to `src/{active_set,warp,diverge,merge,shuffle}.rs`. Retained as research artifact.
//!
//! Q5: "Can we verify lane masks statically?"
//!
//! Building on Q1 (marker types work), this file demonstrates:
//! 1. Compile-time verification that merge() gets complementary sets
//! 2. Compile-time verification that shuffle() requires all lanes active
//! 3. Compile-time tracking of active sets through diverge/merge
//! 4. Type errors for incorrect usage (shown in comments)
//!
//! GOAL: Prove that session-typed divergence can catch real bugs at compile time.
//!
//! # Compile-Time Safety Guarantees
//!
//! ## Bug 1: Shuffle after diverge (caught)
//!
//! ```compile_fail
//! use warp_types::research::static_verify::*;
//!
//! fn buggy_shuffle() {
//!     let warp: Warp<All> = Warp::new();
//!     let (evens, _odds) = warp.diverge_even_odd();
//!
//!     // BUG: shuffle_xor doesn't exist on Warp<Even>
//!     let data = PerLane([1i32; 32]);
//!     let _ = evens.shuffle_xor(data, 1);  // ERROR: method not found
//! }
//! ```
//!
//! ## Bug 2: Merge non-complements (caught)
//!
//! ```compile_fail
//! use warp_types::research::static_verify::*;
//!
//! fn buggy_merge() {
//!     let evens: Warp<Even> = Warp::new();
//!     let low: Warp<LowHalf> = Warp::new();
//!
//!     // BUG: Even and LowHalf are not complements (they overlap)
//!     let _ = merge(evens, low);  // ERROR: ComplementOf<LowHalf> not implemented for Even
//! }
//! ```
//!
//! ## Bug 3: Merge same set (caught)
//!
//! ```compile_fail
//! use warp_types::research::static_verify::*;
//!
//! fn buggy_merge_same() {
//!     let evens1: Warp<Even> = Warp::new();
//!     let evens2: Warp<Even> = Warp::new();
//!
//!     // BUG: Can't merge Even with Even
//!     let _ = merge(evens1, evens2);  // ERROR: ComplementOf<Even> not implemented for Even
//! }
//! ```

use std::marker::PhantomData;

// ============================================================================
// ACTIVE SET TYPE HIERARCHY
// ============================================================================

/// Marker trait for all active set types
pub trait ActiveSet: Copy + 'static {
    /// The bitmask for this set (for runtime/debugging)
    const MASK: u32;

    /// Human-readable name
    const NAME: &'static str;
}

/// Marker trait: S1 and S2 are complements (disjoint AND cover all lanes)
///
/// This is THE key trait for merge verification.
/// Only implemented for valid complement pairs.
pub trait ComplementOf<Other: ActiveSet>: ActiveSet {}

// ============================================================================
// CONCRETE ACTIVE SET TYPES
// ============================================================================

#[derive(Copy, Clone, Debug, Default)]
pub struct All;
impl ActiveSet for All {
    const MASK: u32 = 0xFFFFFFFF;
    const NAME: &'static str = "All";
}

#[derive(Copy, Clone, Debug, Default)]
pub struct None;
impl ActiveSet for None {
    const MASK: u32 = 0x00000000;
    const NAME: &'static str = "None";
}

#[derive(Copy, Clone, Debug, Default)]
pub struct Even;
impl ActiveSet for Even {
    const MASK: u32 = 0x55555555;
    const NAME: &'static str = "Even";
}

#[derive(Copy, Clone, Debug, Default)]
pub struct Odd;
impl ActiveSet for Odd {
    const MASK: u32 = 0xAAAAAAAA;
    const NAME: &'static str = "Odd";
}

#[derive(Copy, Clone, Debug, Default)]
pub struct LowHalf;
impl ActiveSet for LowHalf {
    const MASK: u32 = 0x0000FFFF;
    const NAME: &'static str = "LowHalf";
}

#[derive(Copy, Clone, Debug, Default)]
pub struct HighHalf;
impl ActiveSet for HighHalf {
    const MASK: u32 = 0xFFFF0000;
    const NAME: &'static str = "HighHalf";
}

#[derive(Copy, Clone, Debug, Default)]
pub struct Lane0;
impl ActiveSet for Lane0 {
    const MASK: u32 = 0x00000001;
    const NAME: &'static str = "Lane0";
}

#[derive(Copy, Clone, Debug, Default)]
pub struct NotLane0;
impl ActiveSet for NotLane0 {
    const MASK: u32 = 0xFFFFFFFE;
    const NAME: &'static str = "NotLane0";
}

// ============================================================================
// COMPLEMENT RELATIONSHIPS (Compile-Time Verified)
// ============================================================================

// Even ↔ Odd
impl ComplementOf<Odd> for Even {}
impl ComplementOf<Even> for Odd {}

// LowHalf ↔ HighHalf
impl ComplementOf<HighHalf> for LowHalf {}
impl ComplementOf<LowHalf> for HighHalf {}

// Lane0 ↔ NotLane0
impl ComplementOf<NotLane0> for Lane0 {}
impl ComplementOf<Lane0> for NotLane0 {}

// All ↔ None (degenerate cases)
impl ComplementOf<None> for All {}
impl ComplementOf<All> for None {}

// ============================================================================
// WARP TYPE WITH ACTIVE SET TRACKING
// ============================================================================

/// A warp with compile-time tracked active lanes
#[derive(Copy, Clone)]
pub struct Warp<S: ActiveSet> {
    // In a real implementation, this would hold per-lane data
    _phantom: PhantomData<S>,
}

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

    pub fn active_set_name(&self) -> &'static str {
        S::NAME
    }

    pub fn active_mask(&self) -> u32 {
        S::MASK
    }
}

impl<S: ActiveSet> Default for Warp<S> {
    fn default() -> Self {
        Self::new()
    }
}

// ============================================================================
// DIVERGE OPERATIONS (Split warp by predicate)
// ============================================================================

impl Warp<All> {
    /// Diverge into even/odd lanes
    ///
    /// Returns two warps with COMPLEMENTARY active sets.
    /// The type system tracks that Even and Odd together make All.
    pub fn diverge_even_odd(self) -> (Warp<Even>, Warp<Odd>) {
        (Warp::new(), Warp::new())
    }

    /// Diverge into low/high halves
    pub fn diverge_halves(self) -> (Warp<LowHalf>, Warp<HighHalf>) {
        (Warp::new(), Warp::new())
    }

    /// Extract lane 0 (like reduction result)
    pub fn extract_lane0(self) -> (Warp<Lane0>, Warp<NotLane0>) {
        (Warp::new(), Warp::new())
    }
}

// Nested diverge: Even lanes can be further split
impl Warp<Even> {
    /// Further split even lanes into low-even and high-even
    ///
    /// Even ∩ LowHalf = lanes 0, 2, 4, 6, 8, 10, 12, 14
    /// Even ∩ HighHalf = lanes 16, 18, 20, 22, 24, 26, 28, 30
    pub fn diverge_halves(self) -> (Warp<EvenLow>, Warp<EvenHigh>) {
        (Warp::new(), Warp::new())
    }
}

// Additional types for nested divergence
#[derive(Copy, Clone, Debug, Default)]
pub struct EvenLow;
impl ActiveSet for EvenLow {
    const MASK: u32 = 0x00005555; // Even ∩ LowHalf
    const NAME: &'static str = "EvenLow";
}

#[derive(Copy, Clone, Debug, Default)]
pub struct EvenHigh;
impl ActiveSet for EvenHigh {
    const MASK: u32 = 0x55550000; // Even ∩ HighHalf
    const NAME: &'static str = "EvenHigh";
}

impl ComplementOf<EvenHigh> for EvenLow {}
impl ComplementOf<EvenLow> for EvenHigh {}

// ============================================================================
// MERGE OPERATION (Reconvergence - THE KEY SAFETY CHECK)
// ============================================================================

/// Merge two warps back together.
///
/// **COMPILE-TIME VERIFIED**: Only works if S1 and S2 are complements!
///
/// This is the core safety property of session-typed divergence.
/// You cannot merge overlapping sets. You cannot forget lanes.
pub fn merge<S1, S2>(_left: Warp<S1>, _right: Warp<S2>) -> Warp<All>
where
    S1: ComplementOf<S2>,
    S2: ActiveSet,
{
    Warp::new()
}

/// Merge for nested divergence (returns to parent active set, not All)
pub fn merge_to_even(_left: Warp<EvenLow>, _right: Warp<EvenHigh>) -> Warp<Even> {
    Warp::new()
}

// ============================================================================
// SHUFFLE OPERATION (Requires All Lanes Active)
// ============================================================================

/// Per-lane value (placeholder for actual data)
#[derive(Copy, Clone, Debug)]
pub struct PerLane<T>(pub T);

impl Warp<All> {
    /// Shuffle: exchange data between lanes.
    ///
    /// **ONLY AVAILABLE ON Warp<All>**
    ///
    /// This is enforced by implementing shuffle ONLY for Warp<All>.
    /// Calling shuffle on a diverged warp (Warp<Even>) is a compile error.
    pub fn shuffle_xor<T: Copy>(&self, data: PerLane<T>, _mask: u32) -> PerLane<T> {
        // In real implementation, this would do the shuffle
        data
    }

    /// Shuffle down: lane[i] reads from lane[i+delta].
    ///
    /// **ONLY AVAILABLE ON Warp<All>**
    ///
    /// If the source lane is out of range, result is undefined in CUDA.
    /// In our type system, shuffle_down is simply absent on diverged warps.
    pub fn shuffle_down<T: Copy>(&self, data: PerLane<T>, _delta: u32) -> PerLane<T> {
        // In real implementation: __shfl_down_sync(All::MASK, data, delta)
        data
    }

    /// Reduce: combine all lanes into lane 0.
    ///
    /// **ONLY AVAILABLE ON Warp<All>**
    pub fn reduce_sum<T: Copy + std::ops::Add<Output = T>>(&self, data: PerLane<T>) -> T {
        // Placeholder - real impl would reduce
        data.0
    }

    /// Broadcast: all lanes get the same value.
    ///
    /// Available on All because all lanes need to receive.
    pub fn broadcast<T: Copy>(&self, value: T) -> PerLane<T> {
        PerLane(value)
    }
}

// NOTE: Warp<Even>, Warp<Odd>, etc. do NOT have shuffle/reduce methods.
// Attempting to call them is a compile error:
//
//   let (evens, odds) = all.diverge_even_odd();
//   evens.shuffle_xor(data, 1);  // ERROR: no method named `shuffle_xor` found
//
// This is exactly what we want! Shuffle on a diverged warp is undefined behavior
// in CUDA. Our type system prevents it at compile time.

// ============================================================================
// SYNC OPERATION (Requires matching active sets)
// ============================================================================

impl<S: ActiveSet> Warp<S> {
    /// Sync: barrier for all active lanes in this warp.
    ///
    /// Safe because we only sync the lanes that are actually active.
    /// The active set S tells us exactly which lanes participate.
    pub fn sync(&self) {
        // In real implementation: __syncwarp(S::MASK)
        // Sync with mask S::MASK — type guarantees only active lanes participate
    }
}

// ============================================================================
// TESTS: Demonstrating Compile-Time Verification
// ============================================================================

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

    #[test]
    fn test_diverge_merge_even_odd() {
        let all: Warp<All> = Warp::new();
        let (evens, odds) = all.diverge_even_odd();

        // Type system knows: evens is Warp<Even>, odds is Warp<Odd>
        assert_eq!(evens.active_set_name(), "Even");
        assert_eq!(odds.active_set_name(), "Odd");

        // Merge works because Even: ComplementOf<Odd>
        let merged: Warp<All> = merge(evens, odds);
        assert_eq!(merged.active_set_name(), "All");
    }

    #[test]
    fn test_diverge_merge_halves() {
        let all: Warp<All> = Warp::new();
        let (low, high) = all.diverge_halves();

        // Merge works because LowHalf: ComplementOf<HighHalf>
        let merged = merge(low, high);
        assert_eq!(merged.active_mask(), 0xFFFFFFFF);
    }

    #[test]
    fn test_nested_diverge() {
        let all: Warp<All> = Warp::new();
        let (evens, odds) = all.diverge_even_odd();

        // Further split evens
        let (even_low, even_high) = evens.diverge_halves();
        assert_eq!(even_low.active_mask(), 0x00005555);
        assert_eq!(even_high.active_mask(), 0x55550000);

        // Merge even_low and even_high back to evens
        let evens_restored: Warp<Even> = merge_to_even(even_low, even_high);
        assert_eq!(evens_restored.active_set_name(), "Even");

        // Then merge evens and odds back to all
        let all_restored = merge(evens_restored, odds);
        assert_eq!(all_restored.active_set_name(), "All");
    }

    #[test]
    fn test_shuffle_only_on_all() {
        let all: Warp<All> = Warp::new();
        let data = PerLane(42i32);

        // This compiles: shuffle on Warp<All>
        let _shuffled = all.shuffle_xor(data, 1);

        // This would NOT compile (uncomment to see error):
        // let (evens, _) = all.diverge_even_odd();
        // evens.shuffle_xor(data, 1);
        // Error: no method named `shuffle_xor` found for struct `Warp<Even>`
    }

    #[test]
    fn test_reduce_only_on_all() {
        let all: Warp<All> = Warp::new();
        let data = PerLane(1i32);

        // This compiles: reduce on Warp<All>
        let _sum = all.reduce_sum(data);

        // This would NOT compile (uncomment to see error):
        // let (evens, _) = all.diverge_even_odd();
        // evens.reduce_sum(data);
        // Error: no method named `reduce_sum` found for struct `Warp<Even>`
    }

    #[test]
    fn test_sync_on_any_active_set() {
        let all: Warp<All> = Warp::new();
        let (evens, odds) = all.diverge_even_odd();

        // Sync works on any active set - it just syncs those lanes
        all.sync(); // Syncs all 32 lanes
        evens.sync(); // Syncs only even lanes
        odds.sync(); // Syncs only odd lanes
    }

    // ========================================================================
    // COMPILE-TIME ERROR EXAMPLES (Uncomment to see type errors)
    // ========================================================================

    // fn test_bad_merge_non_complement() {
    //     let all: Warp<All> = Warp::new();
    //     let (evens, _odds) = all.diverge_even_odd();
    //     let (low, _high) = all.diverge_halves();
    //
    //     // ERROR: Even is not ComplementOf<LowHalf>
    //     let _bad = merge(evens, low);
    //     //         ^^^^^ the trait `ComplementOf<LowHalf>` is not implemented for `Even`
    // }

    // fn test_bad_merge_same_twice() {
    //     let all: Warp<All> = Warp::new();
    //     let (evens, _odds) = all.diverge_even_odd();
    //
    //     // ERROR: Even is not ComplementOf<Even>
    //     let _bad = merge(evens, evens);
    //     //         ^^^^^ the trait `ComplementOf<Even>` is not implemented for `Even`
    // }

    // fn test_bad_shuffle_on_diverged() {
    //     let all: Warp<All> = Warp::new();
    //     let (evens, _odds) = all.diverge_even_odd();
    //     let data = PerLane(42i32);
    //
    //     // ERROR: shuffle_xor not available on Warp<Even>
    //     evens.shuffle_xor(data, 1);
    //     //    ^^^^^^^^^^ method not found in `Warp<Even>`
    // }
}

// ============================================================================
// CONCLUSIONS
// ============================================================================
//
// Q5 Answer: **YES, we can verify lane masks statically for predefined patterns.**
//
// What's verified at compile time:
// ✓ merge() only accepts complementary sets (Even+Odd, LowHalf+HighHalf)
// ✓ shuffle/reduce only available when all lanes active (Warp<All>)
// ✓ Active set tracked through diverge/merge chains
// ✓ Nested divergence works (can split Even into EvenLow+EvenHigh)
//
// What's NOT verified (would need dependent types):
// ✗ Arbitrary user-defined predicates (e.g., lanes where data[i] > 0)
// ✗ Runtime-computed masks
// ✗ Data-dependent divergence
//
// The key insight: For the COMMON patterns (even/odd, halves, single lane),
// marker types give us full compile-time verification. This covers:
// - Butterfly reductions (even/odd XOR patterns)
// - Prefix scans (halves)
// - Reduction results (lane 0 extraction)
// - Blocked algorithms (tiles)
//
// For arbitrary predicates, we'd need:
// 1. Dependent types (not Rust)
// 2. Runtime checking (loses compile-time guarantee)
// 3. Restrict to predefined patterns (what we did here)
//
// RECOMMENDATION: The marker type approach is sufficient for a research
// prototype and covers the most important use cases. A production
// implementation could extend with proc-macros for more patterns.