synthra 0.1.0

data/ext/synthra_native/src/lib.rs

@@ -0,0 +1,229 @@
+//! Synthra Native Extension
+//!
+//! High-performance data generation using Rust and fake-rs.
+//! Provides 100-250x speedup over pure Ruby implementation.
+//! Uses Rayon for parallel processing to utilize CPU cores.
+//!
+//! # Thread Control
+//! Set `RAYON_NUM_THREADS` environment variable before loading the gem
+//! to limit the number of threads used for parallel generation.
+//! Example: `RAYON_NUM_THREADS=4 ruby my_script.rb`
+
+use magnus::{function, prelude::*, Error, RArray, Ruby, Value};
+use rayon::ThreadPoolBuilder;
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+// Default thread count (2 threads to leave resources for other work)
+const DEFAULT_THREADS: usize = 2;
+
+// Track configured thread limit for reporting
+static THREAD_LIMIT: AtomicUsize = AtomicUsize::new(DEFAULT_THREADS);
+
+mod generator;
+mod types;
+
+/// Generate records and return as Ruby Array of Hashes
+///
+/// # Arguments
+/// * `schema_spec` - Array of [name, type_name, args_json, optional, nullable]
+/// * `count` - Number of records to generate
+/// * `seed` - Optional random seed for reproducibility
+/// * `locale` - Locale string (en, fr_fr, de_de, etc.)
+/// * `threads` - Number of threads (nil = default 2, 0 = use all cores)
+fn generate(
+    ruby: &Ruby,
+    schema_spec: RArray,
+    count: usize,
+    seed: Option<u64>,
+    locale: Option<String>,
+    threads: Option<usize>,
+) -> Result<RArray, Error> {
+    // Configure thread pool
+    let num_threads = match threads {
+        Some(0) => num_cpus(),           // 0 = use all cores
+        Some(t) => t.min(num_cpus()),    // specified value
+        None => DEFAULT_THREADS,          // default = 2
+    };
+    
+    let _ = ThreadPoolBuilder::new()
+        .num_threads(num_threads)
+        .build_global();
+
+    let gen = generator::NativeGenerator::from_ruby(schema_spec, seed, locale)?;
+    gen.generate_to_ruby(ruby, count)
+}
+
+/// Generate records directly to a file (most efficient for large datasets)
+///
+/// # Arguments
+/// * `schema_spec` - Array of [name, type_name, args_json, optional, nullable]
+/// * `count` - Number of records to generate
+/// * `path` - Output file path
+/// * `seed` - Optional random seed for reproducibility
+/// * `locale` - Locale string
+/// * `format` - Output format: "jsonl", "json", "csv"
+/// * `threads` - Number of threads (nil = default 2, 0 = use all cores)
+fn generate_to_file(
+    schema_spec: RArray,
+    count: usize,
+    path: String,
+    seed: Option<u64>,
+    locale: Option<String>,
+    format: Option<String>,
+    threads: Option<usize>,
+) -> Result<usize, Error> {
+    // Configure thread pool
+    let num_threads = match threads {
+        Some(0) => num_cpus(),           // 0 = use all cores
+        Some(t) => t.min(num_cpus()),    // specified value
+        None => DEFAULT_THREADS,          // default = 2
+    };
+    
+    let _ = ThreadPoolBuilder::new()
+        .num_threads(num_threads)
+        .build_global();
+
+    let gen = generator::NativeGenerator::from_ruby(schema_spec, seed, locale)?;
+    let fmt = format.as_deref().unwrap_or("jsonl");
+    gen.generate_to_file(count, &path, fmt)
+        .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))
+}
+
+/// Generate a single record (for testing/preview)
+fn generate_one(
+    ruby: &Ruby,
+    schema_spec: RArray,
+    seed: Option<u64>,
+    locale: Option<String>,
+) -> Result<Value, Error> {
+    let gen = generator::NativeGenerator::from_ruby(schema_spec, seed, locale)?;
+    gen.generate_one_to_ruby(ruby)
+}
+
+/// Get list of available type names
+fn available_types() -> Vec<String> {
+    types::AVAILABLE_TYPES.iter().map(|s| s.to_string()).collect()
+}
+
+/// Get list of available locale codes
+fn available_locales() -> Vec<String> {
+    vec![
+        "en".to_string(),
+        "fr_fr".to_string(),
+        "de_de".to_string(),
+        "zh_cn".to_string(),
+        "zh_tw".to_string(),
+        "ja_jp".to_string(),
+        "pt_br".to_string(),
+        "it_it".to_string(),
+        "nl_nl".to_string(),
+        "ar_sa".to_string(),
+    ]
+}
+
+/// Get native extension version
+fn version() -> String {
+    env!("CARGO_PKG_VERSION").to_string()
+}
+
+/// Check if extension is loaded and functional
+fn loaded() -> bool {
+    true
+}
+
+/// Get performance info
+fn performance_info() -> String {
+    let cores = num_cpus();
+    let rayon_threads = rayon::current_num_threads();
+    format!(
+        "Native engine powered by fake-rs (Rust). Expected: 1-5M records/sec. \
+         Cores: {}, Threads: {} (default: {}, use 0 for all cores).",
+        cores,
+        rayon_threads,
+        DEFAULT_THREADS
+    )
+}
+
+fn num_cpus() -> usize {
+    std::thread::available_parallelism()
+        .map(|p| p.get())
+        .unwrap_or(1)
+}
+
+/// Get thread count info
+fn thread_count() -> usize {
+    rayon::current_num_threads()
+}
+
+/// Set the number of threads to use for parallel generation
+/// Pass 0 to use all available cores (default)
+/// Returns the actual number of threads that will be used
+///
+/// NOTE: This must be called before any generation operations.
+/// Alternatively, set RAYON_NUM_THREADS environment variable before loading the gem.
+///
+/// Example usage:
+///   # Option 1: Environment variable (recommended)
+///   # RAYON_NUM_THREADS=4 ruby my_script.rb
+///
+///   # Option 2: Call set_threads early in your code
+///   Synthra::Native.set_threads(4)
+fn set_threads(num_threads: usize) -> Result<usize, Error> {
+    let threads = if num_threads == 0 {
+        num_cpus()
+    } else {
+        num_threads.min(num_cpus())
+    };
+
+    THREAD_LIMIT.store(threads, Ordering::SeqCst);
+
+    // Try to configure Rayon's global thread pool
+    // This only works if called before any parallel operations
+    match ThreadPoolBuilder::new()
+        .num_threads(threads)
+        .build_global()
+    {
+        Ok(_) => Ok(threads),
+        Err(_) => {
+            // Pool already initialized - return current thread count
+            Ok(rayon::current_num_threads())
+        }
+    }
+}
+
+/// Get the configured thread limit
+fn get_thread_limit() -> usize {
+    let limit = THREAD_LIMIT.load(Ordering::SeqCst);
+    if limit == 0 {
+        rayon::current_num_threads()
+    } else {
+        limit
+    }
+}
+
+/// Initialize the native extension module
+#[magnus::init]
+fn init(ruby: &Ruby) -> Result<(), Error> {
+    let module = ruby.define_module("Synthra")?;
+    let native = module.define_module("Native")?;
+
+    // Core generation methods
+    native.define_singleton_method("generate", function!(generate, 5))?;
+    native.define_singleton_method("generate_to_file", function!(generate_to_file, 7))?;
+    native.define_singleton_method("generate_one", function!(generate_one, 3))?;
+
+    // Introspection methods
+    native.define_singleton_method("available_types", function!(available_types, 0))?;
+    native.define_singleton_method("available_locales", function!(available_locales, 0))?;
+    native.define_singleton_method("version", function!(version, 0))?;
+    native.define_singleton_method("loaded?", function!(loaded, 0))?;
+    native.define_singleton_method("performance_info", function!(performance_info, 0))?;
+    native.define_singleton_method("thread_count", function!(thread_count, 0))?;
+
+    // Thread control methods
+    native.define_singleton_method("set_threads", function!(set_threads, 1))?;
+    native.define_singleton_method("get_thread_limit", function!(get_thread_limit, 0))?;
+
+    Ok(())
+}
+