clean up

Author: tadghh

Cargo.lock

@@ -14,5 +14,7 @@ curl "http://127.0.0.1:8080/generate?size=100mb&format=csv"
 ```
 
 - The `pretty` parameter enables pretty-printed output (optional).
+  - default: `false`
 - The `size` parameter specifies the target size of the generated content.
 - The `format` parameter supports either JSON or CSV.
+  - default: `json`

readme.md

@@ -1,15 +1,12 @@
 #![feature(get_mut_unchecked, portable_simd)]
 use actix_web::{web, App, Error, HttpResponse, HttpServer};
-
 use parking_lot::Mutex;
 use processing::*;
 use rand::{Rng, SeedableRng};
 use rand_chacha::ChaCha8Rng;
 use rayon::prelude::*;
-
 use std::collections::HashMap;
 use std::io::{stdout, Write};
-
 use std::sync::Arc;
 mod processing;
 
@@ -35,34 +32,30 @@ fn parse_size(size_str: &str) -> Result<usize, String> {
     Ok(number * multiplier)
 }
 
+const BYTE_SIZE: usize = 1024 * 1024;
+
 async fn generate_data(
     web::Query(params): web::Query<HashMap<String, String>>,
     data_pools: web::Data<Arc<DataPools>>,
 ) -> Result<HttpResponse, Error> {
-    const BYTE_SIZE: usize = 1024 * 1024;
-    let target_size = if let Some(size_str) = params.get("size") {
-        parse_size(size_str).unwrap_or(BYTE_SIZE)
-    } else {
-        BYTE_SIZE
+    let target_size = match params.get("size") {
+        Some(size) => match parse_size(size) {
+            Ok(size) => size,
+            Err(_) => BYTE_SIZE,
+        },
+        None => BYTE_SIZE,
     };
     let pretty = params.get("pretty").map_or(false, |v| v == "true");
     let format = OutputFormat::from_str(params.get("format").map_or("json", |s| s));
     let seed: u64 = rand::thread_rng().gen();
     let num_threads = num_cpus::get();
     let chunk_size = target_size / num_threads;
-
-    // Calculate approximate records per chunk based on average record size
     let avg_record_size = if format == OutputFormat::JSON {
-        if pretty {
-            250
-        } else {
-            180
-        } // Approximate sizes
+        250
     } else {
-        120 // CSV approximate record size
+        120
     };
     let records_per_chunk = chunk_size / avg_record_size;
-
     let progress = Arc::new(Mutex::new(ProgressInfo::new(
         (target_size / BYTE_SIZE) as f64,
     )));
@@ -79,6 +72,7 @@ async fn generate_data(
             "CSV"
         }
     );
+
     if format == OutputFormat::JSON {
         println!(
             "Pretty print: {}",

src/main.rs

@@ -1,13 +1,87 @@
-use fake::faker::address::en::*;
-use fake::faker::company::en::*;
-use fake::Fake;
+use crate::ProgressInfo;
+use fake::{
+    faker::{address::en::*, company::en::*},
+    Fake,
+};
 use parking_lot::Mutex;
 use rand::Rng;
 use rand_chacha::ChaCha8Rng;
 use serde::Serialize;
 use std::{simd::u8x32, sync::Arc};
 
-use crate::ProgressInfo;
+const BYTE_COUNT: usize = 32;
+const REFRESH_COUNT: u32 = 1500;
+const POOL_SIZE: i32 = 1000;
+
+// Pre-computed patterns for both pretty and compact modes
+struct JsonPatterns {
+    separator_pretty: [u8; 32],  // ",\n    "
+    separator_compact: [u8; 32], // ","
+    ending_pretty: [u8; 32],     // "\n  }"
+    ending_compact: [u8; 32],    // "}"
+    quoted_field_patterns: [QuotedFieldPattern; 5],
+    unquoted_field_patterns: [UnquotedFieldPattern; 3],
+}
+
+struct QuotedFieldPattern {
+    prefix: [u8; 32], // "\"field\": \""
+    suffix: [u8; 32], // "\""
+    prefix_len: usize,
+    suffix_len: usize,
+}
+
+// Pre-computed pattern for each unquoted field
+struct UnquotedFieldPattern {
+    prefix: [u8; 32],
+    prefix_len: usize,
+}
+
+impl JsonPatterns {
+    fn new() -> Self {
+        let mut field_start_pretty = [0u8; 32];
+        field_start_pretty[..6].copy_from_slice(b"\n    \"");
+
+        let mut field_start_compact = [0u8; 32];
+        field_start_compact[0] = b'"';
+
+        let mut separator_pretty = [0u8; 32];
+        separator_pretty[..6].copy_from_slice(b",\n    ");
+
+        let mut separator_compact = [0u8; 32];
+        separator_compact[..1].copy_from_slice(b",");
+
+        let mut ending_pretty = [0u8; 32];
+        ending_pretty[..4].copy_from_slice(b"\n  }");
+
+        let mut ending_compact = [0u8; 32];
+        ending_compact[0] = b'}';
+
+        let quoted_fields = [
+            ("name", create_quoted_pattern(b"name")),
+            ("industry", create_quoted_pattern(b"industry")),
+            ("city", create_quoted_pattern(b"city")),
+            ("state", create_quoted_pattern(b"state")),
+            ("country", create_quoted_pattern(b"country")),
+        ]
+        .map(|(_, pattern)| pattern);
+
+        let unquoted_fields = [
+            ("id", create_unquoted_pattern(b"id")),
+            ("revenue", create_unquoted_pattern(b"revenue")),
+            ("employees", create_unquoted_pattern(b"employees")),
+        ]
+        .map(|(_, pattern)| pattern);
+
+        Self {
+            separator_pretty,
+            separator_compact,
+            ending_pretty,
+            ending_compact,
+            quoted_field_patterns: quoted_fields,
+            unquoted_field_patterns: unquoted_fields,
+        }
+    }
+}
 
 #[derive(Serialize)]
 struct BusinessLocation {
@@ -21,7 +95,6 @@ struct BusinessLocation {
     country: String,
 }
 
-// Add an enum to represent output formats
 #[derive(PartialEq)]
 pub enum OutputFormat {
     JSON,
@@ -58,12 +131,11 @@ pub struct ChunkResult {
 
 impl DataPools {
     pub fn new() -> Self {
-        let pool_size = 1000;
         DataPools {
-            names: (0..pool_size).map(|_| CompanyName().fake()).collect(),
-            industries: (0..pool_size).map(|_| Industry().fake()).collect(),
-            cities: (0..pool_size).map(|_| CityName().fake()).collect(),
-            states: (0..pool_size).map(|_| StateName().fake()).collect(),
+            names: (0..POOL_SIZE).map(|_| CompanyName().fake()).collect(),
+            industries: (0..POOL_SIZE).map(|_| Industry().fake()).collect(),
+            cities: (0..POOL_SIZE).map(|_| CityName().fake()).collect(),
+            states: (0..POOL_SIZE).map(|_| StateName().fake()).collect(),
             countries: (0..50).map(|_| CountryName().fake()).collect(),
         }
     }
@@ -122,7 +194,7 @@ pub fn generate_chunk(
             progress_locked.get_mut().update(output.len());
         }
 
-        if current_id % 1500 == 0 {
+        if current_id % REFRESH_COUNT == 0 {
             progress.lock().print_progress();
         }
     }
@@ -174,92 +246,21 @@ fn write_location_csv_simd(location: &BusinessLocation, output: &mut Vec<u8>) {
 fn copy_str_simd(output: &mut Vec<u8>, s: &str) {
     let bytes = s.as_bytes();
     let len = bytes.len();
-    let chunks = len / 32;
+    let chunks = len / BYTE_COUNT;
 
     // Process 32 bytes at a time using SIMD
     for chunk in 0..chunks {
-        let offset = chunk * 32;
-        let simd_chunk = u8x32::from_slice(&bytes[offset..offset + 32]);
+        let offset = chunk * BYTE_COUNT;
+        let simd_chunk = u8x32::from_slice(&bytes[offset..offset + BYTE_COUNT]);
         output.extend_from_slice(&simd_chunk.to_array());
     }
 
-    // Handle remaining bytes
-    let remaining_start = chunks * 32;
+    let remaining_start = chunks * BYTE_COUNT;
     if remaining_start < len {
         output.extend_from_slice(&bytes[remaining_start..]);
     }
 }
 
-// Pre-computed patterns for both pretty and compact modes
-struct JsonPatterns {
-    separator_pretty: [u8; 32],  // ",\n    "
-    separator_compact: [u8; 32], // ","
-    ending_pretty: [u8; 32],     // "\n  }"
-    ending_compact: [u8; 32],    // "}"
-    quoted_field_patterns: [QuotedFieldPattern; 5],
-    unquoted_field_patterns: [UnquotedFieldPattern; 3],
-}
-
-struct QuotedFieldPattern {
-    prefix: [u8; 32], // "\"field\": \""
-    suffix: [u8; 32], // "\""
-    prefix_len: usize,
-    suffix_len: usize,
-}
-
-// Pre-computed pattern for each unquoted field
-struct UnquotedFieldPattern {
-    prefix: [u8; 32],
-    prefix_len: usize,
-}
-
-impl JsonPatterns {
-    fn new() -> Self {
-        let mut field_start_pretty = [0u8; 32];
-        field_start_pretty[..6].copy_from_slice(b"\n    \"");
-
-        let mut field_start_compact = [0u8; 32];
-        field_start_compact[0] = b'"';
-
-        let mut separator_pretty = [0u8; 32];
-        separator_pretty[..6].copy_from_slice(b",\n    ");
-
-        let mut separator_compact = [0u8; 32];
-        separator_compact[..1].copy_from_slice(b",");
-
-        let mut ending_pretty = [0u8; 32];
-        ending_pretty[..4].copy_from_slice(b"\n  }");
-
-        let mut ending_compact = [0u8; 32];
-        ending_compact[0] = b'}';
-
-        let quoted_fields = [
-            ("name", create_quoted_pattern(b"name")),
-            ("industry", create_quoted_pattern(b"industry")),
-            ("city", create_quoted_pattern(b"city")),
-            ("state", create_quoted_pattern(b"state")),
-            ("country", create_quoted_pattern(b"country")),
-        ]
-        .map(|(_, pattern)| pattern);
-
-        let unquoted_fields = [
-            ("id", create_unquoted_pattern(b"id")),
-            ("revenue", create_unquoted_pattern(b"revenue")),
-            ("employees", create_unquoted_pattern(b"employees")),
-        ]
-        .map(|(_, pattern)| pattern);
-
-        Self {
-            separator_pretty,
-            separator_compact,
-            ending_pretty,
-            ending_compact,
-            quoted_field_patterns: quoted_fields,
-            unquoted_field_patterns: unquoted_fields,
-        }
-    }
-}
-
 #[inline]
 fn create_quoted_pattern(field_name: &[u8]) -> QuotedFieldPattern {
     let mut prefix = [0u8; 32];