miden-processor 1.0.0-rc.0

use alloc::{boxed::Box, vec, vec::Vec};
use core::{borrow::BorrowMut, mem::size_of};

use miden_air::{
    AeadStreamCols, BitwiseCols,
    trace::{
        CHIPLETS_STREAM_MODE_COL,
        and8_lookup::{BYTE_LOOKUP_COUNT_LEN, BytePairRelation, byte_pair_count_index},
        chiplets::{
            bitwise::{BITWISE_AND, BITWISE_XOR, OP_CYCLE_LEN, TRACE_WIDTH},
            hasher::CONTROLLER_TRACE_ALIGNMENT,
        },
    },
};
use miden_core::{chiplets::eidos_compression, field::Field};

use super::BITWISE_COL_START;
use crate::{Felt, ONE, ZERO, operation::OperationError, trace::ChipletTraceFragment};

#[cfg(test)]
mod tests;

// CONSTANTS
// ================================================================================================

/// Initial capacity, in ops.
const INIT_OPS_CAPACITY: usize = 128;
pub(crate) const AEAD_STREAM_CYCLE_LEN: usize = 8;
const AEAD_STREAM_WIDTH: usize = size_of::<AeadStreamCols<u8>>();
const STREAM_MODE_OFFSET: usize = AEAD_STREAM_WIDTH;
pub(crate) const AEAD_STREAM_FRAGMENT_WIDTH: usize = STREAM_MODE_OFFSET + 1;

const _: () = assert!(BITWISE_COL_START + STREAM_MODE_OFFSET == CHIPLETS_STREAM_MODE_COL);
const _: () = assert!(CONTROLLER_TRACE_ALIGNMENT.is_multiple_of(AEAD_STREAM_CYCLE_LEN));

// BITWISE OPERATION
// ================================================================================================

/// Which bitwise operation a row encodes.
#[derive(Debug, Clone, Copy)]
enum Op {
    And,
    Xor,
}

impl Op {
    fn selector(self) -> Felt {
        match self {
            Self::And => BITWISE_AND,
            Self::Xor => BITWISE_XOR,
        }
    }

    fn apply(self, a: u32, b: u32) -> u32 {
        match self {
            Self::And => a & b,
            Self::Xor => a ^ b,
        }
    }
}

/// A single bitwise operation recorded for later trace materialization.
#[derive(Debug, Clone, Copy)]
struct BitwiseOp {
    op: Op,
    a: u32,
    b: u32,
}

#[derive(Debug, Clone, Copy)]
struct AeadStreamOp {
    ctx: Felt,
    clk: Felt,
    src_ptr: Felt,
    dst_ptr: Felt,
    lane_base: Felt,
    plaintext: [Felt; 4],
    keystream: [Felt; 8],
    ciphertext: [Felt; 8],
}

#[derive(Debug)]
enum Entry {
    Bitwise(BitwiseOp),
    AeadStream(Box<AeadStreamOp>),
}

impl Entry {
    fn row_count(&self) -> usize {
        match self {
            Self::Bitwise(_) => OP_CYCLE_LEN,
            Self::AeadStream(_) => AEAD_STREAM_CYCLE_LEN,
        }
    }

    fn is_aead_stream(&self) -> bool {
        matches!(self, Self::AeadStream(_))
    }
}

// BITWISE
// ================================================================================================

/// Helper for the VM that computes AND and XOR bitwise operations on 32-bit values.
/// It also builds an execution trace of these operations.
///
/// ## Bitwise operation execution trace (AND and XOR)
/// Each operation uses one row. The row stores the operation flag, little-endian bytes of both
/// inputs, and bytewise `a & b` witnesses. Four AND8 lookups prove the byte witnesses. The
/// response bus reconstructs the VM-facing values and derives XOR as `a + b - 2*(a & b)`.
///
/// The layout of the table is illustrated below.
///
///    s     a0    a1    a2    a3    b0    b1    b2    b3    c0    c1    c2    c3
/// |-----+-----+-----+-----+-----+-----+-----+-----+-----+-----+-----+-----+-----|
///
/// In the above, the meaning of the columns is as follows:
/// - `s` selects the bitwise operator: 0 = AND, 1 = XOR.
/// - `a*` and `b*` are little-endian input bytes.
/// - `c*` are bytewise `a & b` witnesses.
#[derive(Debug)]
pub struct Bitwise {
    entries: Vec<Entry>,
    trace_len: usize,
}

impl Bitwise {
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Returns a new [Bitwise] initialized with an empty op log.
    pub fn new() -> Self {
        Self {
            entries: Vec::with_capacity(INIT_OPS_CAPACITY),
            trace_len: 0,
        }
    }

    // PUBLIC ACCESSORS
    // --------------------------------------------------------------------------------------------

    /// Returns length of execution trace required to describe bitwise operations executed on the
    /// VM.
    pub fn trace_len(&self) -> usize {
        self.trace_len
    }

    // TRACE MUTATORS
    // --------------------------------------------------------------------------------------------

    /// Computes a bitwise AND of `a` and `b` and returns the result. We assume that `a` and `b`
    /// are 32-bit values. If that's not the case, the result of the computation is undefined.
    ///
    /// Records the op for later trace generation in [`Self::fill_trace`].
    pub fn u32and(&mut self, a: Felt, b: Felt) -> Result<Felt, OperationError> {
        self.record(Op::And, a, b)
    }

    /// Computes a bitwise XOR of `a` and `b` and returns the result. We assume that `a` and `b`
    /// are 32-bit values. If that's not the case, the result of the computation is undefined.
    ///
    /// Records the op for later trace generation in [`Self::fill_trace`].
    pub fn u32xor(&mut self, a: Felt, b: Felt) -> Result<Felt, OperationError> {
        self.record(Op::Xor, a, b)
    }

    fn record(&mut self, op: Op, a: Felt, b: Felt) -> Result<Felt, OperationError> {
        let a = assert_u32(a)?;
        let b = assert_u32(b)?;
        self.entries.push(Entry::Bitwise(BitwiseOp { op, a, b }));
        self.trace_len += OP_CYCLE_LEN;
        Ok(Felt::from_u32(op.apply(a, b)))
    }

    /// Records one 8-row AEAD stream entry.
    ///
    /// Entries may come from deserialized witnesses. Checks that keystream and ciphertext lanes
    /// are u32 values and that each ciphertext lane is the plaintext lane XOR its keystream lane.
    pub(crate) fn aead_stream(
        &mut self,
        ctx: Felt,
        clk: Felt,
        src_ptr: Felt,
        dst_ptr: Felt,
        lane_base: Felt,
        plaintext: [Felt; 4],
        keystream: [Felt; 8],
        ciphertext: [Felt; 8],
    ) -> Result<(), OperationError> {
        let limbs = plaintext.map(eidos_compression::unpack);
        for (idx, (&ks, &c)) in keystream.iter().zip(&ciphertext).enumerate() {
            let (lo, hi) = limbs[idx / 2];
            let plaintext_limb = if idx % 2 == 0 { lo } else { hi };
            if plaintext_limb ^ assert_u32(ks)? != assert_u32(c)? {
                return Err(OperationError::Internal(
                    "AEAD stream ciphertext does not match plaintext XOR keystream",
                ));
            }
        }

        self.entries.push(Entry::AeadStream(Box::new(AeadStreamOp {
            ctx,
            clk,
            src_ptr,
            dst_ptr,
            lane_base,
            plaintext,
            keystream,
            ciphertext,
        })));
        self.trace_len += AEAD_STREAM_CYCLE_LEN;
        Ok(())
    }

    // EXECUTION TRACE GENERATION
    // --------------------------------------------------------------------------------------------

    /// Fills the provided trace fragment with the row-major trace materialized from the recorded
    /// op log: eight rows per AEAD stream entry followed by one row per normal bitwise op.
    ///
    /// AEAD stream entries are kept in execution order relative to each other and materialized
    /// before the normal bitwise entries. The bitwise trace region starts on an 8-row boundary,
    /// so this stable partition keeps every stream entry aligned with its period-8 constraints.
    /// The interactions emitted by both entry kinds carry their execution metadata rather than a
    /// trace-row address, so their relative placement in this trace is immaterial.
    pub fn fill_trace(self, trace: &mut ChipletTraceFragment) -> Vec<u64> {
        debug_assert_eq!(self.trace_len(), trace.len(), "inconsistent trace lengths");
        debug_assert!(trace.width() >= TRACE_WIDTH, "inconsistent trace widths");
        let has_stream_rows =
            self.entries.iter().any(|entry| matches!(entry, Entry::AeadStream(_)));
        debug_assert!(
            !has_stream_rows || trace.width() >= AEAD_STREAM_FRAGMENT_WIDTH,
            "trace fragment too narrow for AEAD stream rows",
        );

        let row_width = trace.width();
        let mut row_offset = 0;
        let mut and8_counts = vec![0u64; BYTE_LOOKUP_COUNT_LEN];

        let entries = self
            .entries
            .iter()
            .filter(|entry| entry.is_aead_stream())
            .chain(self.entries.iter().filter(|entry| !entry.is_aead_stream()));

        for entry in entries {
            let row_count = entry.row_count();
            if entry.is_aead_stream() {
                debug_assert_eq!(row_offset % AEAD_STREAM_CYCLE_LEN, 0);
            }
            let mut chunk = vec![ZERO; row_width * row_count];
            match entry {
                Entry::Bitwise(op) => {
                    fill_bitwise_chunk(&mut chunk, row_width, *op, &mut and8_counts);
                },
                Entry::AeadStream(op) => {
                    fill_aead_stream_chunk(&mut chunk, row_width, **op, &mut and8_counts);
                },
            }
            trace.copy_rows_into(row_offset, &chunk);
            row_offset += row_count;
        }
        debug_assert_eq!(row_offset, trace.len());

        and8_counts
    }
}

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

// HELPER FUNCTIONS
// --------------------------------------------------------------------------------------------

pub fn assert_u32(value: Felt) -> Result<u32, OperationError> {
    u32::try_from(value.as_canonical_u64())
        .map_err(|_| OperationError::NotU32Values { values: vec![value] })
}

fn fill_bitwise_chunk(
    chunk: &mut [Felt],
    row_width: usize,
    BitwiseOp { op, a, b }: BitwiseOp,
    and8_counts: &mut [u64],
) {
    debug_assert_eq!(chunk.len(), row_width);

    let a_bytes = a.to_le_bytes();
    let b_bytes = b.to_le_bytes();
    let selector = op.selector();

    let row = &mut chunk[..TRACE_WIDTH];
    let cols: &mut BitwiseCols<Felt> = row.borrow_mut();
    cols.op_flag = selector;
    cols.a_bytes = a_bytes.map(Felt::from_u8);
    cols.b_bytes = b_bytes.map(Felt::from_u8);
    cols.and_bytes = core::array::from_fn(|idx| {
        let and = a_bytes[idx] & b_bytes[idx];
        count_and8(and8_counts, a_bytes[idx], b_bytes[idx], and);
        Felt::from_u8(and)
    });
}

fn fill_aead_stream_chunk(
    chunk: &mut [Felt],
    row_width: usize,
    op: AeadStreamOp,
    and8_counts: &mut [u64],
) {
    debug_assert_eq!(chunk.len(), row_width * AEAD_STREAM_CYCLE_LEN);

    let plaintext = op.plaintext;
    let limbs = plaintext.map(eidos_compression::unpack);
    let limb_data: [(u32, Felt, Felt); 8] = core::array::from_fn(|idx| {
        let (lo, hi) = limbs[idx / 2];
        let plaintext_limb = if idx % 2 == 0 { lo } else { hi };
        let ks = op.keystream[idx];
        let ciphertext = op.ciphertext[idx];
        debug_assert_eq!(
            Felt::from_u32(plaintext_limb ^ felt_to_u32(ks)),
            ciphertext,
            "AEAD stream ciphertext mismatch at limb {idx}",
        );
        (plaintext_limb, ks, ciphertext)
    });

    let witnesses = limb_data.map(|(plaintext_limb, ks, ciphertext)| {
        limb_xor_witness(plaintext_limb, felt_to_u32(ks), felt_to_u32(ciphertext))
    });
    for witness in witnesses {
        count_and8_witness(and8_counts, witness.bytes);
    }

    for row_idx in 0..AEAD_STREAM_CYCLE_LEN {
        let row = &mut chunk[row_idx * row_width..(row_idx + 1) * row_width];
        row[STREAM_MODE_OFFSET] = ONE;
    }

    // Each four-row half reads the same plaintext word and handles four consecutive u32 lanes.
    fill_stream_word_pair(
        &mut chunk[..4 * row_width],
        row_width,
        op,
        plaintext,
        0,
        0,
        [limbs[0], limbs[1]],
        &[witnesses[0], witnesses[1], witnesses[2], witnesses[3]],
    );

    fill_stream_word_pair(
        &mut chunk[4 * row_width..8 * row_width],
        row_width,
        op,
        plaintext,
        2,
        4,
        [limbs[2], limbs[3]],
        &[witnesses[4], witnesses[5], witnesses[6], witnesses[7]],
    );
}

fn fill_stream_word_pair(
    chunk: &mut [Felt],
    row_width: usize,
    op: AeadStreamOp,
    plaintext: [Felt; 4],
    plaintext_offset: usize,
    lane_offset: usize,
    limbs: [(u32, u32); 2],
    witnesses: &[LimbXorWitness; 4],
) {
    let lane_base = op.lane_base + Felt::new_unchecked(lane_offset as u64);
    let dst_ptr = op.dst_ptr + Felt::new_unchecked(lane_offset as u64);

    {
        let row = &mut chunk[0..row_width];
        let cols: &mut AeadStreamCols<Felt> = row[..AEAD_STREAM_WIDTH].borrow_mut();
        let cols = cols.read_mut();
        cols.ctx = op.ctx;
        cols.clk = op.clk;
        cols.src_ptr = op.src_ptr;
        cols.lane_base = lane_base;
        cols.plaintext = plaintext;
        cols.bytes = witnesses[0].bytes;
    }

    {
        let row = &mut chunk[row_width..2 * row_width];
        let cols: &mut AeadStreamCols<Felt> = row[..AEAD_STREAM_WIDTH].borrow_mut();
        let cols = cols.high_first_mut();
        cols.ctx = op.ctx;
        cols.clk = op.clk;
        cols.src_ptr = op.src_ptr;
        cols.lane_base = lane_base;
        cols.next_plaintext = plaintext[plaintext_offset + 1];
        cols.c_prev0 = op.ciphertext[lane_offset];
        cols.hi_quotient = canonical_hi_quotient(limbs[0]);
        cols.bytes = witnesses[1].bytes;
    }

    {
        let row = &mut chunk[2 * row_width..3 * row_width];
        let cols: &mut AeadStreamCols<Felt> = row[..AEAD_STREAM_WIDTH].borrow_mut();
        let cols = cols.low_second_mut();
        cols.ctx = op.ctx;
        cols.clk = op.clk;
        cols.src_ptr = op.src_ptr;
        cols.dst_ptr = dst_ptr;
        cols.lane_base = lane_base;
        cols.active_plaintext = plaintext[plaintext_offset + 1];
        cols.c_prev0 = op.ciphertext[lane_offset];
        cols.c_prev1 = op.ciphertext[lane_offset + 1];
        cols.bytes = witnesses[2].bytes;
    }

    {
        let row = &mut chunk[3 * row_width..4 * row_width];
        let cols: &mut AeadStreamCols<Felt> = row[..AEAD_STREAM_WIDTH].borrow_mut();
        let cols = cols.high_second_mut();
        cols.ctx = op.ctx;
        cols.clk = op.clk;
        cols.dst_ptr = dst_ptr;
        cols.lane_base = lane_base;
        cols.c_prev0 = op.ciphertext[lane_offset];
        cols.c_prev1 = op.ciphertext[lane_offset + 1];
        cols.c_prev2 = op.ciphertext[lane_offset + 2];
        cols.hi_quotient = canonical_hi_quotient(limbs[1]);
        cols.bytes = witnesses[3].bytes;
    }
}

#[derive(Debug, Clone, Copy)]
struct LimbXorWitness {
    bytes: [Felt; 12],
}

fn limb_xor_witness(a: u32, b: u32, c: u32) -> LimbXorWitness {
    debug_assert_eq!(a ^ b, c, "invalid u32 XOR witness");
    let a = a.to_le_bytes();
    let b = b.to_le_bytes();

    LimbXorWitness { bytes: and8_bytes(a, b) }
}

fn and8_bytes(a: [u8; 4], b: [u8; 4]) -> [Felt; 12] {
    [
        Felt::from_u8(a[0]),
        Felt::from_u8(a[1]),
        Felt::from_u8(a[2]),
        Felt::from_u8(a[3]),
        Felt::from_u8(b[0]),
        Felt::from_u8(b[1]),
        Felt::from_u8(b[2]),
        Felt::from_u8(b[3]),
        Felt::from_u8(a[0] & b[0]),
        Felt::from_u8(a[1] & b[1]),
        Felt::from_u8(a[2] & b[2]),
        Felt::from_u8(a[3] & b[3]),
    ]
}

fn count_and8_witness(counts: &mut [u64], bytes: [Felt; 12]) {
    for idx in 0..4 {
        count_and8(
            counts,
            felt_to_u8(bytes[idx]),
            felt_to_u8(bytes[4 + idx]),
            felt_to_u8(bytes[8 + idx]),
        );
    }
}

fn count_and8(counts: &mut [u64], a: u8, b: u8, result: u8) {
    debug_assert_eq!(a & b, result, "AEAD stream witness does not match the byte-pair table");
    let count_index = byte_pair_count_index(BytePairRelation::CanonicalXor, a, b);
    counts[count_index] += 1;
}

fn canonical_hi_quotient((lo, hi): (u32, u32)) -> Felt {
    let gap = Felt::from_u32(u32::MAX) - Felt::from_u32(hi);
    Felt::from_u32(lo) * gap.try_inverse().unwrap_or(ZERO)
}

fn felt_to_u32(value: Felt) -> u32 {
    u32::try_from(value.as_canonical_u64()).expect("AEAD stream value is not u32")
}

fn felt_to_u8(value: Felt) -> u8 {
    u8::try_from(value.as_canonical_u64()).expect("AEAD stream byte is not u8")
}