miden-processor 1.0.0-rc.0

//! Memory-chiplet bus tests.
//!
//! Exercises stack-issued memory opcodes (`MStoreW`, `MLoadW`, `MLoad`, `MStore`, `MStream`)
//! plus the `AeadStream` plaintext read / ciphertext write rows, and verifies the chiplet-requests
//! / chiplet-responses bus pair.
//!
//! For each stack-level memory op the test registers an expected `-1` push of a
//! [`MemoryMsg`] (the "request" side). For each memory chiplet row the test registers an
//! expected `+1` push of a [`MemoryResponseMsg`] (the "response" side). The subset matcher in
//! `lookup_harness` is column-blind, so a `(mult, denom)` pair on the response side pairs up
//! with a matching `(-mult, denom)` on the request side regardless of which aux column the
//! framework routes them onto.
//!
//! # Scope
//!
//! Coverage is limited to the stack-only memory ops above plus `AeadStream`. The DYN,
//! DYNCALL, CALL-FMP-write, and PIPE memory-request paths in
//! `air/src/constraints/lookup/buses/chiplet_requests.rs` are deferred to integration tests —
//! each is heavyweight to set up and small in algebraic surface. A bug in those paths would
//! escape this module.
//!
//! The programs run at ctx = 0 throughout (no CALL/SYSCALL), so a request/response bug that
//! mismatches stack-side `ctx` vs chiplet-side `mem_ctx` is not caught here.

use alloc::vec::Vec;

use miden_air::{
    logup::{MemoryMsg, MemoryResponseMsg},
    trace::MainTrace,
};
use miden_core::{
    Felt, ONE, ZERO,
    chiplets::eidos_compression,
    operations::{Operation, opcodes},
};

use super::super::{
    build_trace_from_ops,
    lookup_harness::{Expectations, InteractionLog},
};
use crate::RowIndex;

const FOUR: Felt = Felt::new_unchecked(4);
const TWO: Felt = Felt::new_unchecked(2);

/// Covers `MStoreW`, `MLoad`, `MLoadW`, `MStore`, `MStream` — every memory opcode issuable
/// directly from the stack — asserting the chiplet-bus request/response pair fires at every
/// memory row.
#[test]
fn memory_chiplet_bus_request_response_pairs() {
    let stack = [0, 1, 2, 3, 4];
    let operations = vec![
        Operation::MStoreW, // store [1, 2, 3, 4] at addr 0
        Operation::Drop,
        Operation::Drop,
        Operation::Drop,
        Operation::Drop,
        Operation::MLoad,      // read first element at addr 0
        Operation::MovDn5,     // reshape stack
        Operation::MLoadW,     // load word from addr 0
        Operation::Push(ONE),  // value = 1
        Operation::Push(FOUR), // addr = 4
        Operation::MStore,     // store 1 at addr 4
        Operation::Drop,
        Operation::MStream, // two-word read starting at stack[12]
    ];
    let trace = build_trace_from_ops(operations, &stack);
    let log = InteractionLog::new(&trace);
    let main = trace.main_trace();

    let mut exp = Expectations::new(&log);

    // ---- Request side: stack rows emit `-1 × MemoryMsg` when their opcode is a memory op.
    //
    // We also count the number of `remove` calls and assert it matches the expected total at
    // the end. Without this, a bug that stopped the emitter entirely would pass vacuously:
    // the request-opcode loop iterates the trace, so no memory rows → no expectations →
    // trivial subset match.
    let mut request_exps_added = 0usize;
    for row in 0..main.core_height() {
        let idx = RowIndex::from(row);
        let ctx = main.ctx(idx);
        let clk = main.clk(idx);
        let next = RowIndex::from(row + 1);
        let op = main.get_op_code(idx).as_canonical_u64();
        if op == opcodes::MLOAD as u64 {
            let addr = main.stack_element(0, idx);
            let value = main.stack_element(0, next);
            exp.remove(row, &MemoryMsg::read_element(ctx, addr, clk, value));
            request_exps_added += 1;
        } else if op == opcodes::MSTORE as u64 {
            let addr = main.stack_element(0, idx);
            let value = main.stack_element(1, idx);
            exp.remove(row, &MemoryMsg::write_element(ctx, addr, clk, value));
            request_exps_added += 1;
        } else if op == opcodes::MLOADW as u64 {
            let addr = main.stack_element(0, idx);
            let word = next_word(main, next, 0);
            exp.remove(row, &MemoryMsg::read_word(ctx, addr, clk, word));
            request_exps_added += 1;
        } else if op == opcodes::MSTOREW as u64 {
            let addr = main.stack_element(0, idx);
            let word = [
                main.stack_element(1, idx),
                main.stack_element(2, idx),
                main.stack_element(3, idx),
                main.stack_element(4, idx),
            ];
            exp.remove(row, &MemoryMsg::write_word(ctx, addr, clk, word));
            request_exps_added += 1;
        } else if op == opcodes::MSTREAM as u64 {
            let base = main.stack_element(12, idx);
            let word0 = next_word(main, next, 0);
            let word1 = next_word(main, next, 4);
            exp.remove(row, &MemoryMsg::read_word(ctx, base, clk, word0));
            exp.remove(row, &MemoryMsg::read_word(ctx, base + FOUR, clk, word1));
            request_exps_added += 2;
        }
    }
    // 5 stack opcodes (MStoreW, MLoad, MLoadW, MStore, MStream) + 1 extra for MStream's 2nd read.
    assert_eq!(request_exps_added, 6, "expected 6 memory request expectations");

    // ---- Response side: every memory chiplet row emits `+1 × MemoryResponseMsg`.
    let mut mem_rows_seen = 0usize;
    for row in 0..main.chiplets_height() {
        let idx = RowIndex::from(row);
        if !main.is_memory_row(idx) {
            continue;
        }
        mem_rows_seen += 1;

        let mem = main.chiplet_cols(idx).memory();
        let is_read = mem.is_read;
        let is_word = mem.is_word;
        let mem_ctx = main.chiplet_memory_ctx(idx);
        let word_addr = main.chiplet_memory_word(idx);
        let idx0 = main.chiplet_memory_idx0(idx);
        let idx1 = main.chiplet_memory_idx1(idx);
        let addr = word_addr + idx1.double() + idx0;
        let mem_clk = main.chiplet_memory_clk(idx);
        let word = [
            main.chiplet_memory_value_0(idx),
            main.chiplet_memory_value_1(idx),
            main.chiplet_memory_value_2(idx),
            main.chiplet_memory_value_3(idx),
        ];
        // `element` is ignored by `MemoryResponseMsg::encode` when `is_word = 1`, so on
        // word-access rows the fallback `ZERO` is harmless. `element_idx` uses `u64`
        // arithmetic (native `usize` indexing) while `addr` above uses felt arithmetic —
        // same math, different domain required by the consumer.
        let element = if is_word == ZERO {
            let element_idx = (idx1.as_canonical_u64() * 2 + idx0.as_canonical_u64()) as usize;
            word[element_idx]
        } else {
            ZERO
        };

        exp.add(
            row,
            &MemoryResponseMsg {
                is_read,
                ctx: mem_ctx,
                addr,
                clk: mem_clk,
                is_word,
                element,
                word,
            },
        );
    }
    // 6 memory operations: MStoreW, MLoad, MLoadW, MStore, MStream (2 rows).
    assert_eq!(mem_rows_seen, 6, "expected 6 memory chiplet rows");

    log.assert_contains(&exp);
}

/// Verifies that `AeadStream` emits its source reads and ciphertext writes.
#[test]
fn aead_stream_emits_memory_requests() {
    // Stack layout: [K_CTR(4), counter, src_ptr, dst_ptr, remaining, tail(8)].
    let stack = [
        1, 2, 3, 4, // K_CTR
        0, // counter
        0, // src_ptr
        8, // dst_ptr
        1, // remaining
        0, 0, 0, 0, 0, 0, 0, 0, // tail
    ];

    let trace = build_trace_from_ops(vec![Operation::CryptoStream], &stack);
    let log = InteractionLog::new(&trace);
    let main = trace.main_trace();

    let mut exp = Expectations::new(&log);

    let aead_rows: Vec<_> = (0..main.core_height())
        .filter(|&row| main.chiplet_aead_stream_active(RowIndex::from(row)) == ONE)
        .collect();
    assert_eq!(aead_rows.len(), 16, "expected two 8-row stream entries");

    // AeadStream runs at cycle 1 (cycle 0 is SPAN), ctx = 0. Source memory is uninitialized, so
    // both plaintext reads return zeros.
    let input_state = [
        ZERO,
        ZERO,
        ZERO,
        ZERO,
        ZERO,
        ZERO,
        ZERO,
        ZERO,
        Felt::new_unchecked(1),
        Felt::new_unchecked(2),
        Felt::new_unchecked(3),
        Felt::new_unchecked(4),
    ];
    let keystream = eidos_compression::compress_raw_xof_lanes(&input_state);
    let ciphertext: [Felt; 16] = keystream.map(Felt::from_u32);
    let zero_word = [ZERO, ZERO, ZERO, ZERO];
    let cipher_words: [[Felt; 4]; 4] =
        core::array::from_fn(|i| core::array::from_fn(|j| ciphertext[4 * i + j]));

    let mut request_exps_added = 0usize;
    exp.remove(aead_rows[0], &MemoryMsg::read_word(ZERO, ZERO, ONE, zero_word));
    request_exps_added += 1;
    exp.remove(aead_rows[4], &MemoryMsg::read_word(ZERO, ZERO, ONE, zero_word));
    request_exps_added += 1;
    exp.remove(
        aead_rows[3],
        &MemoryMsg::write_word(ZERO, Felt::new_unchecked(8), ONE, cipher_words[0]),
    );
    request_exps_added += 1;
    exp.remove(
        aead_rows[7],
        &MemoryMsg::write_word(ZERO, Felt::new_unchecked(12), ONE, cipher_words[1]),
    );
    request_exps_added += 1;
    exp.remove(aead_rows[8], &MemoryMsg::read_word(ZERO, FOUR, ONE, zero_word));
    request_exps_added += 1;
    exp.remove(aead_rows[12], &MemoryMsg::read_word(ZERO, FOUR, ONE, zero_word));
    request_exps_added += 1;
    exp.remove(
        aead_rows[11],
        &MemoryMsg::write_word(ZERO, Felt::new_unchecked(16), ONE, cipher_words[2]),
    );
    request_exps_added += 1;
    exp.remove(
        aead_rows[15],
        &MemoryMsg::write_word(ZERO, Felt::new_unchecked(20), ONE, cipher_words[3]),
    );
    request_exps_added += 1;

    assert_eq!(request_exps_added, 8, "expected 8 AeadStream memory expectations");

    log.assert_contains(&exp);
}

/// Verifies that `HornerBase` emits one zero-padded word read on the chiplet-requests bus.
#[test]
fn hornerbase_emits_one_memory_request() {
    // Use a point for which helpers[2] and helpers[3] are nonzero. This makes either helper
    // distinguishable from the zero padding in the memory request.
    let stack = [
        1, 2, 3, 4, 5, 6, 7, 8, // coeffs[0..8]
        0, 0, 0, 0, 0, // pad (5 slots)
        0, // alpha_ptr (stack[13])
        0, 0, // acc_low, acc_high
    ];

    let operations = vec![
        Operation::Push(ZERO),
        Operation::Push(ZERO),
        Operation::Push(TWO),
        Operation::Push(ONE),
        Operation::Push(ZERO),
        Operation::MStoreW,
        Operation::Drop,
        Operation::Drop,
        Operation::Drop,
        Operation::Drop,
        Operation::HornerBase,
    ];
    let trace = build_trace_from_ops(operations, &stack);
    const HORNER_CLK: u32 = 11;
    let helpers = trace.get_user_op_helpers_at(HORNER_CLK);
    assert_ne!(helpers[2], ZERO);
    assert_ne!(helpers[3], ZERO);

    let log = InteractionLog::new(&trace);

    let mut exp = Expectations::new(&log);

    // HornerBase runs at cycle 11 (cycle 0 is SPAN).
    exp.remove(
        HORNER_CLK as usize,
        &MemoryMsg::read_word(ZERO, ZERO, Felt::from_u32(HORNER_CLK), [ONE, TWO, ZERO, ZERO]),
    );

    log.assert_contains(&exp);
}

/// Verifies that `HornerExt`'s word read lands on the chiplet-requests bus.
#[test]
fn hornerext_emits_one_memory_request() {
    // HornerExt stack layout: [s0_lo, s0_hi, ..., s3_lo, s3_hi, _, _, _, _, _, alpha_ptr,
    // acc_low, acc_high]
    let stack = [
        1, 2, 3, 4, 5, 6, 7, 8, // 4 extension coeffs (s0..s3)
        0, 0, 0, 0, 0, // pad (5 slots)
        0, // alpha_ptr (stack[13])
        0, 0, // acc_low, acc_high
    ];

    let trace = build_trace_from_ops(vec![Operation::HornerExt], &stack);
    let log = InteractionLog::new(&trace);

    let mut exp = Expectations::new(&log);

    const ROW: usize = 1;
    exp.remove(ROW, &MemoryMsg::read_word(ZERO, ZERO, ONE, [ZERO; 4]));

    log.assert_contains(&exp);
}

fn next_word(main: &MainTrace, next: RowIndex, start: usize) -> [Felt; 4] {
    [
        main.stack_element(start, next),
        main.stack_element(start + 1, next),
        main.stack_element(start + 2, next),
        main.stack_element(start + 3, next),
    ]
}