json-extractor 0.1.0

//! Stage 1: Character Classification and Structural Index Building
//!
//! This module implements the first stage of the two-stage JSON parsing pipeline.
//! It performs bulk character classification using SIMD to identify all structural
//! characters in the document, building an index of their positions.
//!
//! # Architecture
//!
//! Uses the "shufti" technique from simd-json:
//! - Process 64 bytes at once (two 32-byte AVX2 registers)
//! - Nibble-based lookup tables for character classification
//! - Bitmask extraction to find structural character positions
//! - Returns Vec<u32> of byte positions for Stage 2
//!
//! # Structural Characters
//!
//! Stage 1 identifies: `{`, `}`, `[`, `]`, `:`, `,`, `"`
//!
//! # Performance
//!
//! Expected throughput: 5-10 GiB/s on character classification

#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::*;

use smallvec::SmallVec;

/// Check if a quote at the given position is escaped by preceding backslashes
///
/// A quote is escaped if it's preceded by an odd number of backslashes.
/// Examples:
/// - `\"` → escaped (1 backslash)
/// - `\\"` → not escaped (2 backslashes, the quote is real)
/// - `\\\"` → escaped (3 backslashes)
#[inline]
fn is_escaped_quote(data: &[u8], quote_pos: usize) -> bool {
    // Count consecutive backslashes before the quote
    let mut backslash_count = 0;
    let mut pos = quote_pos;
    while pos > 0 && data[pos - 1] == b'\\' {
        backslash_count += 1;
        pos -= 1;
    }
    // Quote is escaped if odd number of backslashes precede it
    backslash_count % 2 == 1
}

/// Result of Stage 1 processing
pub struct Stage1Output {
    /// Positions of structural characters in the document
    pub structural_indices: Vec<u32>,
    /// Matched bracket pairs: (opening_position, closing_position)
    /// Enables O(1) bracket matching in Stage 2
    pub bracket_pairs: Vec<(u32, u32)>,
}

impl Stage1Output {
    /// Create a new empty Stage1Output with no pre-allocated capacity
    pub fn new() -> Self {
        Self {
            structural_indices: Vec::new(),
            bracket_pairs: Vec::new(),
        }
    }

    /// Clear all buffers while preserving capacity
    fn clear(&mut self) {
        self.structural_indices.clear();
        self.bracket_pairs.clear();
    }
}

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

/// Process a complete JSON document and identify all structural character positions (reusable buffer)
///
/// This version reuses the provided Stage1Output buffer, clearing it first to preserve capacity.
///
/// # Arguments
/// * `data` - Complete JSON document bytes
/// * `output` - Mutable reference to Stage1Output buffer to reuse
#[inline]
pub fn find_structural_indices(data: &[u8], output: &mut Stage1Output) {
    // Clear buffers while preserving capacity
    output.clear();

    // Fast path for small inputs - SIMD overhead not worth it
    const SMALL_INPUT_THRESHOLD: usize = 64;
    if data.len() < SMALL_INPUT_THRESHOLD {
        find_structural_indices_scalar(data, output);
        return;
    }

    #[cfg(target_arch = "x86_64")]
    {
        if is_x86_feature_detected!("avx2") {
            // Safety: We just checked that AVX2 is available
            unsafe { find_structural_indices_avx2(data, output) }
        } else if is_x86_feature_detected!("sse4.2") {
            // Safety: We just checked that SSE4.2 is available
            unsafe { find_structural_indices_sse42(data, output) }
        } else {
            find_structural_indices_scalar(data, output)
        }
    }

    #[cfg(not(target_arch = "x86_64"))]
    {
        find_structural_indices_scalar(data, output)
    }
}

/// AVX2 implementation with reusable buffer: Process 64 bytes at once (2x 32-byte registers)
///
/// # Safety
/// Caller must ensure AVX2 is supported on the CPU
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn find_structural_indices_avx2(data: &[u8], output: &mut Stage1Output) {
    const CHUNK_SIZE: usize = 64;

    // Note: No need to pre-allocate capacity - clear() preserves existing capacity
    // and Vec will grow as needed. Pre-allocation can cause unbounded growth.

    let mut pos = 0;

    // Process 64-byte chunks with SIMD
    while pos + CHUNK_SIZE <= data.len() {
        unsafe {
            // Load two 32-byte chunks
            let ptr = data.as_ptr().add(pos);

            // Prefetch next chunk (3-5 cache lines ahead)
            if pos + CHUNK_SIZE * 2 <= data.len() {
                _mm_prefetch(
                    data.as_ptr().add(pos + CHUNK_SIZE * 2) as *const i8,
                    _MM_HINT_T0,
                );
            }

            let v0 = _mm256_loadu_si256(ptr as *const __m256i);
            let v1 = _mm256_loadu_si256(ptr.add(32) as *const __m256i);

            // Apply shufti classification to find structural characters
            let structural_mask0 = classify_structural_avx2(v0);
            let structural_mask1 = classify_structural_avx2(v1);

            // Convert to bitmasks (one bit per byte)
            let bits0 = _mm256_movemask_epi8(structural_mask0) as u32;
            let bits1 = _mm256_movemask_epi8(structural_mask1) as u32;

            // Combine into 64-bit mask
            let combined_bits = (bits1 as u64) << 32 | bits0 as u64;

            // Extract positions of set bits
            extract_indices(combined_bits, pos as u32, &mut output.structural_indices);

            pos += CHUNK_SIZE;
        }
    }

    // After main AVX2 loop, process remaining 32+ bytes with SSE
    if pos + 32 <= data.len() {
        // Use SSE4.2 for one more chunk
        unsafe {
            let ptr = data.as_ptr().add(pos);
            let v0 = _mm_loadu_si128(ptr as *const __m128i);
            let v1 = _mm_loadu_si128(ptr.add(16) as *const __m128i);

            let mask0 = classify_structural_sse42(v0);
            let mask1 = classify_structural_sse42(v1);

            let bits =
                ((_mm_movemask_epi8(mask1) as u32) << 16) | (_mm_movemask_epi8(mask0) as u32);

            extract_indices(bits as u64, pos as u32, &mut output.structural_indices);
            pos += 32;
        }
    }

    // Handle remaining bytes with scalar code
    while pos < data.len() {
        if crate::charclass::is_structural(data[pos]) {
            output.structural_indices.push(pos as u32);
        }
        pos += 1;
    }

    // Pair brackets, skipping those inside strings
    pair_brackets(data, &output.structural_indices, &mut output.bracket_pairs);
}

/// Classify characters as structural using shufti technique
///
/// # Safety
/// Requires AVX2 support
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn classify_structural_avx2(chunk: __m256i) -> __m256i {
    // Shufti technique: Use low and high nibbles as lookup indices
    //
    // Structural characters we need to find: { } [ ] : , "
    // ASCII codes: 0x7B, 0x7D, 0x5B, 0x5D, 0x3A, 0x2C, 0x22
    //
    // We split each byte into low and high nibbles and use them to index
    // into lookup tables. A character is structural if both lookups produce
    // non-zero results after AND operation.

    // Shufti lookup tables with unique bit patterns to avoid false positives
    // Each structural char needs a unique bit combo so only exact matches have MSB set after AND
    //
    // Bit assignments:
    // - Bit 7 (0x80/-128): '"' (0x22) → low=2, high=2
    // - Bit 6 (0x40/-64):  ':' (0x3A) → low=A, high=3
    // - Bit 5 (0x20/-32):  ',' (0x2C) → low=C, high=2
    // - Bit 4 (0x10/-16):  '[' (0x5B) and ']' (0x5D) → low=B/D, high=5
    // - Bit 3 (0x08/-8):   '{' (0x7B) and '}' (0x7D) → low=B/D, high=7

    // Low nibble table (index = low 4 bits)
    // Index: 0  1  2 ("|,)              3  4  5  6  7  8  9  A (:)  B ([{)       C (,)   D (]{)       E  F
    let low_nibble_mask = _mm256_setr_epi8(
        0, 0, -96, 0, 0, 0, 0, 0, 0, 0, 64, 24, 32, 24, 0, 0, // First 16 bytes
        //      -96 = 0xA0 = bits 7,5 (quote, comma)
        //                                      64 = 0x40 = bit 6 (colon)
        //                                          24 = 0x18 = bits 4,3 (brackets/braces)
        //                                              32 = 0x20 = bit 5 (comma)
        //                                                  24 = 0x18 = bits 4,3 (brackets/braces)
        0, 0, -96, 0, 0, 0, 0, 0, 0, 0, 64, 24, 32, 24, 0, 0, // Second 16 bytes (repeat)
    );

    // High nibble table (index = high 4 bits)
    // Index: 0  1  2 ("|,)              3 (:)  4  5 ([])       6  7 ({})        8  9  A  B  C  D  E  F
    let high_nibble_mask = _mm256_setr_epi8(
        0, 0, -96, 64, 0, 16, 0, 8, 0, 0, 0, 0, 0, 0, 0, 0, // First 16 bytes
        //      -96 = 0xA0 = bits 7,5 (quote, comma)
        //             64 = 0x40 = bit 6 (colon)
        //                      16 = 0x10 = bit 4 (square brackets)
        //                              8 = 0x08 = bit 3 (curly braces)
        0, 0, -96, 64, 0, 16, 0, 8, 0, 0, 0, 0, 0, 0, 0, 0, // Second 16 bytes (repeat)
    );

    // Extract low and high nibbles
    let low_mask = _mm256_set1_epi8(0x0F);
    let low_nibbles = _mm256_and_si256(chunk, low_mask);
    let high_nibbles = _mm256_and_si256(_mm256_srli_epi32(chunk, 4), low_mask);

    // Shuffle lookup tables using nibbles as indices
    let low_result = _mm256_shuffle_epi8(low_nibble_mask, low_nibbles);
    let high_result = _mm256_shuffle_epi8(high_nibble_mask, high_nibbles);

    // Character is structural if BOTH nibbles match (AND has non-zero bits)
    let result = _mm256_and_si256(low_result, high_result);

    // movemask only checks bit 7, but our result may have bits 3-7 set
    // Check if byte != 0 by comparing == 0 and then inverting
    let zeros = _mm256_setzero_si256();
    let eq_zero = _mm256_cmpeq_epi8(result, zeros); // 0xFF where zero, 0x00 where non-zero
    _mm256_xor_si256(eq_zero, _mm256_set1_epi8(-1)) // Invert: 0x00 where zero, 0xFF where non-zero
}

/// SSE4.2 implementation with reusable buffer: Process 32 bytes at once (2x 16-byte registers)
///
/// # Safety
/// Caller must ensure SSE4.2 is supported on the CPU
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "sse4.2")]
unsafe fn find_structural_indices_sse42(data: &[u8], output: &mut Stage1Output) {
    const CHUNK_SIZE: usize = 32;

    // Note: No need to pre-allocate capacity - clear() preserves existing capacity
    // and Vec will grow as needed. Pre-allocation can cause unbounded growth.

    let mut pos = 0;

    // Process 32-byte chunks with SIMD
    while pos + CHUNK_SIZE <= data.len() {
        unsafe {
            // Load two 16-byte chunks
            let ptr = data.as_ptr().add(pos);

            // Prefetch next chunk (2 chunks ahead)
            if pos + CHUNK_SIZE * 2 <= data.len() {
                _mm_prefetch(
                    data.as_ptr().add(pos + CHUNK_SIZE * 2) as *const i8,
                    _MM_HINT_T0,
                );
            }

            let v0 = _mm_loadu_si128(ptr as *const __m128i);
            let v1 = _mm_loadu_si128(ptr.add(16) as *const __m128i);

            // Apply shufti classification
            let structural_mask0 = classify_structural_sse42(v0);
            let structural_mask1 = classify_structural_sse42(v1);

            // Convert to bitmasks
            let bits0 = _mm_movemask_epi8(structural_mask0) as u16;
            let bits1 = _mm_movemask_epi8(structural_mask1) as u16;

            // Combine into 32-bit mask
            let combined_bits = (bits1 as u32) << 16 | bits0 as u32;

            // Extract positions
            extract_indices(
                combined_bits as u64,
                pos as u32,
                &mut output.structural_indices,
            );

            pos += CHUNK_SIZE;
        }
    }

    // After main SSE loop, process remaining 16+ bytes with a single SSE chunk
    if pos + 16 <= data.len() {
        unsafe {
            let ptr = data.as_ptr().add(pos);
            let v0 = _mm_loadu_si128(ptr as *const __m128i);

            let mask0 = classify_structural_sse42(v0);
            let bits = _mm_movemask_epi8(mask0) as u16;

            extract_indices(bits as u64, pos as u32, &mut output.structural_indices);
            pos += 16;
        }
    }

    // Handle remaining bytes with scalar code
    while pos < data.len() {
        if crate::charclass::is_structural(data[pos]) {
            output.structural_indices.push(pos as u32);
        }
        pos += 1;
    }

    // Pair brackets, skipping those inside strings
    pair_brackets(data, &output.structural_indices, &mut output.bracket_pairs);
}

/// Classify characters as structural using shufti technique (SSE4.2 version)
///
/// # Safety
/// Requires SSE4.2 support
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "sse4.2")]
unsafe fn classify_structural_sse42(chunk: __m128i) -> __m128i {
    // Same bit-pattern approach as AVX2 to avoid false positives
    // Index: 0  1  2 ("|,)  3  4  5  6  7  8  9  A (:)  B ([{)  C (,)   D (]{)  E  F
    let low_nibble_mask = _mm_setr_epi8(0, 0, -96, 0, 0, 0, 0, 0, 0, 0, 64, 24, 32, 24, 0, 0);

    let high_nibble_mask = _mm_setr_epi8(0, 0, -96, 64, 0, 16, 0, 8, 0, 0, 0, 0, 0, 0, 0, 0);

    // Extract nibbles
    let low_mask = _mm_set1_epi8(0x0F);
    let low_nibbles = _mm_and_si128(chunk, low_mask);
    let high_nibbles = _mm_and_si128(_mm_srli_epi32(chunk, 4), low_mask);

    // Shuffle and combine
    let low_result = _mm_shuffle_epi8(low_nibble_mask, low_nibbles);
    let high_result = _mm_shuffle_epi8(high_nibble_mask, high_nibbles);

    let result = _mm_and_si128(low_result, high_result);

    // Check if byte != 0 by comparing == 0 and then inverting
    let zeros = _mm_setzero_si128();
    let eq_zero = _mm_cmpeq_epi8(result, zeros);
    _mm_xor_si128(eq_zero, _mm_set1_epi8(-1))
}

/// Extract bit positions from a bitmask and add them as indices
///
/// Processes up to 64 bits, extracting positions where bits are set.
#[inline]
fn extract_indices(mut bitmask: u64, base_pos: u32, indices: &mut Vec<u32>) {
    while bitmask != 0 {
        let offset = bitmask.trailing_zeros();
        indices.push(base_pos + offset);
        bitmask &= bitmask - 1; // Clear lowest bit
    }
}

/// Scalar implementation with reusable buffer: Byte-by-byte scanning
///
/// Uses character classification lookup table for fast classification
fn find_structural_indices_scalar(data: &[u8], output: &mut Stage1Output) {
    // Note: No need to pre-allocate capacity - clear() preserves existing capacity
    // and Vec will grow as needed. Pre-allocation can cause unbounded growth.

    for (pos, byte) in data.iter().enumerate() {
        if crate::charclass::is_structural(*byte) {
            output.structural_indices.push(pos as u32);
        }
    }

    // Pair brackets, skipping those inside strings
    pair_brackets(data, &output.structural_indices, &mut output.bracket_pairs);
}

/// Pair up opening and closing brackets (reusable buffer), skipping brackets inside strings
///
/// Uses a stack-based approach to match brackets, handling both {} and [] pairs.
/// Skips brackets that are inside string ranges to avoid false matches.
fn pair_brackets(data: &[u8], indices: &[u32], output: &mut Vec<(u32, u32)>) {
    let mut brace_stack: SmallVec<[u32; 16]> = SmallVec::new(); // Stack for { positions (8 on stack, heap fallback)
    let mut bracket_stack: SmallVec<[u32; 16]> = SmallVec::new(); // Stack for [ positions (8 on stack, heap fallback)

    // Track string state inline as we iterate (O(1) per index instead of O(log m) binary search)
    let mut in_string = false;

    for &idx in indices {
        let byte = data[idx as usize];

        // Update string state when we encounter quotes
        if byte == b'"' {
            // Check if this quote is escaped
            if !is_escaped_quote(data, idx as usize) {
                in_string = !in_string; // Toggle string state
            }
        }

        // Skip brackets inside strings (simple O(1) boolean check)
        if in_string && (byte == b'{' || byte == b'}' || byte == b'[' || byte == b']') {
            continue;
        }

        match byte {
            b'{' => brace_stack.push(idx),
            b'}' => {
                if let Some(open_pos) = brace_stack.pop() {
                    output.push((open_pos, idx));
                }
            }
            b'[' => bracket_stack.push(idx),
            b']' => {
                if let Some(open_pos) = bracket_stack.pop() {
                    output.push((open_pos, idx));
                }
            }
            _ => {} // Ignore other structural chars
        }
    }

    // Sort pairs by opening position for efficient binary search in Stage 2
    output.sort_unstable_by_key(|(open, _close)| *open);
}

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

    #[test]
    fn test_find_structural_simple_object() {
        let json = br#"{"key":"value"}"#;
        let mut output = Stage1Output::new();
        find_structural_indices(json, &mut output);

        // Expected structural positions: { " " : " " } = positions 0, 1, 5, 6, 7, 13, 14
        assert_eq!(output.structural_indices.len(), 7);
        assert_eq!(output.structural_indices[0], 0); // {
        assert_eq!(output.structural_indices[1], 1); // " (start of "key")
        assert_eq!(output.structural_indices[2], 5); // " (end of "key")
        assert_eq!(output.structural_indices[3], 6); // :
        assert_eq!(output.structural_indices[4], 7); // " (start of "value")
        assert_eq!(output.structural_indices[5], 13); // " (end of "value")
        assert_eq!(output.structural_indices[6], 14); // }
    }

    #[test]
    fn test_find_structural_array() {
        let json = br#"[1,2,3]"#;
        let mut output = Stage1Output::new();
        find_structural_indices(json, &mut output);

        // Expected: [ , , ] = positions 0, 2, 4, 6
        assert_eq!(output.structural_indices.len(), 4);
        assert_eq!(output.structural_indices[0], 0); // [
        assert_eq!(output.structural_indices[1], 2); // ,
        assert_eq!(output.structural_indices[2], 4); // ,
        assert_eq!(output.structural_indices[3], 6); // ]
    }

    #[test]
    fn test_find_structural_nested() {
        let json = br#"{"a":[1,2]}"#;
        let mut output = Stage1Output::new();
        find_structural_indices(json, &mut output);

        // Expected: { " " : [ , ] } = positions 0, 1, 3, 4, 5, 7, 9, 10
        assert_eq!(output.structural_indices.len(), 8);
    }

    #[test]
    fn test_scalar_matches_simd() {
        let json = br#"{"test":[1,2,3],"nested":{"key":"value"}}"#;

        let mut scalar_output = Stage1Output::new();
        find_structural_indices_scalar(json, &mut scalar_output);

        let mut simd_output = Stage1Output::new();
        find_structural_indices(json, &mut simd_output);

        assert_eq!(
            scalar_output.structural_indices,
            simd_output.structural_indices
        );
    }

    #[test]
    fn test_large_document() {
        // Generate a large JSON document
        let mut json = String::from("[");
        for i in 0..1000 {
            if i > 0 {
                json.push(',');
            }
            json.push_str(&format!(r#"{{"id":{}}}"#, i));
        }
        json.push(']');

        let mut output = Stage1Output::new();
        find_structural_indices(json.as_bytes(), &mut output);

        // Should find many structural characters
        assert!(output.structural_indices.len() > 1000);
    }

    // ==== GRANULAR UNIT TESTS FOR SIMD DEBUGGING ====

    #[test]
    fn test_shufti_each_structural_char() {
        // Test that each structural character is detected individually
        let test_cases = vec![
            (b'{', "left brace"),
            (b'}', "right brace"),
            (b'[', "left bracket"),
            (b']', "right bracket"),
            (b':', "colon"),
            (b',', "comma"),
            (b'"', "quote"),
        ];

        for (ch, name) in test_cases {
            let data = vec![ch];
            let mut output = Stage1Output::new();
            find_structural_indices(&data, &mut output);
            assert_eq!(
                output.structural_indices.len(),
                1,
                "{} (0x{:02X}) should be detected as structural",
                name,
                ch
            );
            assert_eq!(
                output.structural_indices[0], 0,
                "{} should be at position 0",
                name
            );
        }
    }

    #[test]
    fn test_extract_indices_known_bitmasks() {
        // Test extract_indices with known bitmask patterns
        let mut indices = Vec::new();

        // Bitmask with bit 0 set (position 0)
        extract_indices(0b1, 0, &mut indices);
        assert_eq!(indices, vec![0]);

        // Bitmask with bit 5 set (position 5)
        indices.clear();
        extract_indices(0b100000, 10, &mut indices);
        assert_eq!(indices, vec![15]); // base 10 + offset 5

        // Bitmask with multiple bits: positions 0, 3, 7
        indices.clear();
        extract_indices(0b10001001, 0, &mut indices);
        assert_eq!(indices, vec![0, 3, 7]);

        // Bitmask with bits in high word (32-bit offset)
        indices.clear();
        extract_indices(0b1_00000000_00000000_00000000_00000000, 0, &mut indices);
        assert_eq!(indices, vec![32]);
    }

    #[test]
    fn test_single_chunk_exactly_64_bytes() {
        // Exactly 64 bytes of complete JSON objects
        let json = br#"{"k0":0}{"k1":1}{"k2":2}{"k3":3}{"k4":4}{"k5":5}{"k6":6}{"k7":7}"#;
        assert_eq!(json.len(), 64, "Test JSON must be exactly 64 bytes");

        let mut output = Stage1Output::new();
        find_structural_indices(json, &mut output);

        // Count expected structural chars
        let expected_count = json
            .iter()
            .filter(|&&b| {
                b == b'{'
                    || b == b'}'
                    || b == b'"'
                    || b == b':'
                    || b == b','
                    || b == b'['
                    || b == b']'
            })
            .count();

        println!(
            "Single chunk (64 bytes): Found {} structural chars, expected {}",
            output.structural_indices.len(),
            expected_count
        );

        assert_eq!(
            output.structural_indices.len(),
            expected_count,
            "Single 64-byte chunk should find all structural characters"
        );
    }

    #[test]
    fn test_two_chunks_exactly_128_bytes() {
        // Exactly 128 bytes: 124 bytes JSON + 4 spaces padding
        let json = br#"{"k0":0}{"k1":1}{"k2":2}{"k3":3}{"k4":4}{"k5":5}{"k6":6}{"k7":7}{"k100":100}{"k101":101}{"k102":102}{"k103":103}{"k104":104}    "#;
        assert_eq!(json.len(), 128, "Test JSON must be exactly 128 bytes");

        let mut output = Stage1Output::new();
        find_structural_indices(json, &mut output);

        let expected_count = json
            .iter()
            .filter(|&&b| {
                b == b'{'
                    || b == b'}'
                    || b == b'"'
                    || b == b':'
                    || b == b','
                    || b == b'['
                    || b == b']'
            })
            .count();

        println!(
            "Two chunks (128 bytes): Found {} structural chars, expected {}",
            output.structural_indices.len(),
            expected_count
        );

        assert_eq!(
            output.structural_indices.len(),
            expected_count,
            "Two 64-byte chunks should find all structural characters"
        );
    }

    #[test]
    fn test_progressive_sizes() {
        // Test progressively larger sizes to isolate where the bug appears
        let sizes = vec![7, 32, 64, 96, 128, 160, 192, 256];

        for size in sizes {
            // Generate JSON of approximately this size
            let mut json = String::from("[");
            let mut current_len = 1;
            let mut item_num = 0;

            while current_len < size - 10 {
                if item_num > 0 {
                    json.push(',');
                    current_len += 1;
                }
                json.push_str(r#"{"x":1}"#);
                current_len += 7;
                item_num += 1;
            }
            json.push(']');

            let json_bytes = json.as_bytes();
            let mut output = Stage1Output::new();
            find_structural_indices(json_bytes, &mut output);

            let expected_count = json_bytes
                .iter()
                .filter(|&&b| {
                    b == b'{'
                        || b == b'}'
                        || b == b'['
                        || b == b']'
                        || b == b'"'
                        || b == b':'
                        || b == b','
                })
                .count();

            println!(
                "Size {}: Found {} structural chars, expected {}",
                json_bytes.len(),
                output.structural_indices.len(),
                expected_count
            );

            assert_eq!(
                output.structural_indices.len(),
                expected_count,
                "Size {} should find all {} structural characters",
                json_bytes.len(),
                expected_count
            );
        }
    }

    #[test]
    fn test_chunk_boundary_structural_chars() {
        // Place structural characters exactly at chunk boundaries (64, 128, etc.)
        let mut json = vec![b' '; 192]; // 3 chunks worth

        // Place structural chars at strategic positions
        json[0] = b'{'; // Start of chunk 0
        json[32] = b'"'; // Middle of chunk 0
        json[63] = b':'; // End of chunk 0
        json[64] = b'['; // Start of chunk 1
        json[96] = b','; // Middle of chunk 1
        json[127] = b']'; // End of chunk 1
        json[128] = b'{'; // Start of chunk 2
        json[160] = b'"'; // Middle of chunk 2
        json[191] = b'}'; // End of chunk 2

        let mut output = Stage1Output::new();
        find_structural_indices(&json, &mut output);

        let expected_positions = vec![0, 32, 63, 64, 96, 127, 128, 160, 191];

        println!(
            "Chunk boundary test: Found {} indices",
            output.structural_indices.len()
        );
        println!("Expected positions: {:?}", expected_positions);
        println!("Found positions: {:?}", output.structural_indices);

        assert_eq!(
            output.structural_indices.len(),
            expected_positions.len(),
            "Should find exactly {} structural characters at known positions",
            expected_positions.len()
        );

        for (i, &expected_pos) in expected_positions.iter().enumerate() {
            assert_eq!(
                output.structural_indices[i], expected_pos,
                "Structural char {} should be at position {}",
                i, expected_pos
            );
        }
    }
}