morok-schedule 0.1.0-alpha.2

//! Tests for partial contiguous buffer optimization.
//!
//! These tests validate the cost-based buffer removal with partial contiguous support.
//! Tests cover:
//! - Configuration levels
//! - Cost heuristics (accessed_buffers, out_in_ratio, buffer_in_reduce)
//! - Transformation correctness (full removal vs partial contiguous)
//! - Edge cases (empty ranges, symbolic sizes, etc.)
//! - Integration with rangeify pipeline

use std::sync::Arc;

use morok_dtype::DType;
use morok_ir::{AddrSpace, AxisId, AxisType, BufferizeOpts, Op, SInt, UOp};
use test_case::test_case;

use crate::rangeify::indexing::IndexingContext;
use crate::rangeify::kernel::PcontigConfig;
use crate::rangeify::patterns::buffer_removal_with_pcontig;
use crate::rewrite::graph_rewrite;

// ============================================================================
// Helper Functions
// ============================================================================

/// Create a test PARAM with given size and dtype.
///
/// Uses Op::Param to match the real pipeline where normalize_buffers_to_params
/// converts Buffer → Param before buffer removal runs. This ensures buffer
/// counting tests reflect actual behavior (Param is counted, Buffer is not).
fn create_test_buffer(size: usize, dtype: DType, id: usize) -> Arc<UOp> {
    let device = UOp::device(morok_device::DeviceSpec::Cpu);
    UOp::new(Op::Param { slot: id, size, device: Some(device) }, dtype.ptr(Some(size), AddrSpace::Global))
}

/// Create a test INDEX(BUFFERIZE(...), ...) pattern.
///
/// This is the core pattern that buffer_removal_with_pcontig matches against.
fn create_index_bufferize(
    src: Arc<UOp>,
    buf_ranges: Vec<Arc<UOp>>,
    idx_ranges: Vec<Arc<UOp>>,
    opts: BufferizeOpts,
) -> Arc<UOp> {
    let bufferized = UOp::bufferize(src, buf_ranges, opts);
    UOp::index().buffer(bufferized).indices(idx_ranges).call().expect("Failed to create INDEX")
}

/// Create a simple computation graph for testing.
///
/// Returns: (buffer, range1, range2, compute)
/// - buffer: BUFFER operation
/// - range1, range2: RANGE operations
/// - compute: Simple ADD operation on buffer
fn create_simple_graph(ctx: &mut IndexingContext) -> (Arc<UOp>, Arc<UOp>, Arc<UOp>, Arc<UOp>) {
    let buffer = create_test_buffer(100, DType::Float32, 1);
    let range1 = ctx.new_range(&SInt::Const(10), AxisType::Loop);
    let range2 = ctx.new_range(&SInt::Const(10), AxisType::Loop);

    // Create INDEX(buffer, [r1, r2])
    let indexed = UOp::index()
        .buffer(buffer.clone())
        .indices(vec![range1.clone(), range2.clone()])
        .call()
        .expect("Failed to create INDEX");

    // Create simple ADD operation
    let one = UOp::native_const(1.0f32);
    let compute = indexed.try_add(&one).expect("Failed to create ADD");

    (buffer, range1, range2, compute)
}

/// Create a computation graph that accesses multiple buffers.
///
/// Creates a computation that adds together `num_buffers` different buffers.
/// Returns: (buffers, ranges, compute)
/// - buffers: Vector of BUFFER operations
/// - ranges: Vector of RANGE operations used for indexing
/// - compute: Computation that accesses all buffers (nested ADD operations)
fn create_multi_buffer_graph(
    ctx: &mut IndexingContext,
    num_buffers: usize,
) -> (Vec<Arc<UOp>>, Vec<Arc<UOp>>, Arc<UOp>) {
    assert!(num_buffers > 0, "Must have at least one buffer");

    let mut buffers = Vec::new();
    let range1 = ctx.new_range(&SInt::Const(10), AxisType::Loop);
    let range2 = ctx.new_range(&SInt::Const(10), AxisType::Loop);
    let ranges = vec![range1.clone(), range2.clone()];

    // Create buffers and index them
    let mut compute = {
        let buffer = create_test_buffer(100, DType::Float32, 0);
        let indexed =
            UOp::index().buffer(buffer.clone()).indices(ranges.clone()).call().expect("Failed to create INDEX");
        buffers.push(buffer);
        indexed
    };

    // Add more buffers to the computation
    for i in 1..num_buffers {
        let buffer = create_test_buffer(100, DType::Float32, i);
        let indexed =
            UOp::index().buffer(buffer.clone()).indices(ranges.clone()).call().expect("Failed to create INDEX");
        compute = compute.try_add(&indexed).expect("Failed to create ADD");
        buffers.push(buffer);
    }

    (buffers, ranges, compute)
}

/// Create a buffer with specific size for ratio testing.
///
/// Returns a BUFFER operation with the given size in bytes.
fn create_buffer_with_size(size: usize, dtype: DType, id: usize) -> Arc<UOp> {
    create_test_buffer(size, dtype, id)
}

/// Create a computation graph with controlled output/input ratio.
///
/// Creates a pattern where one input buffer feeds into output buffer,
/// allowing precise control over the out_in_ratio for testing.
/// Returns: (input_buffer, output_buffer_size, ranges, compute)
fn create_ratio_test_graph(
    ctx: &mut IndexingContext,
    input_size: usize,
    output_size: usize,
) -> (Arc<UOp>, usize, Vec<Arc<UOp>>, Arc<UOp>) {
    let input_buffer = create_buffer_with_size(input_size, DType::Float32, 1);

    // Create ranges that would produce the desired output size
    // output_size = range_product * element_size
    // For Float32, element_size = 4 bytes
    let elements = output_size / 4;
    let range_size = (elements as f64).sqrt() as usize;

    let range1 = ctx.new_range(&SInt::Const(range_size), AxisType::Loop);
    let range2 = ctx.new_range(&SInt::Const(range_size), AxisType::Loop);
    let ranges = vec![range1, range2];

    // Create simple computation: just index the buffer
    let compute =
        UOp::index().buffer(input_buffer.clone()).indices(ranges.clone()).call().expect("Failed to create INDEX");

    (input_buffer, output_size, ranges, compute)
}

/// Create a computation graph with REDUCE operations.
///
/// Creates a reduce-sum pattern that either accesses a buffer or not.
/// Returns: (all_ranges, compute)
/// - all_ranges: Vector of all RANGE operations (mix of Loop and Reduce axes)
/// - compute: Computation that may or may not access buffers
fn create_reduce_graph(ctx: &mut IndexingContext, has_buffer_access: bool) -> (Vec<Arc<UOp>>, Arc<UOp>) {
    let loop_range = ctx.new_range(&SInt::Const(10), AxisType::Loop);
    let reduce_range = ctx.new_range(&SInt::Const(20), AxisType::Reduce);
    let all_ranges = vec![loop_range.clone(), reduce_range.clone()];

    let compute = if has_buffer_access {
        // Create computation that accesses a buffer
        let buffer = create_test_buffer(800, DType::Float32, 1); // 10 * 20 * 4 bytes
        let indexed = UOp::index().buffer(buffer).indices(all_ranges.clone()).call().expect("Failed to create INDEX");

        // Create REDUCE operation that accesses the buffer
        // Op::Reduce takes: src, ranges (the reduce axes), reduce_op
        UOp::new(
            Op::Reduce { src: indexed, ranges: vec![reduce_range].into(), reduce_op: morok_ir::ReduceOp::Add },
            DType::Float32,
        )
    } else {
        // Create REDUCE without buffer access (just reduce a constant)
        let const_val = UOp::native_const(1.0f32);
        UOp::new(
            Op::Reduce { src: const_val, ranges: vec![reduce_range].into(), reduce_op: morok_ir::ReduceOp::Add },
            DType::Float32,
        )
    };

    (all_ranges, compute)
}

// ============================================================================
// Phase 5.1: Basic Tests
// ============================================================================

#[test]
fn test_config_default() {
    // Test that default config has sensible values
    let config = PcontigConfig::default();
    assert_eq!(config.level, 2); // Enabled by default
    assert_eq!(config.max_buffers_threshold, 3);
    assert_eq!(config.out_in_ratio_threshold, 10.0);
}

#[test]
fn test_config_levels() {
    // Test different configuration levels
    let disabled = PcontigConfig { level: 0, ..Default::default() };
    let basic = PcontigConfig { level: 1, ..Default::default() };
    let enabled = PcontigConfig { level: 2, ..Default::default() };
    let aggressive = PcontigConfig { level: 3, ..Default::default() };

    assert_eq!(disabled.level, 0);
    assert_eq!(basic.level, 1);
    assert_eq!(enabled.level, 2);
    assert_eq!(aggressive.level, 3);
}

#[test]
fn test_pattern_matcher_creation() {
    // Test that pattern matcher can be created
    let matcher = buffer_removal_with_pcontig();

    // Pattern matcher should be created successfully
    // (We can't inspect internal state, but we can verify it doesn't panic)
    drop(matcher);
}

#[test]
fn test_disabled_config_no_rewrite() {
    // When level=0, pattern should not match
    let mut config = PcontigConfig { level: 0, ..Default::default() };
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let (_buffer, range1, range2, compute) = create_simple_graph(&mut ctx);

    // Create INDEX(BUFFERIZE(compute, [r1, r2]), [r1, r2])
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, vec![range1.clone(), range2.clone()], vec![range1, range2], opts);

    // Apply rewrite - should not change anything
    let rewritten = graph_rewrite(&matcher, idx_buf.clone(), &mut config);

    // With level=0, no rewrite should occur
    assert!(Arc::ptr_eq(&rewritten, &idx_buf), "Expected no rewrite with level=0");
}

#[test]
fn test_cheap_inline_removal() {
    // Test Pattern 1: cheap operations should be inlined
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    // Create BUFFERIZE(const, [range])
    let mut ctx = IndexingContext::new();
    let range = ctx.new_range(&SInt::Const(10), AxisType::Loop);
    let const_val = UOp::native_const(1.0f32);
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let bufferized = UOp::bufferize(const_val.clone(), vec![range], opts);

    // Apply rewrite - should remove BUFFERIZE and return const
    let rewritten = graph_rewrite(&matcher, bufferized, &mut config);

    // Should inline the constant
    assert!(matches!(rewritten.op(), Op::Const(_)), "Expected const to be inlined");
}

#[test]
fn test_nested_bufferize_removal() {
    // Test Pattern 3: nested BUFFERIZE should be flattened
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let range = ctx.new_range(&SInt::Const(10), AxisType::Loop);

    // Create nested BUFFERIZE(BUFFERIZE(const, r1), r2)
    let const_val = UOp::native_const(1.0f32);
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let inner = UOp::bufferize(const_val.clone(), vec![range.clone()], opts.clone());
    let outer = UOp::bufferize(inner, vec![range.clone()], opts);

    // Apply rewrite - should flatten to single BUFFERIZE or inline const
    let rewritten = graph_rewrite(&matcher, outer, &mut config);

    // After multiple rewrites, const should be fully inlined
    // (First rewrite removes nested bufferize, second inlines const)
    let final_result = graph_rewrite(&matcher, rewritten, &mut config);
    assert!(matches!(final_result.op(), Op::Const(_)), "Expected const to be fully inlined");
}

#[test]
fn test_simple_index_bufferize_pattern() {
    // Test that INDEX(BUFFERIZE) pattern is recognized
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let (_buffer, range1, range2, compute) = create_simple_graph(&mut ctx);

    // Create INDEX(BUFFERIZE(compute, [r1, r2]), [r1, r2])
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, vec![range1.clone(), range2.clone()], vec![range1, range2], opts);

    // Apply rewrite
    let rewritten = graph_rewrite(&matcher, idx_buf.clone(), &mut config);

    // Pattern should match, but exact behavior depends on heuristics
    // For now, just verify it doesn't panic
    drop(rewritten);
}

// ============================================================================
// Phase 5.2.1: accessed_buffers Heuristic Tests
// ============================================================================

/// Test accessed_buffers heuristic with different buffer counts.
///
/// The heuristic should keep buffers when there are >3 input buffers accessed,
/// as this indicates a complex multi-input operation.
///
/// Note: Uses level=0 to disable cheap inline pattern so we can test buffer count heuristic.
#[test_case(1 ; "one buffer - should optimize")]
#[test_case(2 ; "two buffers - should optimize")]
#[test_case(3 ; "three buffers - at threshold should optimize")]
#[test_case(4 ; "four buffers - above threshold should keep")]
#[test_case(5 ; "five buffers - above threshold should keep")]
fn test_accessed_buffers_threshold(num_buffers: usize) {
    let mut config = PcontigConfig { level: 0, ..PcontigConfig::default() };
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let (_buffers, ranges, compute) = create_multi_buffer_graph(&mut ctx, num_buffers);

    // Create INDEX(BUFFERIZE(compute, ranges), ranges)
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, ranges.clone(), ranges, opts);

    // Apply rewrite
    let rewritten = graph_rewrite(&matcher, idx_buf.clone(), &mut config);

    // Verify behavior based on threshold
    let threshold = config.max_buffers_threshold;
    if num_buffers > threshold {
        // Above threshold → should keep buffer (no rewrite)
        assert!(
            Arc::ptr_eq(&rewritten, &idx_buf),
            "Expected no rewrite with {} buffers (>{} threshold)",
            num_buffers,
            threshold
        );
    } else {
        // At or below threshold → may optimize (rewrite occurs)
        // We can't guarantee rewrite happens (depends on other heuristics)
        // But we verify it doesn't crash
        drop(rewritten);
    }
}

/// Test that duplicate buffer accesses are counted correctly.
///
/// If the same buffer is accessed multiple times, it should only be counted once.
#[test]
fn test_accessed_buffers_with_duplicates() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let buffer = create_test_buffer(100, DType::Float32, 1);
    let range1 = ctx.new_range(&SInt::Const(10), AxisType::Loop);
    let range2 = ctx.new_range(&SInt::Const(10), AxisType::Loop);
    let ranges = vec![range1.clone(), range2.clone()];

    // Access the same buffer multiple times
    let idx1 = UOp::index().buffer(buffer.clone()).indices(ranges.clone()).call().expect("Failed to create INDEX");
    let idx2 = UOp::index().buffer(buffer.clone()).indices(ranges.clone()).call().expect("Failed to create INDEX");
    let idx3 = UOp::index().buffer(buffer).indices(ranges.clone()).call().expect("Failed to create INDEX");

    // Create compute: idx1 + idx2 + idx3 (all same buffer)
    let compute = idx1.try_add(&idx2).expect("Failed to create ADD").try_add(&idx3).expect("Failed to create ADD");

    // Create INDEX(BUFFERIZE(compute, ranges), ranges)
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, ranges.clone(), ranges, opts);

    // Apply rewrite - should see this as 1 buffer accessed (not 3)
    let rewritten = graph_rewrite(&matcher, idx_buf, &mut config);

    // With only 1 unique buffer, should be eligible for optimization
    // (exact behavior depends on other heuristics, but shouldn't be blocked by accessed_buffers)
    drop(rewritten);
}

/// Test accessed_buffers with nested operations.
///
/// Buffers accessed in nested computations should all be counted.
///
/// Note: Uses level=0 to disable cheap inline pattern so we can test buffer count heuristic.
#[test]
fn test_accessed_buffers_nested_computation() {
    let mut config = PcontigConfig { level: 0, ..PcontigConfig::default() };
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let ranges = vec![ctx.new_range(&SInt::Const(10), AxisType::Loop), ctx.new_range(&SInt::Const(10), AxisType::Loop)];

    // Create nested structure: (buf1 + buf2) * (buf3 + buf4)
    let buf1 = create_test_buffer(100, DType::Float32, 1);
    let buf2 = create_test_buffer(100, DType::Float32, 2);
    let buf3 = create_test_buffer(100, DType::Float32, 3);
    let buf4 = create_test_buffer(100, DType::Float32, 4);

    let idx1 = UOp::index().buffer(buf1).indices(ranges.clone()).call().expect("Failed to create INDEX");
    let idx2 = UOp::index().buffer(buf2).indices(ranges.clone()).call().expect("Failed to create INDEX");
    let idx3 = UOp::index().buffer(buf3).indices(ranges.clone()).call().expect("Failed to create INDEX");
    let idx4 = UOp::index().buffer(buf4).indices(ranges.clone()).call().expect("Failed to create INDEX");

    let left = idx1.try_add(&idx2).expect("Failed to create ADD");
    let right = idx3.try_add(&idx4).expect("Failed to create ADD");
    let compute = left.try_mul(&right).expect("Failed to create MUL");

    // Create INDEX(BUFFERIZE(compute, ranges), ranges)
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, ranges.clone(), ranges, opts);

    // Apply rewrite - with 4 buffers (> threshold of 3), should keep buffer
    let rewritten = graph_rewrite(&matcher, idx_buf.clone(), &mut config);

    assert!(Arc::ptr_eq(&rewritten, &idx_buf), "Expected no rewrite with 4 buffers in nested computation");
}

// ============================================================================
// Phase 5.2.2: out_in_ratio Heuristic Tests
// ============================================================================

/// Test out_in_ratio heuristic with efficient buffer (ratio < threshold).
///
/// Buffers with low out/in ratio are memory-efficient and should be kept.
#[test]
fn test_out_in_ratio_efficient_buffer() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();

    // Create a buffer with ratio ~8.8 (< threshold of 10.0)
    // input_size = 1000 bytes, output_size = 9000 bytes target
    // Actual: 47*47*4 = 8836 bytes output, ratio = 8836/1000 ≈ 8.8
    let input_size = 1000;
    let output_size = 9000;
    let (_input_buffer, _output_size, ranges, compute) = create_ratio_test_graph(&mut ctx, input_size, output_size);

    // Create INDEX(BUFFERIZE(compute, ranges), ranges)
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, ranges.clone(), ranges, opts);

    // Apply rewrite — like Tinygrad, simple path (no buffer in reduce) always inlines
    // regardless of ratio. The ratio check only applies to the partial contiguous path.
    let rewritten = graph_rewrite(&matcher, idx_buf.clone(), &mut config);

    assert!(!Arc::ptr_eq(&rewritten, &idx_buf), "Simple path should always inline (no ratio guard, matching Tinygrad)");
    assert_eq!(count_bufferizes(&rewritten), 0, "BUFFERIZE should be removed after inlining");
}

/// Test out_in_ratio heuristic at threshold boundary.
///
/// Buffers at exactly the threshold (ratio = 10.0) test edge case behavior.
#[test]
fn test_out_in_ratio_at_threshold() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();

    // Create a buffer with ratio = 10.0 (exactly at threshold)
    // input_size = 1000 bytes, output_size = 10000 bytes
    let input_size = 1000;
    let output_size = 10000;
    let (_input_buffer, _output_size, ranges, compute) = create_ratio_test_graph(&mut ctx, input_size, output_size);

    // Create INDEX(BUFFERIZE(compute, ranges), ranges)
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, ranges.clone(), ranges, opts);

    // Apply rewrite
    let rewritten = graph_rewrite(&matcher, idx_buf, &mut config);

    // At threshold, behavior is implementation-defined
    // Just verify it doesn't crash
    drop(rewritten);
}

/// Test out_in_ratio heuristic with wasteful buffer (ratio >> threshold).
///
/// Buffers with high out/in ratio waste memory and should be optimized.
#[test]
fn test_out_in_ratio_wasteful_buffer() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();

    // Create a buffer with ratio = 50.0 (>> threshold of 10.0)
    // input_size = 1000 bytes, output_size = 50000 bytes
    let input_size = 1000;
    let output_size = 50000;
    let (_input_buffer, _output_size, ranges, compute) = create_ratio_test_graph(&mut ctx, input_size, output_size);

    // Create INDEX(BUFFERIZE(compute, ranges), ranges)
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, ranges.clone(), ranges, opts);

    // Apply rewrite - wasteful buffer should be optimized
    let rewritten = graph_rewrite(&matcher, idx_buf.clone(), &mut config);

    // With ratio >> 10, may optimize (depends on other heuristics)
    // At minimum, verify it doesn't crash and doesn't prevent optimization
    // We can't use ptr_eq here because other heuristics may apply
    drop(rewritten);
}

/// Test out_in_ratio with flash attention-like pattern.
///
/// Simulates flash attention where a 25M element buffer reduces to 512 elements.
/// This creates an extremely high ratio (~48828), demonstrating why partial
/// contiguous is beneficial for attention mechanisms.
#[test]
fn test_out_in_ratio_flash_attention_simulation() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();

    // Flash attention pattern: 25M → 512 elements
    // Input: 25M float32 = 100MB
    // Output: 512 float32 = 2KB
    // Ratio = 100MB / 2KB = 50000
    let input_size = 100_000_000; // 100 MB

    let input_buffer = create_buffer_with_size(input_size, DType::Float32, 1);

    // Create ranges for output (512 elements = 16x32)
    let range1 = ctx.new_range(&SInt::Const(16), AxisType::Loop);
    let range2 = ctx.new_range(&SInt::Const(32), AxisType::Loop);
    let ranges = vec![range1, range2];

    // Simple computation accessing the large buffer
    let compute = UOp::index().buffer(input_buffer).indices(ranges.clone()).call().expect("Failed to create INDEX");

    // Create INDEX(BUFFERIZE(compute, ranges), ranges)
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, ranges.clone(), ranges, opts);

    // Apply rewrite - extreme ratio should be eligible for optimization
    let rewritten = graph_rewrite(&matcher, idx_buf, &mut config);

    // This tests that extremely high ratios are handled correctly
    // The actual optimization depends on other heuristics (buffer_in_reduce)
    drop(rewritten);
}

/// Test out_in_ratio with symbolic (variable) sizes.
///
/// When sizes are symbolic, ratio calculation returns None,
/// and the heuristic should fall back to safe behavior (keep buffer).
#[test]
fn test_out_in_ratio_symbolic_sizes() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    // Create buffer with symbolic size via codegen PARAM
    let n = UOp::param(1, 1, DType::Index, None);
    let buffer = UOp::param(2, 1, DType::Float32, None);

    // Create ranges with symbolic size
    let range = UOp::range(n, 0);
    let ranges = vec![range.clone()];

    // Create computation
    let compute = UOp::index().buffer(buffer).indices(ranges.clone()).call().expect("Failed to create INDEX");

    // Create INDEX(BUFFERIZE(compute, ranges), ranges)
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, ranges.clone(), ranges, opts);

    // Apply rewrite - with symbolic sizes, ratio is None, should keep buffer
    let rewritten = graph_rewrite(&matcher, idx_buf.clone(), &mut config);

    // Symbolic sizes should cause safe fallback (keep buffer)
    // We can't guarantee NoMatch due to other patterns, but verify no crash
    drop(rewritten);
}

/// Test out_in_ratio with no input buffers (edge case).
///
/// When there are no input buffers, ratio calculation should handle gracefully.
/// The cheap_inline pattern should inline the constant.
///
/// Note: This test uses graph_rewrite_bottom_up because cheap_inline pattern
/// requires bottom-up traversal to see the optimized BUFFERIZE before INDEX.
#[test]
fn test_out_in_ratio_no_inputs() {
    use crate::rewrite::graph_rewrite_bottom_up;

    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let range = ctx.new_range(&SInt::Const(10), AxisType::Loop);
    let ranges = vec![range];

    // Compute with no buffer inputs (just a constant)
    let const_val = UOp::native_const(1.0f32);

    // Create INDEX(BUFFERIZE(const, ranges), ranges)
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(const_val, ranges.clone(), ranges, opts);

    // Apply rewrite - should be handled by cheap_inline pattern instead
    let rewritten = graph_rewrite_bottom_up(&matcher, idx_buf, &mut config);

    // Const should be inlined (cheap operation)
    assert!(matches!(rewritten.op(), Op::Const(_)), "Expected constant to be inlined");
}

// ============================================================================
// Phase 5.2.3: buffer_in_reduce Heuristic Tests
// ============================================================================

/// Test buffer_in_reduce heuristic with no REDUCE operations.
///
/// When there are no REDUCE operations in the computation, the buffer
/// should be eligible for full removal via substitution.
#[test]
fn test_buffer_not_in_reduce_full_removal() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();

    // Create simple computation WITHOUT reduce
    let buffer = create_test_buffer(100, DType::Float32, 1);
    let range1 = ctx.new_range(&SInt::Const(10), AxisType::Loop);
    let range2 = ctx.new_range(&SInt::Const(10), AxisType::Loop);
    let ranges = vec![range1.clone(), range2.clone()];

    // Just index + add (no reduce)
    let indexed = UOp::index().buffer(buffer).indices(ranges.clone()).call().expect("Failed to create INDEX");
    let one = UOp::native_const(1.0f32);
    let compute = indexed.try_add(&one).expect("Failed to create ADD");

    // Create INDEX(BUFFERIZE(compute, ranges), ranges)
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, ranges.clone(), ranges, opts);

    // Apply rewrite - without reduce, should do full removal
    let rewritten = graph_rewrite(&matcher, idx_buf, &mut config);

    // Verify transformation occurred (full removal via substitution)
    // The result should not contain BUFFERIZE
    drop(rewritten);
}

/// Test buffer_in_reduce heuristic with REDUCE accessing buffer.
///
/// When REDUCE operations access the buffer, partial contiguous should be applied
/// to materialize only the necessary dimensions.
#[test]
fn test_buffer_in_reduce_partial_contiguous() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();

    // Create computation WITH reduce that accesses buffer
    let (ranges, compute) = create_reduce_graph(&mut ctx, true);

    // Create INDEX(BUFFERIZE(compute, ranges), ranges)
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let bufferized = UOp::bufferize(compute, ranges.clone(), opts.clone());

    // For REDUCE patterns, we need an INDEX to trigger Pattern 4
    // But REDUCE results typically don't get indexed in the same way
    // So this test validates that REDUCE with buffer is detected correctly

    // Apply rewrite directly to bufferized (tests Pattern 1-3)
    let rewritten = graph_rewrite(&matcher, bufferized, &mut config);

    // With REDUCE accessing buffer, may apply partial contiguous or keep buffer
    // Depends on LOCAL index detection
    drop(rewritten);
}

/// Test buffer_in_reduce with REDUCE that doesn't access buffer.
///
/// If REDUCE exists but doesn't access the buffer, should still do full removal.
#[test]
fn test_reduce_without_buffer_access_full_removal() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();

    // Create computation with REDUCE but NO buffer access
    let (ranges, reduce_compute) = create_reduce_graph(&mut ctx, false);

    // Create a separate buffer access
    let buffer = create_test_buffer(100, DType::Float32, 2);
    let buffer_indexed = UOp::index().buffer(buffer).indices(ranges.clone()).call().expect("Failed to create INDEX");

    // Combine: buffer_access + reduce_const
    let compute = buffer_indexed.try_add(&reduce_compute).expect("Failed to create ADD");

    // Create INDEX(BUFFERIZE(compute, ranges), ranges)
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, ranges.clone(), ranges, opts);

    // Apply rewrite - REDUCE exists but doesn't access buffer, so full removal
    let rewritten = graph_rewrite(&matcher, idx_buf, &mut config);

    // Should still be eligible for optimization
    drop(rewritten);
}

/// Test buffer_in_reduce with multiple REDUCE axes.
///
/// Multiple REDUCE operations accessing the buffer should still be detected.
#[test]
fn test_multiple_reduces_with_buffer() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();

    // Create buffer and ranges
    let buffer = create_test_buffer(1600, DType::Float32, 1); // 10 * 10 * 4 * 4 bytes
    let loop_range = ctx.new_range(&SInt::Const(10), AxisType::Loop);
    let reduce_range1 = ctx.new_range(&SInt::Const(10), AxisType::Reduce);
    let reduce_range2 = ctx.new_range(&SInt::Const(4), AxisType::Reduce);
    let ranges = vec![loop_range.clone(), reduce_range1.clone(), reduce_range2.clone()];

    // Create INDEX(buffer, ranges)
    let indexed = UOp::index().buffer(buffer).indices(ranges.clone()).call().expect("Failed to create INDEX");

    // Create nested REDUCE: reduce over axis1, then axis2
    let reduce1 = UOp::new(
        Op::Reduce { src: indexed, ranges: vec![reduce_range1].into(), reduce_op: morok_ir::ReduceOp::Add },
        DType::Float32,
    );

    let reduce2 = UOp::new(
        Op::Reduce { src: reduce1, ranges: vec![reduce_range2].into(), reduce_op: morok_ir::ReduceOp::Add },
        DType::Float32,
    );

    // Create INDEX(BUFFERIZE(reduce2, ranges), ranges)
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let bufferized = UOp::bufferize(reduce2, ranges, opts);

    // Apply rewrite - multiple reduces accessing buffer
    let rewritten = graph_rewrite(&matcher, bufferized, &mut config);

    // Should detect buffer in reduce and apply appropriate transformation
    drop(rewritten);
}

/// Test buffer_in_reduce with nested REDUCE operations.
///
/// Nested reduces should be analyzed correctly for buffer access.
#[test]
fn test_nested_reduce_with_buffer() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();

    // Create outer reduce with buffer access
    let buffer = create_test_buffer(200, DType::Float32, 1);
    let outer_reduce_range = ctx.new_range(&SInt::Const(10), AxisType::Reduce);
    let ranges = vec![outer_reduce_range.clone()];

    let indexed = UOp::index().buffer(buffer).indices(ranges.clone()).call().expect("Failed to create INDEX");

    // Inner reduce (on the indexed buffer)
    let inner_reduce_range = ctx.new_range(&SInt::Const(5), AxisType::Reduce);
    let inner_reduce = UOp::new(
        Op::Reduce {
            src: indexed.clone(),
            ranges: vec![inner_reduce_range].into(),
            reduce_op: morok_ir::ReduceOp::Add,
        },
        DType::Float32,
    );

    // Outer reduce (wraps inner reduce)
    let outer_reduce = UOp::new(
        Op::Reduce { src: inner_reduce, ranges: vec![outer_reduce_range].into(), reduce_op: morok_ir::ReduceOp::Add },
        DType::Float32,
    );

    // Create BUFFERIZE(outer_reduce, ranges)
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let bufferized = UOp::bufferize(outer_reduce, ranges, opts);

    // Apply rewrite - nested reduces with buffer access
    let rewritten = graph_rewrite(&matcher, bufferized, &mut config);

    // Should detect buffer in nested reduce structure
    drop(rewritten);
}

// ============================================================================
// Phase 5.3.1: Full Removal Transformation Tests
// ============================================================================

/// Count the number of BUFFERIZE operations in a UOp tree.
#[allow(clippy::mutable_key_type)] // UOpKey contains interior mutability (hash-consed IR)
fn count_bufferizes(uop: &Arc<UOp>) -> usize {
    let mut count = 0;
    let mut stack = vec![uop.clone()];
    let mut visited = std::collections::HashSet::new();

    while let Some(current) = stack.pop() {
        if !visited.insert(morok_ir::UOpKey(current.clone())) {
            continue;
        }

        if matches!(current.op(), Op::Bufferize { .. }) {
            count += 1;
        }

        for src in current.op().sources() {
            stack.push(src.clone());
        }
    }

    count
}

/// Test Pattern 1: BUFFERIZE(const) → const.
///
/// Cheap operations should always be inlined.
#[test]
fn test_pattern1_cheap_inline() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let range = ctx.new_range(&SInt::Const(10), AxisType::Loop);

    // BUFFERIZE(const) - Pattern 1 should inline this
    let const_val = UOp::native_const(42.0f32);
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let bufferized = UOp::bufferize(const_val.clone(), vec![range], opts);

    // Apply rewrite - Pattern 1 should match
    let rewritten = graph_rewrite(&matcher, bufferized, &mut config);

    // Verify: BUFFERIZE removed, const returned
    assert!(matches!(rewritten.op(), Op::Const(_)), "Pattern 1 should inline const");
    let count = count_bufferizes(&rewritten);
    assert_eq!(count, 0, "No BUFFERIZE should remain after Pattern 1");
}

/// Test Pattern 4: Full removal transformation with permissive config.
///
/// Uses permissive thresholds to test that the full removal path works correctly.
#[test]
fn test_pattern4_full_removal_with_permissive_config() {
    // Use permissive config that allows optimization
    let mut config = PcontigConfig {
        level: 2,
        max_buffers_threshold: 10,   // Allow many buffers
        out_in_ratio_threshold: 1.0, // Allow any ratio >= 1.0
    };
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();

    // Create simple case: INDEX(BUFFERIZE(INDEX(buffer)))
    let buffer = create_test_buffer(40, DType::Float32, 1);
    let range = ctx.new_range(&SInt::Const(10), AxisType::Loop);

    // Compute: just index the buffer
    let indexed = UOp::index().buffer(buffer).indices(vec![range.clone()]).call().expect("Failed to create INDEX");
    let compute = indexed;

    // Create INDEX(BUFFERIZE(compute, [range]), [range])
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, vec![range.clone()], vec![range], opts);

    let bufferizes_before = count_bufferizes(&idx_buf);

    // Apply rewrite - with permissive config, Pattern 4 should match
    let rewritten = graph_rewrite(&matcher, idx_buf.clone(), &mut config);

    let bufferizes_after = count_bufferizes(&rewritten);

    // Verify optimization occurred
    assert!(
        bufferizes_after < bufferizes_before,
        "Expected BUFFERIZE removal with permissive config, before={}, after={}",
        bufferizes_before,
        bufferizes_after
    );
}

/// Test Pattern 4: Efficient buffer is kept (ratio < 10.0).
///
/// When ratio < threshold, buffer should be KEPT (not optimized).
#[test]
fn test_pattern4_keeps_efficient_buffer() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();

    // Create buffer with ratio = 1.0 (efficient)
    // Input: 40 bytes, Output: 40 bytes
    let input_buffer = create_test_buffer(40, DType::Float32, 1);
    let range = ctx.new_range(&SInt::Const(10), AxisType::Loop); // 10 * 4 = 40 bytes

    let indexed =
        UOp::index().buffer(input_buffer).indices(vec![range.clone()]).call().expect("Failed to create INDEX");
    let compute = indexed;

    // Create INDEX(BUFFERIZE(compute, [range]), [range])
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, vec![range.clone()], vec![range], opts);

    // Apply rewrite — like Tinygrad, simple path (no buffer in reduce) always inlines
    let rewritten = graph_rewrite(&matcher, idx_buf.clone(), &mut config);

    assert!(!Arc::ptr_eq(&rewritten, &idx_buf), "Simple path should always inline (matching Tinygrad)");
}

/// Test that Pattern 1 preserves dtype.
///
/// Inlining should maintain the dtype of the original computation.
#[test]
fn test_pattern1_preserves_dtype() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let range = ctx.new_range(&SInt::Const(10), AxisType::Loop);

    // Create BUFFERIZE(const) with specific dtype
    let const_val = UOp::native_const(42.0f32);
    let original_dtype = const_val.dtype();
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let bufferized = UOp::bufferize(const_val, vec![range], opts);

    // Apply rewrite
    let rewritten = graph_rewrite(&matcher, bufferized, &mut config);

    // Verify dtype preserved
    assert_eq!(rewritten.dtype(), original_dtype, "Dtype should be preserved after Pattern 1");
}

/// Test that full removal doesn't occur when heuristics prevent it.
///
/// If accessed_buffers > threshold, full removal should not occur.
///
/// Note: Uses level=0 to disable cheap inline pattern so we can test buffer count heuristic.
#[test]
fn test_full_removal_blocked_by_heuristics() {
    let mut config = PcontigConfig { level: 0, ..PcontigConfig::default() };
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();

    // Create computation with >3 buffers (exceeds threshold)
    let (_buffers, ranges, compute) = create_multi_buffer_graph(&mut ctx, 4);

    // Create INDEX(BUFFERIZE(compute, ranges), ranges)
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, ranges.clone(), ranges, opts);

    // Apply rewrite - should NOT do full removal due to accessed_buffers heuristic
    let rewritten = graph_rewrite(&matcher, idx_buf.clone(), &mut config);

    // Should keep buffer (no rewrite)
    assert!(Arc::ptr_eq(&rewritten, &idx_buf), "Expected no rewrite when heuristics prevent optimization");
}

// ============================================================================
// Phase 5.3.2: Partial Contiguous Transformation Tests
// ============================================================================

/// Test partial contiguous with single REDUCE dimension.
///
/// When computation includes REDUCE accessing buffer, should materialize reduce dimension
/// while inlining LOOP dimensions.
#[test]
fn test_partial_contiguous_single_reduce() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let buffer = create_test_buffer(4000, DType::Float32, 1); // 100 * 10 * 4 bytes

    // Create LOOP and REDUCE ranges
    let loop_range = ctx.new_range(&SInt::Const(100), AxisType::Loop);
    let reduce_range = ctx.new_range(&SInt::Const(10), AxisType::Reduce);
    let all_ranges = vec![loop_range.clone(), reduce_range.clone()];

    // Create computation: REDUCE(INDEX(buffer, [loop, reduce]), [reduce])
    let indexed = UOp::index().buffer(buffer).indices(all_ranges.clone()).call().expect("Failed to create INDEX");
    let reduce = UOp::new(
        Op::Reduce { src: indexed, ranges: vec![reduce_range.clone()].into(), reduce_op: morok_ir::ReduceOp::Add },
        DType::Float32,
    );

    // Create INDEX(BUFFERIZE(reduce, [loop]), [loop])
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let bufferized = UOp::bufferize(reduce.clone(), vec![loop_range.clone()], opts);
    let idx_buf = UOp::index().buffer(bufferized).indices(vec![loop_range]).call().expect("Failed to create INDEX");

    // Apply rewrite - should apply partial contiguous
    let rewritten = graph_rewrite(&matcher, idx_buf.clone(), &mut config);

    // Verify transformation occurred (may be partial contiguous or full removal depending on heuristics)
    // This is primarily a smoke test to ensure reduce handling works
    drop(rewritten);
}

/// Test partial contiguous with LOCAL dimension.
///
/// LOCAL axes should be materialized, LOOP axes should be inlined.
#[test]
fn test_partial_contiguous_local_axis() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let buffer = create_test_buffer(400, DType::Float32, 1); // 10 * 10 * 4 bytes

    // Create LOOP and LOCAL ranges
    let loop_range = ctx.new_range(&SInt::Const(10), AxisType::Loop);
    let local_range = ctx.new_range(&SInt::Const(10), AxisType::Local);
    let all_ranges = vec![loop_range.clone(), local_range.clone()];

    // Create computation that indexes buffer with both dimensions
    let indexed = UOp::index().buffer(buffer).indices(all_ranges.clone()).call().expect("Failed to create INDEX");
    let two = UOp::native_const(2.0f32);
    let compute = indexed.try_mul(&two).expect("Failed to create MUL");

    // Create INDEX(BUFFERIZE(compute, all_ranges), all_ranges)
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, all_ranges.clone(), all_ranges, opts);

    // Apply rewrite
    let rewritten = graph_rewrite(&matcher, idx_buf.clone(), &mut config);

    // Verify transformation (LOCAL should trigger partial contiguous)
    // The exact behavior depends on heuristics
    drop(rewritten);
}

/// Test partial contiguous with mixed axes (LOOP + REDUCE).
///
/// Should materialize REDUCE dimension, inline LOOP dimension.
#[test]
fn test_partial_contiguous_mixed_axes() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let buffer = create_test_buffer(8000, DType::Float32, 1); // 100 * 20 * 4 bytes

    let loop_range = ctx.new_range(&SInt::Const(100), AxisType::Loop);
    let reduce_range = ctx.new_range(&SInt::Const(20), AxisType::Reduce);
    let all_ranges = vec![loop_range.clone(), reduce_range.clone()];

    // Create REDUCE accessing buffer
    let indexed = UOp::index().buffer(buffer).indices(all_ranges.clone()).call().expect("Failed to create INDEX");
    let reduce = UOp::new(
        Op::Reduce { src: indexed, ranges: vec![reduce_range].into(), reduce_op: morok_ir::ReduceOp::Max },
        DType::Float32,
    );

    // Create INDEX(BUFFERIZE(reduce, [loop]), [loop])
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let bufferized = UOp::bufferize(reduce, vec![loop_range.clone()], opts);
    let idx_buf = UOp::index().buffer(bufferized).indices(vec![loop_range]).call().expect("Failed to create INDEX");

    // Apply rewrite
    let rewritten = graph_rewrite(&matcher, idx_buf, &mut config);

    // Verify transformation
    drop(rewritten);
}

/// Test partial contiguous with different reduce operations.
///
/// Should work correctly with Max, Mul, etc.
#[test]
fn test_partial_contiguous_different_reduce_ops() {
    use morok_ir::ReduceOp;

    for reduce_op in [ReduceOp::Add, ReduceOp::Max, ReduceOp::Mul] {
        let mut config = PcontigConfig::default();
        let matcher = buffer_removal_with_pcontig();

        let mut ctx = IndexingContext::new();
        let buffer = create_test_buffer(400, DType::Float32, 1);

        let loop_range = ctx.new_range(&SInt::Const(10), AxisType::Loop);
        let reduce_range = ctx.new_range(&SInt::Const(10), AxisType::Reduce);

        let indexed = UOp::index()
            .buffer(buffer)
            .indices(vec![loop_range.clone(), reduce_range.clone()])
            .call()
            .expect("Failed to create INDEX");
        let reduce = indexed.reduce(smallvec::smallvec![reduce_range], reduce_op);

        let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
        let bufferized = UOp::bufferize(reduce, vec![loop_range.clone()], opts);
        let idx_buf = UOp::index().buffer(bufferized).indices(vec![loop_range]).call().expect("Failed to create INDEX");

        // Apply rewrite - should not panic regardless of reduce op
        let rewritten = graph_rewrite(&matcher, idx_buf, &mut config);
        drop(rewritten);
    }
}

/// Test partial contiguous with multi-dimensional reduce.
///
/// Multiple REDUCE dimensions should all be materialized.
#[test]
fn test_partial_contiguous_multi_dimensional_reduce() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let buffer = create_test_buffer(8000, DType::Float32, 1); // 10 * 20 * 10 * 4 bytes

    let loop_range = ctx.new_range(&SInt::Const(10), AxisType::Loop);
    let reduce_range1 = ctx.new_range(&SInt::Const(20), AxisType::Reduce);
    let reduce_range2 = ctx.new_range(&SInt::Const(10), AxisType::Reduce);

    // Index buffer with all three dimensions
    let indexed = UOp::index()
        .buffer(buffer)
        .indices(vec![loop_range.clone(), reduce_range1.clone(), reduce_range2.clone()])
        .call()
        .expect("Failed to create INDEX");

    // Reduce over both reduce dimensions
    let reduce = UOp::new(
        Op::Reduce {
            src: indexed,
            ranges: vec![reduce_range1, reduce_range2].into(),
            reduce_op: morok_ir::ReduceOp::Add,
        },
        DType::Float32,
    );

    // Create INDEX(BUFFERIZE(reduce, [loop]), [loop])
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let bufferized = UOp::bufferize(reduce, vec![loop_range.clone()], opts);
    let idx_buf = UOp::index().buffer(bufferized).indices(vec![loop_range]).call().expect("Failed to create INDEX");

    // Apply rewrite
    let rewritten = graph_rewrite(&matcher, idx_buf, &mut config);
    drop(rewritten);
}

/// Test that partial contiguous respects heuristics.
///
/// Even with REDUCE, if accessed_buffers > threshold, should not optimize.
#[test]
fn test_partial_contiguous_blocked_by_heuristics() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();

    // Create computation with >3 buffers AND reduce
    let buf1 = create_test_buffer(40, DType::Float32, 1);
    let buf2 = create_test_buffer(40, DType::Float32, 2);
    let buf3 = create_test_buffer(40, DType::Float32, 3);
    let buf4 = create_test_buffer(40, DType::Float32, 4);

    let loop_range = ctx.new_range(&SInt::Const(10), AxisType::Loop);
    let reduce_range = ctx.new_range(&SInt::Const(10), AxisType::Reduce);

    // Index all buffers
    let idx1 = UOp::index().buffer(buf1).indices(vec![loop_range.clone()]).call().expect("Failed to create INDEX");
    let idx2 = UOp::index().buffer(buf2).indices(vec![loop_range.clone()]).call().expect("Failed to create INDEX");
    let idx3 = UOp::index().buffer(buf3).indices(vec![loop_range.clone()]).call().expect("Failed to create INDEX");
    let idx4 = UOp::index().buffer(buf4).indices(vec![loop_range.clone()]).call().expect("Failed to create INDEX");

    // Combine them
    let add1 = idx1.try_add(&idx2).expect("Failed to create ADD");
    let add2 = idx3.try_add(&idx4).expect("Failed to create ADD");
    let combined = add1.try_add(&add2).expect("Failed to create ADD");

    // Add reduce
    let reduce = UOp::new(
        Op::Reduce { src: combined, ranges: vec![reduce_range].into(), reduce_op: morok_ir::ReduceOp::Add },
        DType::Float32,
    );

    // Create INDEX(BUFFERIZE(reduce, [loop]), [loop])
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let bufferized = UOp::bufferize(reduce, vec![loop_range.clone()], opts);
    let idx_buf = UOp::index().buffer(bufferized).indices(vec![loop_range]).call().expect("Failed to create INDEX");

    // Apply rewrite - should be blocked by accessed_buffers heuristic
    let rewritten = graph_rewrite(&matcher, idx_buf.clone(), &mut config);

    // Should not rewrite (4 buffers > 3 threshold)
    assert!(Arc::ptr_eq(&rewritten, &idx_buf), "Expected no rewrite when accessed_buffers > threshold");
}

// ============================================================================
// Phase 5.3.3: Edge Case Transformation Tests
// ============================================================================

/// Test edge case: BUFFERIZE with no computation (just buffer passthrough).
///
/// Should handle gracefully without panicking.
#[test]
fn test_edge_case_empty_computation() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let buffer = create_test_buffer(40, DType::Float32, 1);
    let range = ctx.new_range(&SInt::Const(10), AxisType::Loop);

    // Create INDEX(buffer) - minimal computation
    let indexed = UOp::index().buffer(buffer).indices(vec![range.clone()]).call().expect("Failed to create INDEX");

    // Bufferize it
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(indexed, vec![range.clone()], vec![range], opts);

    // Apply rewrite - should not panic
    let rewritten = graph_rewrite(&matcher, idx_buf, &mut config);
    drop(rewritten);
}

/// Test edge case: All const operations (no buffer access).
///
/// Should inline completely via Pattern 1.
#[test]
fn test_edge_case_all_const_operations() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let range = ctx.new_range(&SInt::Const(10), AxisType::Loop);

    // Create computation with only constants
    let const1 = UOp::native_const(1.0f32);
    let const2 = UOp::native_const(2.0f32);
    let compute = const1.try_add(&const2).expect("Failed to create ADD");

    // Bufferize it
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let bufferized = UOp::bufferize(compute.clone(), vec![range], opts);

    // Like Tinygrad, only bare BUFFERIZE(CONST) is folded — BUFFERIZE(ADD(const, const))
    // is NOT a bare constant. In the real pipeline, ADD(const, const) would be folded by
    // symbolic patterns first, then BUFFERIZE(CONST) would be caught.
    let rewritten = graph_rewrite(&matcher, bufferized.clone(), &mut config);

    assert!(Arc::ptr_eq(&rewritten, &bufferized), "Bare BUFFERIZE(non-const compute) should not be removed");
}

/// Test edge case: Deeply nested BUFFERIZE operations.
///
/// Multiple levels of BUFFERIZE should be flattened correctly.
#[test]
fn test_edge_case_deeply_nested_bufferize() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let range = ctx.new_range(&SInt::Const(5), AxisType::Loop);

    // Create deeply nested: BUFFERIZE(BUFFERIZE(BUFFERIZE(const)))
    let const_val = UOp::native_const(42.0f32);
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };

    let level1 = UOp::bufferize(const_val, vec![range.clone()], opts.clone());
    let level2 = UOp::bufferize(level1, vec![range.clone()], opts.clone());
    let level3 = UOp::bufferize(level2, vec![range], opts);

    // Apply multiple rewrites to flatten
    let rewritten1 = graph_rewrite(&matcher, level3, &mut config);
    let rewritten2 = graph_rewrite(&matcher, rewritten1, &mut config);

    // Should eventually inline to const
    drop(rewritten2);
}

/// Test edge case: Zero-sized buffer.
///
/// Should handle edge case without panicking.
#[test]
fn test_edge_case_zero_sized_range() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let buffer = create_test_buffer(40, DType::Float32, 1);
    let zero_range = ctx.new_range(&SInt::Const(0), AxisType::Loop);

    // Create computation with zero-sized range
    let indexed = UOp::index().buffer(buffer).indices(vec![zero_range.clone()]).call().expect("Failed to create INDEX");
    let two = UOp::native_const(2.0f32);
    let compute = indexed.try_mul(&two).expect("Failed to create MUL");

    // Bufferize with zero range
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, vec![zero_range.clone()], vec![zero_range], opts);

    // Apply rewrite - should handle gracefully
    let rewritten = graph_rewrite(&matcher, idx_buf, &mut config);
    drop(rewritten);
}

// ============================================================================
// Phase 5.4.1: Configuration Level Tests
// ============================================================================

/// Test that custom max_buffers_threshold changes behavior.
///
/// A config with higher threshold should optimize more cases.
///
/// Note: Uses level=0 to disable cheap inline pattern so we can test buffer count heuristic.
#[test]
fn test_config_custom_max_buffers_threshold() {
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();

    // Create computation with 4 buffers (above default threshold, below permissive)
    let (_buffers, ranges, compute) = create_multi_buffer_graph(&mut ctx, 4);

    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, ranges.clone(), ranges, opts);

    // Default config should NOT optimize (4 > 3)
    let mut default_config = PcontigConfig { level: 0, ..PcontigConfig::default() };
    let default_rewritten = graph_rewrite(&matcher, idx_buf.clone(), &mut default_config);
    assert!(Arc::ptr_eq(&default_rewritten, &idx_buf), "Default config should block 4 buffers");

    // Permissive config MIGHT optimize (4 < 10), depends on ratio heuristic
    let mut permissive_config = PcontigConfig {
        level: 0,
        max_buffers_threshold: 10, // Allow up to 10 buffers
        ..Default::default()
    };
    let permissive_rewritten = graph_rewrite(&matcher, idx_buf, &mut permissive_config);
    // Don't assert specific behavior since ratio heuristic might still block
    drop(permissive_rewritten);
}

/// Test that custom out_in_ratio_threshold changes behavior.
///
/// A config with lower threshold should optimize fewer cases.
/// Uses REDUCE in compute to avoid cheap_inline pattern (REDUCE is not cheap).
#[test]
fn test_config_custom_ratio_threshold() {
    use morok_ir::types::ReduceOp;

    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let buffer = create_test_buffer(40, DType::Float32, 1);
    let outer_range = ctx.new_range(&SInt::Const(10), AxisType::Loop);
    let reduce_range = ctx.new_range(&SInt::Const(5), AxisType::Reduce);

    // Create computation with REDUCE (not cheap to inline): REDUCE(INDEX(buffer), range, Add)
    let indexed = UOp::index()
        .buffer(buffer)
        .indices(vec![outer_range.clone(), reduce_range.clone()])
        .call()
        .expect("Failed to create INDEX");
    let compute = indexed.reduce(smallvec::smallvec![reduce_range], ReduceOp::Add);

    // Create INDEX(BUFFERIZE(compute, [outer_range]), [outer_range])
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, vec![outer_range.clone()], vec![outer_range], opts);

    // Strict config: low ratio → should KEEP buffer
    let mut strict_config = PcontigConfig {
        out_in_ratio_threshold: 100.0, // Only optimize if ratio >= 100
        ..Default::default()
    };
    let strict_rewritten = graph_rewrite(&matcher, idx_buf.clone(), &mut strict_config);
    assert!(Arc::ptr_eq(&strict_rewritten, &idx_buf), "Strict config should keep buffer (low ratio)");

    // Permissive config: may optimize depending on ratio
    let mut permissive_config = PcontigConfig {
        out_in_ratio_threshold: 0.1, // Optimize if ratio >= 0.1
        ..Default::default()
    };
    let permissive_rewritten = graph_rewrite(&matcher, idx_buf.clone(), &mut permissive_config);
    // Don't assert specific behavior since buffer_in_reduce heuristic might still block
    drop(permissive_rewritten);
}

/// Test level 0 (disabled) vs level 2 (enabled) behavior.
///
/// Level 0 should disable Pattern 4 (complex heuristic-based optimization).
/// Note: Patterns 1-3 (cheap inline, nested bufferize) still run with level 0.
#[test]
fn test_config_level_0_vs_2() {
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let buffer = create_test_buffer(40, DType::Float32, 1);
    let range = ctx.new_range(&SInt::Const(10), AxisType::Loop);

    // Create INDEX(BUFFERIZE(...)) pattern that triggers Pattern 4
    let indexed = UOp::index().buffer(buffer).indices(vec![range.clone()]).call().expect("Failed to create INDEX");
    let one = UOp::native_const(1.0f32);
    let compute = indexed.try_add(&one).expect("Failed to create ADD");

    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, vec![range.clone()], vec![range], opts);

    // Level 0: Pattern 4 disabled → should NOT optimize
    let mut disabled_config = PcontigConfig { level: 0, ..Default::default() };
    let disabled_rewritten = graph_rewrite(&matcher, idx_buf.clone(), &mut disabled_config);
    assert!(Arc::ptr_eq(&disabled_rewritten, &idx_buf), "Level 0 should disable Pattern 4 optimizations");

    // Level 2: Pattern 4 enabled → should optimize (ratio=1.0 >= threshold)
    let mut enabled_config = PcontigConfig {
        level: 2,
        out_in_ratio_threshold: 1.0, // Permissive to ensure optimization happens
        ..Default::default()
    };
    let enabled_rewritten = graph_rewrite(&matcher, idx_buf.clone(), &mut enabled_config);
    assert!(!Arc::ptr_eq(&enabled_rewritten, &idx_buf), "Level 2 should enable Pattern 4 optimizations");
}

/// Test that the same matcher with different configs produces different results.
///
/// Each graph_rewrite call should respect its own config.
///
/// Note: Uses level=0 to disable cheap inline pattern so we can test buffer count heuristic.
#[test]
fn test_config_different_configs_produce_different_results() {
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();

    // Create computation with 3 buffers
    let (_buffers, ranges, compute) = create_multi_buffer_graph(&mut ctx, 3);

    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, ranges.clone(), ranges, opts);

    // Config1 (threshold=2): should NOT optimize (3 > 2)
    let mut config1 = PcontigConfig { level: 0, max_buffers_threshold: 2, ..Default::default() };
    let rewritten1 = graph_rewrite(&matcher, idx_buf.clone(), &mut config1);
    assert!(Arc::ptr_eq(&rewritten1, &idx_buf), "Config1 should block 3 buffers (threshold=2)");

    // Config2 (threshold=5): MIGHT optimize (3 < 5), depends on ratio
    let mut config2 = PcontigConfig { level: 0, max_buffers_threshold: 5, ..Default::default() };
    let rewritten2 = graph_rewrite(&matcher, idx_buf, &mut config2);
    // Don't assert specific behavior since ratio might still block
    drop(rewritten2);
}

// ============================================================================
// Phase 5.4.2: Pipeline Integration Tests
// ============================================================================

/// Test that buffer removal integrates with full rangeify pipeline.
///
/// This verifies Step 7 (buffer removal) works correctly after Steps 1-6.
#[test]
fn test_pipeline_integration_full_rangeify() {
    use crate::rangeify::rangeify;

    // Create a simple PERMUTE operation using a Buffer (pre-kernel pipeline uses Buffer, not codegen Param)
    let src = UOp::new_buffer(morok_device::DeviceSpec::Cpu, 6, DType::Float32);
    let reshaped = src.try_reshape(&smallvec::smallvec![morok_ir::SInt::Const(2), morok_ir::SInt::Const(3)]).unwrap();
    let permute = reshaped.try_permute(vec![1, 0]).unwrap();

    // Run full rangeify pipeline (includes buffer removal at Step 7)
    let (result, _ctx) = rangeify(permute, None).expect("Rangeify should succeed");

    // Verify pipeline completed without panicking
    assert_eq!(result.dtype(), DType::Float32);
}

/// Test that Patterns 1-4 can fire in sequence within one graph_rewrite call.
///
/// Complex graphs may need multiple patterns to fully optimize.
#[test]
fn test_pipeline_multiple_patterns_in_sequence() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let range = ctx.new_range(&SInt::Const(10), AxisType::Loop);

    // Create nested BUFFERIZE(BUFFERIZE(const))
    // Pattern 3 should remove inner BUFFERIZE
    // Pattern 1 should then remove outer BUFFERIZE
    let const_val = UOp::native_const(1.0f32);
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let inner = UOp::bufferize(const_val.clone(), vec![range.clone()], opts.clone());
    let outer = UOp::bufferize(inner, vec![range], opts);

    // Single rewrite should apply Pattern 3
    let rewritten1 = graph_rewrite(&matcher, outer.clone(), &mut config);
    assert!(!Arc::ptr_eq(&rewritten1, &outer), "Pattern 3 should fire");

    // Second rewrite should apply Pattern 1
    let rewritten2 = graph_rewrite(&matcher, rewritten1.clone(), &mut config);
    // Might inline to const depending on graph_rewrite's bottom-up traversal
    drop(rewritten2);
}

/// Test that buffer removal preserves graph structure correctness.
///
/// After optimization, the graph should still be semantically valid.
#[test]
fn test_pipeline_preserves_graph_structure() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let buffer = create_test_buffer(40, DType::Float32, 1);
    let range = ctx.new_range(&SInt::Const(10), AxisType::Loop);

    // Create a computation chain: INDEX(buf) → MUL → ADD
    let indexed = UOp::index().buffer(buffer).indices(vec![range.clone()]).call().expect("Failed to create INDEX");
    let two = UOp::native_const(2.0f32);
    let mul = indexed.try_mul(&two).expect("Failed to create MUL");
    let one = UOp::native_const(1.0f32);
    let add = mul.try_add(&one).expect("Failed to create ADD");

    // Bufferize and index
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(add, vec![range.clone()], vec![range], opts);

    // Original dtype
    let original_dtype = idx_buf.dtype();

    // Apply rewrite
    let rewritten = graph_rewrite(&matcher, idx_buf, &mut config);

    // Dtype should be preserved
    assert_eq!(rewritten.dtype(), original_dtype, "Dtype should be preserved");
}

/// Test interaction with cheap inline operations.
///
/// Pattern 1 (cheap inline) should work correctly with Pattern 4.
#[test]
fn test_pipeline_cheap_inline_interaction() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let range = ctx.new_range(&SInt::Const(10), AxisType::Loop);

    // Create BUFFERIZE(unary_op(const)) - both unary and const are cheap
    let const_val = UOp::native_const(5.0f32);
    let neg = const_val.neg();

    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let bufferized = UOp::bufferize(neg.clone(), vec![range], opts);

    // Like Tinygrad, only bare BUFFERIZE(CONST) is folded — not BUFFERIZE(unary(CONST)).
    // Non-const bare BUFFERIZEs survive to become STORE/BUFFER in later stages.
    let rewritten = graph_rewrite(&matcher, bufferized.clone(), &mut config);
    assert!(Arc::ptr_eq(&rewritten, &bufferized), "Bare BUFFERIZE(non-const) should not be removed");
}

// ============================================================================
// Phase 5.5.1: Symbolic & Dynamic Size Tests
// ============================================================================

/// Test that symbolic buffer sizes are handled correctly.
///
/// When buffer size cannot be computed (symbolic ranges), heuristics should skip.
#[test]
fn test_symbolic_buffer_size_handling() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    // Create symbolic range using DEFINE_VAR
    let batch_size = UOp::define_var("batch".to_string(), 0, 128);
    let symbolic_range = UOp::range_axis(batch_size, AxisId::Renumbered(0), AxisType::Loop);

    let buffer = create_test_buffer(40, DType::Float32, 1);

    // Create computation with symbolic output range
    let concrete_range = UOp::new(
        Op::Range {
            end: UOp::index_const(10),
            axis_id: AxisId::Renumbered(1),
            axis_type: AxisType::Loop,
            deps: smallvec::SmallVec::new(),
        },
        DType::Index,
    );

    let indexed = UOp::index().buffer(buffer).indices(vec![concrete_range]).call().expect("Failed to create INDEX");
    let one = UOp::native_const(1.0f32);
    let compute = indexed.try_add(&one).expect("Failed to create ADD");

    // Bufferize with symbolic range
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let bufferized = UOp::bufferize(compute, vec![symbolic_range.clone()], opts);
    let idx_buf = UOp::index().buffer(bufferized).indices(vec![symbolic_range]).call().expect("Failed to create INDEX");

    // Apply rewrite - should handle gracefully (ratio calculation returns None)
    let rewritten = graph_rewrite(&matcher, idx_buf, &mut config);
    drop(rewritten);
}

/// Test with all symbolic inputs and outputs.
///
/// Pattern 4 should skip when all sizes are symbolic.
#[test]
fn test_all_symbolic_sizes() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    // Create symbolic ranges
    let n = UOp::define_var("n".to_string(), 0, 1024);
    let m = UOp::define_var("m".to_string(), 0, 1024);
    let range_n = UOp::range_axis(n, AxisId::Renumbered(0), AxisType::Loop);
    let range_m = UOp::range_axis(m, AxisId::Renumbered(1), AxisType::Loop);

    // Create symbolic buffer (we can't, so use concrete buffer)
    let buffer = create_test_buffer(4096, DType::Float32, 1);

    let indexed = UOp::index().buffer(buffer).indices(vec![range_n.clone()]).call().expect("Failed to create INDEX");
    let two = UOp::native_const(2.0f32);
    let compute = indexed.try_mul(&two).expect("Failed to create MUL");

    // Bufferize with symbolic ranges
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let bufferized = UOp::bufferize(compute, vec![range_m.clone()], opts);
    let idx_buf = UOp::index().buffer(bufferized).indices(vec![range_m]).call().expect("Failed to create INDEX");

    // Apply rewrite - should not crash with symbolic sizes
    let rewritten = graph_rewrite(&matcher, idx_buf, &mut config);
    drop(rewritten);
}

/// Test mixed concrete and symbolic sizes.
///
/// Heuristics should handle mixed scenarios gracefully.
#[test]
fn test_mixed_concrete_symbolic_sizes() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    // Concrete range
    let concrete_range = UOp::new(
        Op::Range {
            end: UOp::index_const(10),
            axis_id: AxisId::Renumbered(0),
            axis_type: AxisType::Loop,
            deps: smallvec::SmallVec::new(),
        },
        DType::Index,
    );

    // Symbolic range
    let batch = UOp::define_var("batch".to_string(), 0, 64);
    let symbolic_range = UOp::range_axis(batch, AxisId::Renumbered(1), AxisType::Loop);

    let buffer = create_test_buffer(40, DType::Float32, 1);
    let indexed =
        UOp::index().buffer(buffer).indices(vec![concrete_range.clone()]).call().expect("Failed to create INDEX");
    let one = UOp::native_const(1.0f32);
    let compute = indexed.try_add(&one).expect("Failed to create ADD");

    // Bufferize with mixed ranges
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let bufferized = UOp::bufferize(compute, vec![concrete_range.clone(), symbolic_range.clone()], opts);
    let idx_buf = UOp::index()
        .buffer(bufferized)
        .indices(vec![concrete_range, symbolic_range])
        .call()
        .expect("Failed to create INDEX");

    // Apply rewrite
    let rewritten = graph_rewrite(&matcher, idx_buf, &mut config);
    drop(rewritten);
}

// ============================================================================
// Phase 5.5.2: Complex Graph Structure Tests
// ============================================================================

/// Test diamond pattern in computation graph.
///
/// A value used multiple times should be handled correctly.
#[test]
fn test_complex_diamond_pattern() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let buffer = create_test_buffer(40, DType::Float32, 1);
    let range = ctx.new_range(&SInt::Const(10), AxisType::Loop);

    // Create diamond: indexed → mul1, mul2 → add
    let indexed =
        UOp::index().buffer(buffer.clone()).indices(vec![range.clone()]).call().expect("Failed to create INDEX");
    let two = UOp::native_const(2.0f32);
    let three = UOp::native_const(3.0f32);

    let mul1 = indexed.try_mul(&two).expect("Failed to create MUL");
    let mul2 = indexed.try_mul(&three).expect("Failed to create MUL");
    let add = mul1.try_add(&mul2).expect("Failed to create ADD");

    // Bufferize and index
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(add, vec![range.clone()], vec![range], opts);

    // Apply rewrite - should handle diamond correctly
    let rewritten = graph_rewrite(&matcher, idx_buf, &mut config);
    drop(rewritten);
}

/// Test deeply nested computation chain.
///
/// Long chains should be optimized without stack overflow.
#[test]
fn test_complex_deep_computation_chain() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let buffer = create_test_buffer(40, DType::Float32, 1);
    let range = ctx.new_range(&SInt::Const(10), AxisType::Loop);

    // Create chain: INDEX(buf) → +1 → +2 → +3 → +4 → +5
    let indexed = UOp::index().buffer(buffer).indices(vec![range.clone()]).call().expect("Failed to create INDEX");
    let mut current = indexed;

    for i in 1..=5 {
        let const_val = UOp::native_const(i as f32);
        current = current.try_add(&const_val).expect("Failed to create ADD");
    }

    // Bufferize the chain
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(current, vec![range.clone()], vec![range], opts);

    // Apply rewrite
    let rewritten = graph_rewrite(&matcher, idx_buf, &mut config);
    drop(rewritten);
}

/// Test computation with multiple independent buffer accesses.
///
/// Should trigger accessed_buffers heuristic correctly.
///
/// Note: Uses level=0 to disable cheap inline pattern so we can test buffer count heuristic.
#[test]
fn test_complex_multiple_independent_buffers() {
    let mut config = PcontigConfig { level: 0, ..PcontigConfig::default() };
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let range = ctx.new_range(&SInt::Const(10), AxisType::Loop);

    // Create 5 independent buffers
    let mut adds = vec![];
    for i in 0..5 {
        let buf = create_test_buffer(40, DType::Float32, i);
        let idx = UOp::index().buffer(buf).indices(vec![range.clone()]).call().expect("Failed to create INDEX");
        adds.push(idx);
    }

    // Chain additions: ((((a + b) + c) + d) + e)
    let mut compute = adds[0].clone();
    for add in &adds[1..] {
        compute = compute.try_add(add).expect("Failed to create ADD");
    }

    // Bufferize
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, vec![range.clone()], vec![range], opts);

    // Apply rewrite - should block due to >3 buffers
    let rewritten = graph_rewrite(&matcher, idx_buf.clone(), &mut config);
    assert!(Arc::ptr_eq(&rewritten, &idx_buf), "Should block optimization with 5 buffers");
}

/// Test with multiple REDUCE operations.
///
/// Complex reduce patterns should be handled correctly.
#[test]
fn test_complex_multiple_sequential_reduces() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let buffer = create_test_buffer(8000, DType::Float32, 1);

    let loop_range = ctx.new_range(&SInt::Const(10), AxisType::Loop);
    let reduce_range1 = ctx.new_range(&SInt::Const(20), AxisType::Reduce);
    let reduce_range2 = ctx.new_range(&SInt::Const(40), AxisType::Reduce);

    // First reduce
    let indexed1 = UOp::index()
        .buffer(buffer.clone())
        .indices(vec![loop_range.clone(), reduce_range1.clone()])
        .call()
        .expect("Failed to create INDEX");
    let reduce1 = UOp::new(
        Op::Reduce { src: indexed1, ranges: vec![reduce_range1].into(), reduce_op: morok_ir::ReduceOp::Add },
        DType::Float32,
    );

    // Second reduce (conceptually, though this doesn't make practical sense)
    let indexed2 = UOp::index()
        .buffer(buffer)
        .indices(vec![loop_range.clone(), reduce_range2.clone()])
        .call()
        .expect("Failed to create INDEX");
    let reduce2 = UOp::new(
        Op::Reduce { src: indexed2, ranges: vec![reduce_range2].into(), reduce_op: morok_ir::ReduceOp::Max },
        DType::Float32,
    );

    // Combine reduces
    let combined = reduce1.try_add(&reduce2).expect("Failed to create ADD");

    // Bufferize
    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let bufferized = UOp::bufferize(combined, vec![loop_range.clone()], opts);
    let idx_buf = UOp::index().buffer(bufferized).indices(vec![loop_range]).call().expect("Failed to create INDEX");

    // Apply rewrite
    let rewritten = graph_rewrite(&matcher, idx_buf, &mut config);
    drop(rewritten);
}

// ============================================================================
// Phase 5.5.3: Boundary & Error Condition Tests
// ============================================================================

/// Test with very large buffer sizes.
///
/// Should handle large numbers without overflow.
#[test]
fn test_boundary_very_large_buffer() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();

    // Large range: 10000 elements * 4 bytes = 40000 bytes
    let large_range = ctx.new_range(&SInt::Const(10000), AxisType::Loop);
    let buffer = create_test_buffer(40000, DType::Float32, 1);

    let indexed =
        UOp::index().buffer(buffer).indices(vec![large_range.clone()]).call().expect("Failed to create INDEX");
    let one = UOp::native_const(1.0f32);
    let compute = indexed.try_add(&one).expect("Failed to create ADD");

    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, vec![large_range.clone()], vec![large_range], opts);

    // Apply rewrite
    let rewritten = graph_rewrite(&matcher, idx_buf, &mut config);
    drop(rewritten);
}

/// Test with size-1 dimensions.
///
/// Degenerate cases should be handled correctly.
#[test]
fn test_boundary_size_one_dimension() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();
    let range1 = ctx.new_range(&SInt::Const(1), AxisType::Loop);
    let buffer = create_test_buffer(4, DType::Float32, 1);

    let indexed = UOp::index().buffer(buffer).indices(vec![range1.clone()]).call().expect("Failed to create INDEX");
    let two = UOp::native_const(2.0f32);
    let compute = indexed.try_mul(&two).expect("Failed to create MUL");

    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, vec![range1.clone()], vec![range1], opts);

    // Apply rewrite
    let rewritten = graph_rewrite(&matcher, idx_buf, &mut config);
    drop(rewritten);
}

/// Test exact threshold boundary conditions.
///
/// Values exactly at thresholds should behave correctly.
#[test]
fn test_boundary_exact_threshold_values() {
    // Test exact ratio threshold (10.0)
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    let mut ctx = IndexingContext::new();

    // Create buffer where ratio will be exactly 10.0
    // input: 40 bytes, output: 409 bytes → ratio = (409+1)/(40+1) = 410/41 = 10.0
    let buffer = create_test_buffer(40, DType::Float32, 1);
    let input_range = ctx.new_range(&SInt::Const(10), AxisType::Loop);
    let output_range = ctx.new_range(&SInt::Const(102), AxisType::Loop); // 102 * 4 = 408 bytes

    let indexed = UOp::index().buffer(buffer).indices(vec![input_range]).call().expect("Failed to create INDEX");
    let one = UOp::native_const(1.0f32);
    let compute = indexed.try_add(&one).expect("Failed to create ADD");

    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let idx_buf = create_index_bufferize(compute, vec![output_range.clone()], vec![output_range], opts);

    // Apply rewrite - at threshold should optimize (ratio >= 10.0)
    let rewritten = graph_rewrite(&matcher, idx_buf.clone(), &mut config);
    // The actual behavior depends on exact ratio calculation
    drop(rewritten);
}

/// Test with empty/trivial computations.
///
/// Minimal graphs should not cause issues.
#[test]
fn test_boundary_minimal_computation() {
    let mut config = PcontigConfig::default();
    let matcher = buffer_removal_with_pcontig();

    // Just a constant
    let const_val = UOp::native_const(42.0f32);

    let mut ctx = IndexingContext::new();
    let range = ctx.new_range(&SInt::Const(10), AxisType::Loop);

    let opts = BufferizeOpts { device: None, addrspace: AddrSpace::Global, removable: true };
    let bufferized = UOp::bufferize(const_val, vec![range], opts);

    // Apply rewrite - Pattern 1 should inline const
    let rewritten = graph_rewrite(&matcher, bufferized.clone(), &mut config);
    assert!(!Arc::ptr_eq(&rewritten, &bufferized), "Const should be inlined via Pattern 1");
}