miden-processor 1.0.0-rc.0

use alloc::vec::Vec;

use miden_air::trace::{
    CHIPLETS_CLK_COL, CHIPLETS_MODE_COL, CHIPLETS_WIDTH, TRACE_WIDTH,
    chiplets::{
        KERNEL_ROM_TRACE_WIDTH, NUM_BITWISE_SELECTORS, NUM_KERNEL_ROM_SELECTORS,
        NUM_MEMORY_SELECTORS,
        bitwise::{self, BITWISE_XOR, OP_CYCLE_LEN},
        hasher::{CONTROLLER_ROWS_PER_HASHER_OP, CONTROLLER_TRACE_ALIGNMENT, LINEAR_HASH},
        memory,
    },
};
use miden_core::{
    Felt, ONE, Word, ZERO,
    mast::{BasicBlockNodeBuilder, CallNodeBuilder, MastForest},
    program::{Program, StackInputs},
};

use crate::{
    AdviceInputs, DefaultHost, ExecutionOptions, FastProcessor, KernelDescriptor,
    operation::Operation,
};

type ChipletsTrace = [Vec<Felt>; CHIPLETS_WIDTH];

// HASHER TRACE LENGTH HELPERS
// ================================================================================================

/// Computes the chiplets-side hasher trace length from the number of controller rows.
fn hasher_trace_len(controller_rows: usize) -> usize {
    controller_rows.next_multiple_of(CONTROLLER_TRACE_ALIGNMENT)
}

// TESTS
// ================================================================================================

#[test]
fn hasher_chiplet_trace() {
    // --- single hasher compression with no stack manipulation ---
    // The program is a single basic block containing Compress.
    // This produces:
    //   - 1 span hash controller row
    //   - 1 COMPRESS controller row
    // Total: 2 controller rows padded to the chiplet alignment boundary.
    let stack = [2, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0];
    let operations = vec![Operation::Compress];
    let (chiplets_trace, _trace_len) = build_trace(&stack, operations, KernelDescriptor::default());

    let controller_rows = 2 * CONTROLLER_ROWS_PER_HASHER_OP; // span hash + Compress
    let hasher_len = hasher_trace_len(controller_rows);
    assert_eq!(hasher_len, CONTROLLER_TRACE_ALIGNMENT);

    validate_hasher_trace(&chiplets_trace, hasher_len, controller_rows);
}

#[test]
fn bitwise_chiplet_trace() {
    // --- single bitwise operation with no stack manipulation ---
    // This produces: 1 span hash controller row, then 1 bitwise row.
    let stack = [4, 8];
    let operations = vec![Operation::U32xor];
    let (chiplets_trace, _trace_len) = build_trace(&stack, operations, KernelDescriptor::default());

    let controller_rows = CONTROLLER_ROWS_PER_HASHER_OP; // span hash only
    let hasher_len = hasher_trace_len(controller_rows);
    assert_eq!(hasher_len, CONTROLLER_TRACE_ALIGNMENT);

    let bitwise_start = hasher_len;
    let bitwise_end = bitwise_start + OP_CYCLE_LEN;
    validate_bitwise_trace(&chiplets_trace, bitwise_start, bitwise_end);
}

#[test]
fn memory_chiplet_trace() {
    // --- single memory operation with no stack manipulation ---
    // This produces: 1 span hash, then 1 memory row.
    let addr = Felt::from_u32(4);
    let stack = [1, 2, 3, 4];
    let operations = vec![Operation::Push(addr), Operation::MStoreW];
    let (chiplets_trace, _trace_len) = build_trace(&stack, operations, KernelDescriptor::default());

    let controller_rows = CONTROLLER_ROWS_PER_HASHER_OP;
    let hasher_len = hasher_trace_len(controller_rows);
    assert_eq!(hasher_len, CONTROLLER_TRACE_ALIGNMENT);

    let memory_start = hasher_len;
    validate_memory_trace(&chiplets_trace, memory_start, memory_start + 1);
}

#[test]
fn stacked_chiplet_trace() {
    // --- operations in hasher, bitwise, and memory processors ---
    // Operations: U32xor, Push(0), MStoreW, Compress
    // This produces:
    //   - 1 span hash controller row for the basic block
    //   - 1 COMPRESS controller row
    // Total hasher: 2 controller rows padded to the chiplet alignment boundary.
    // Then: 1 bitwise row (U32xor), then 1 memory row (MStoreW)
    let stack = [8, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 1];
    let ops = vec![
        Operation::U32xor,
        Operation::Push(ZERO),
        Operation::MStoreW,
        Operation::Compress,
    ];
    let kernel = build_kernel();
    let (chiplets_trace, _trace_len) = build_trace(&stack, ops, kernel);

    let controller_rows = 2 * CONTROLLER_ROWS_PER_HASHER_OP; // span hash + Compress
    let hasher_len = hasher_trace_len(controller_rows);
    assert_eq!(hasher_len, CONTROLLER_TRACE_ALIGNMENT);

    // Validate hasher region
    validate_hasher_trace(&chiplets_trace, hasher_len, controller_rows);

    // Bitwise starts right after hasher
    let bitwise_start = hasher_len;
    let bitwise_end = bitwise_start + OP_CYCLE_LEN;
    validate_bitwise_trace(&chiplets_trace, bitwise_start, bitwise_end);

    // Memory starts right after bitwise
    let memory_start = bitwise_end;
    validate_memory_trace(&chiplets_trace, memory_start, memory_start + 1);

    // After memory comes kernel ROM (2 entries from build_kernel) then padding
    let kernel_rom_start = memory_start + 1;
    let kernel_rom_end = kernel_rom_start + 2; // 2 kernel procedures
    validate_kernel_rom_trace(&chiplets_trace, kernel_rom_start, kernel_rom_end);

    // Padding fills the remainder
    let padding_start = kernel_rom_end;
    let trace_rows = chiplets_trace[0].len();
    validate_padding(&chiplets_trace, padding_start, trace_rows);
}

#[test]
fn regression_trace_build_does_not_panic_when_first_memory_access_clk_is_zero() {
    let processor = FastProcessor::new(StackInputs::default());
    let mut host = DefaultHost::default();

    // A CALL entrypoint records the callee frame pointer write before the processor clock is
    // incremented, so the first memory access is captured at clk = 0.
    let program = {
        let mut forest = MastForest::new();

        let callee = BasicBlockNodeBuilder::new(vec![Operation::Noop])
            .add_to_forest(&mut forest)
            .unwrap();
        forest.make_root(callee);

        let entry = CallNodeBuilder::new(callee).add_to_forest(&mut forest).unwrap();
        forest.make_root(entry);

        Program::with_kernel(forest.into(), entry, KernelDescriptor::default())
    };

    let execution_witness = processor.execute_for_proving_sync(&program, &mut host).unwrap();
    let (vm_witness, _) = execution_witness.into_parts();

    let _trace = crate::trace::build_trace(vm_witness).unwrap();
}

// HELPER FUNCTIONS
// ================================================================================================

fn build_kernel() -> KernelDescriptor {
    let proc_hash1 = Word::from([1_u32, 0, 1, 0]);
    let proc_hash2 = Word::from([1_u32, 1, 1, 1]);
    KernelDescriptor::new(&[proc_hash1, proc_hash2]).unwrap()
}

fn build_trace(
    stack_inputs: &[u64],
    operations: Vec<Operation>,
    kernel: KernelDescriptor,
) -> (ChipletsTrace, usize) {
    let stack_inputs: Vec<Felt> = stack_inputs.iter().map(|v| Felt::new_unchecked(*v)).collect();
    let processor = FastProcessor::new_with_options(
        StackInputs::new(&stack_inputs).unwrap(),
        AdviceInputs::default(),
        ExecutionOptions::default().with_core_trace_fragment_size(1 << 10).unwrap(),
    )
    .expect("processor advice inputs should fit advice map limits");

    let mut host = DefaultHost::default();
    let program = {
        let mut mast_forest = MastForest::new();
        let basic_block_id =
            BasicBlockNodeBuilder::new(operations).add_to_forest(&mut mast_forest).unwrap();
        mast_forest.make_root(basic_block_id);
        Program::with_kernel(mast_forest.into(), basic_block_id, kernel)
    };

    let execution_witness = processor.execute_for_proving_sync(&program, &mut host).unwrap();
    let (vm_witness, _) = execution_witness.into_parts();
    let trace = crate::trace::build_trace(vm_witness).unwrap();

    let trace_len = trace.get_trace_len();
    (
        trace
            .get_column_range((TRACE_WIDTH - CHIPLETS_WIDTH)..TRACE_WIDTH)
            .try_into()
            .expect("failed to convert vector to array"),
        trace_len,
    )
}

// VALIDATION FUNCTIONS
// ================================================================================================

/// Validates the hasher region of the chiplets trace.
///
/// Checks:
/// - s_ctrl (column 0) = 1 on controller rows
/// - Controller rows have the correct selectors for operation type and boundary flags
/// - Padding rows: selectors [0, 1, 0], non-selector columns are zero
fn validate_hasher_trace(trace: &ChipletsTrace, expected_len: usize, controller_rows: usize) {
    // Column indices within chiplets trace.
    // Column 0 = s_ctrl. Hasher internal columns start at column 1.
    let s0_col = 1; // hasher selector s0
    let s1_col = 2; // hasher selector s1
    let s2_col = 3; // hasher selector s2

    let controller_padded = hasher_trace_len(controller_rows);

    assert_eq!(expected_len, controller_padded);

    // Controller rows (including padding): s_ctrl=1.
    for row in 0..controller_padded {
        assert_eq!(trace[0][row], ONE, "s_ctrl should be 1 for controller row {row}");
    }

    // --- Check controller rows ---
    // Each controller row carries one compression request.
    for row in 0..controller_rows {
        assert_eq!(
            trace[s0_col][row], LINEAR_HASH[0],
            "controller row {row}: s0 should be {} (LINEAR_HASH)",
            LINEAR_HASH[0]
        );
    }

    for row in controller_rows..controller_padded {
        assert_eq!(trace[s0_col][row], ZERO, "padding row {row}: s0 should be 0");
        assert_eq!(trace[s1_col][row], ONE, "padding row {row}: s1 should be 1");
        assert_eq!(trace[s2_col][row], ZERO, "padding row {row}: s2 should be 0");

        // Non-selector hasher columns should be zero on padding rows, except the shared
        // controller discriminator. Controller constraints bind this cell to
        // The trailing column (chip_clk, CHIPLETS_WIDTH - 1) is the chiplet-trace row counter
        // and is non-zero on every row by design; see `air/src/constraints/chiplets/chip_clk.rs`.
        for col in 4..CHIPLETS_WIDTH - 1 {
            if col == CHIPLETS_MODE_COL {
                assert_eq!(trace[col][row], ONE, "padding row {row}: mode cell should be 1");
                continue;
            }
            assert_eq!(trace[col][row], ZERO, "padding row {row}, col {col} should be zero");
        }
    }
}

/// Validates the bitwise region of the chiplets trace.
///
/// Checks:
/// - Chiplet selectors: s_ctrl=0, s1=0, stream_mode=0
/// - Bitwise operation selector = XOR
/// - Columns beyond bitwise trace width + selectors are zero
fn validate_bitwise_trace(trace: &ChipletsTrace, start: usize, end: usize) {
    // Bitwise uses NUM_BITWISE_SELECTORS (2) chiplet selector columns + bitwise::TRACE_WIDTH (13)
    // internal columns = 15 columns total. Columns 15..CHIPLETS_WIDTH should be zero.
    let bitwise_used_cols = NUM_BITWISE_SELECTORS + bitwise::TRACE_WIDTH;

    for row in start..end {
        // Chiplet selectors: s_ctrl=0, s1=0 (active via virtual s0 * !s1)
        assert_eq!(ZERO, trace[0][row], "bitwise s_ctrl at row {row}");
        assert_eq!(ZERO, trace[1][row], "bitwise s1 at row {row}");

        // Internal bitwise operation selector (XOR)
        assert_eq!(BITWISE_XOR, trace[NUM_BITWISE_SELECTORS][row], "bitwise op at row {row}");

        // Columns beyond bitwise trace should be zero (chip_clk excluded; see chip_clk.rs).
        for col in bitwise_used_cols..CHIPLETS_WIDTH - 1 {
            assert_eq!(
                trace[col][row], ZERO,
                "bitwise padding col {col} at row {row} should be zero"
            );
        }
    }
}

/// Validates the memory region of the chiplets trace.
///
/// Checks:
/// - Chiplet selectors: s_ctrl=0, s1=1, s2=0, stream_mode=0
/// - Column beyond memory trace width + selectors is zero
fn validate_memory_trace(trace: &ChipletsTrace, start: usize, end: usize) {
    // Memory uses NUM_MEMORY_SELECTORS (3) chiplet selector columns + memory::TRACE_WIDTH (17)
    // internal columns = 20 columns total. Column 20 should be zero.
    let memory_used_cols = NUM_MEMORY_SELECTORS + memory::TRACE_WIDTH;

    for row in start..end {
        // Chiplet selectors: s_ctrl=0, s1=1, s2=0 (active via virtual s0 * s1 * !s2)
        assert_eq!(ZERO, trace[0][row], "memory s_ctrl at row {row}");
        assert_eq!(ONE, trace[1][row], "memory s1 at row {row}");
        assert_eq!(ZERO, trace[2][row], "memory s2 at row {row}");

        // Columns beyond memory trace should be zero (chip_clk excluded; see chip_clk.rs).
        for col in memory_used_cols..CHIPLETS_WIDTH - 1 {
            assert_eq!(
                trace[col][row], ZERO,
                "memory padding col {col} at row {row} should be zero"
            );
        }
    }
}

/// Validates the kernel ROM region of the chiplets trace.
///
/// Checks:
/// - Chiplet selectors: s_ctrl=0, s1=1, s2=1, s3=1, s4=0, stream_mode=0
/// - Columns beyond kernel ROM trace width + selectors are zero
fn validate_kernel_rom_trace(trace: &ChipletsTrace, start: usize, end: usize) {
    // Kernel ROM uses NUM_KERNEL_ROM_SELECTORS (5) chiplet selector columns +
    // KERNEL_ROM_TRACE_WIDTH (5) internal columns = 10 columns total.
    let kernel_rom_used_cols = NUM_KERNEL_ROM_SELECTORS + KERNEL_ROM_TRACE_WIDTH;

    for row in start..end {
        // Chiplet selectors: s_ctrl=0, s1=1, s2=1, s3=1, s4=0
        // (active via virtual s0 * s1 * s2 * s3 * !s4)
        assert_eq!(ZERO, trace[0][row], "kernel_rom s_ctrl at row {row}");
        assert_eq!(ONE, trace[1][row], "kernel_rom s1 at row {row}");
        assert_eq!(ONE, trace[2][row], "kernel_rom s2 at row {row}");
        assert_eq!(ONE, trace[3][row], "kernel_rom s3 at row {row}");
        assert_eq!(ZERO, trace[4][row], "kernel_rom s4 at row {row}");

        // Columns beyond kernel ROM trace should be zero (chip_clk excluded; see chip_clk.rs).
        for col in kernel_rom_used_cols..CHIPLETS_WIDTH - 1 {
            assert_eq!(
                trace[col][row], ZERO,
                "kernel_rom padding col {col} at row {row} should be zero"
            );
        }
    }
}

/// Validates the padding region at the end of the chiplets trace.
///
/// Checks:
/// - s_ctrl (column 0) = 0, s1..s4 (columns 1-4) = 1
/// - payload columns are zero
fn validate_padding(trace: &ChipletsTrace, start: usize, end: usize) {
    for row in start..end {
        // s_ctrl = 0 on padding rows
        assert_eq!(ZERO, trace[0][row], "padding s_ctrl at row {row}");
        // s1..s4 = 1 on padding rows
        for col in 1..5 {
            assert_eq!(ONE, trace[col][row], "padding s{col} at row {row}");
        }
        // chip_clk at CHIPLETS_WIDTH - 1 is non-zero by design.
        for col in 5..CHIPLETS_WIDTH - 1 {
            assert_eq!(ZERO, trace[col][row], "padding data col {col} at row {row} should be zero");
        }
        assert_ne!(ZERO, trace[CHIPLETS_CLK_COL][row], "padding chip_clk at row {row}");
    }
}