gpu_scatter_gather/

lib.rs

//! GPU Scatter-Gather Wordlist Generator
//!
//! The world's fastest wordlist generator using GPU acceleration.
//!
//! # Core Innovation
//!
//! Instead of sequential odometer iteration (like maskprocessor), this uses
//! direct index-to-word mapping via mixed-radix arithmetic on the GPU:
//!
//! ```text
//! Index → Mixed-Radix Decomposition → Word
//! ```
//!
//! This enables:
//! - Massively parallel generation (every thread works independently)
//! - O(1) random access to any position in the keyspace
//! - Perfect GPU utilization with coalesced memory access
//! - 500M-1B+ words/second throughput (3-7x faster than maskprocessor)
//! - Multi-GPU support for linear scaling (v1.1.0+)
//!
//! # Example Usage
//!
//! ```rust,no_run
//! use gpu_scatter_gather::{WordlistGenerator, Charset};
//!
//! # fn main() -> anyhow::Result<()> {
//! let mut generator = WordlistGenerator::builder()
//!     .charset(1, Charset::from("abc"))
//!     .charset(2, Charset::from("123"))
//!     .mask(&[1, 2, 1, 2])  // Pattern: ?1?2?1?2
//!     .build()?;
//!
//! // Generate wordlist
//! for word in generator.iter() {
//!     println!("{}", String::from_utf8_lossy(&word));
//! }
//! # Ok(())
//! # }
//! ```

pub mod charset;
pub mod keyspace;
pub mod mask;
pub mod transpose;

// CUDA-dependent modules (only available with "cuda" feature)
#[cfg(feature = "cuda")]
pub mod gpu;
#[cfg(feature = "cuda")]
pub mod multigpu;

pub mod bindings;

#[cfg(feature = "cuda")]
pub mod ffi;

// Re-exports for convenience
pub use charset::Charset;
pub use keyspace::{calculate_keyspace, index_to_word};
pub use mask::Mask;

use anyhow::Result;
use std::collections::HashMap;

/// Main wordlist generator builder
pub struct WordlistGeneratorBuilder {
    charsets: HashMap<usize, Charset>,
    mask: Option<Vec<usize>>,
    batch_size: usize,
}

impl WordlistGeneratorBuilder {
    /// Create a new builder
    pub fn new() -> Self {
        Self {
            charsets: HashMap::new(),
            mask: None,
            batch_size: 10_000_000, // 10M words default batch size
        }
    }

    /// Add a charset with the given ID
    pub fn charset(mut self, id: usize, charset: Charset) -> Self {
        self.charsets.insert(id, charset);
        self
    }

    /// Set the mask pattern (array of charset IDs)
    pub fn mask(mut self, mask: &[usize]) -> Self {
        self.mask = Some(mask.to_vec());
        self
    }

    /// Set the batch size for GPU generation
    pub fn batch_size(mut self, size: usize) -> Self {
        self.batch_size = size;
        self
    }

    /// Build the generator (currently CPU-only, GPU support coming)
    pub fn build(self) -> Result<WordlistGenerator> {
        let mask = self.mask.ok_or_else(|| anyhow::anyhow!("Mask not set"))?;

        // Validate that all charset IDs in mask exist
        for &charset_id in &mask {
            if !self.charsets.contains_key(&charset_id) {
                anyhow::bail!("Mask references undefined charset ID: {charset_id}");
            }
        }

        Ok(WordlistGenerator {
            charsets: self.charsets,
            mask,
            batch_size: self.batch_size,
        })
    }
}

impl Default for WordlistGeneratorBuilder {
    fn default() -> Self {
        Self::new()
    }
}

/// Main wordlist generator
pub struct WordlistGenerator {
    charsets: HashMap<usize, Charset>,
    mask: Vec<usize>,
    #[allow(dead_code)]
    batch_size: usize,
}

impl WordlistGenerator {
    /// Create a new builder
    pub fn builder() -> WordlistGeneratorBuilder {
        WordlistGeneratorBuilder::new()
    }

    /// Calculate the total keyspace size
    pub fn keyspace_size(&self) -> u128 {
        // Calculate product of all charset sizes in the mask
        self.mask
            .iter()
            .map(|&id| self.charsets[&id].len() as u128)
            .product()
    }

    /// Convert index to word using CPU reference implementation
    pub fn index_to_word(&self, index: u64) -> Vec<u8> {
        let mut output = vec![0u8; self.mask.len()];
        let mut remaining = index;

        // Process positions from right to left (least significant to most significant)
        for pos in (0..self.mask.len()).rev() {
            let charset_id = self.mask[pos];
            let charset = &self.charsets[&charset_id];
            let charset_size = charset.len() as u64;

            // Mixed-radix decomposition
            let char_idx = (remaining % charset_size) as usize;
            output[pos] = charset.as_bytes()[char_idx];
            remaining /= charset_size;
        }

        output
    }

    /// Create an iterator over all words (CPU reference implementation)
    pub fn iter(&self) -> WordlistIterator<'_> {
        WordlistIterator {
            generator: self,
            current_index: 0,
            total_keyspace: self.keyspace_size(),
        }
    }
}

/// Iterator over wordlist entries
pub struct WordlistIterator<'a> {
    generator: &'a WordlistGenerator,
    current_index: u64,
    total_keyspace: u128,
}

impl<'a> Iterator for WordlistIterator<'a> {
    type Item = Vec<u8>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.current_index as u128 >= self.total_keyspace {
            return None;
        }

        let word = self.generator.index_to_word(self.current_index);
        self.current_index += 1;
        Some(word)
    }
}

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

    #[test]
    fn test_builder_pattern() {
        let generator = WordlistGenerator::builder()
            .charset(1, Charset::from("abc"))
            .charset(2, Charset::from("123"))
            .mask(&[1, 2])
            .build()
            .unwrap();

        assert_eq!(generator.keyspace_size(), 9); // 3 * 3
    }

    #[test]
    fn test_missing_charset() {
        let result = WordlistGenerator::builder()
            .charset(1, Charset::from("abc"))
            .mask(&[1, 2]) // References charset 2 which doesn't exist
            .build();

        assert!(result.is_err());
    }
}