miden-processor 0.22.1

use alloc::vec::Vec;
use core::ops::Range;

use miden_air::trace::{
    RowIndex,
    chiplets::ace::{
        CLK_IDX, CTX_IDX, EVAL_OP_IDX, ID_0_IDX, ID_1_IDX, ID_2_IDX, M_0_IDX, M_1_IDX, PTR_IDX,
        READ_NUM_EVAL_IDX, SELECTOR_BLOCK_IDX, SELECTOR_START_IDX, V_0_0_IDX, V_0_1_IDX, V_1_0_IDX,
        V_1_1_IDX, V_2_0_IDX, V_2_1_IDX,
    },
};
use miden_core::{
    Felt, Word,
    field::{BasedVectorSpace, QuadFelt},
};

use super::{
    MAX_NUM_ACE_WIRES,
    instruction::{Op, decode_instruction},
};
use crate::{ContextId, errors::AceError};

/// Number of LogUp fractions in the wiring bus for rows in the `READ` section.
pub const NUM_ACE_LOGUP_FRACTIONS_READ: usize = 2;
/// Number of LogUp fractions in the wiring bus for rows in the `EVAL` section.
pub const NUM_ACE_LOGUP_FRACTIONS_EVAL: usize = 3;

/// Contains the variable and evaluation nodes resulting from the evaluation of a circuit.
/// The output value is checked to be equal to 0.
///
/// The set of nodes is used to fill the ACE chiplet trace.
#[derive(Debug, Clone)]
pub struct CircuitEvaluation {
    ctx: ContextId,
    clk: RowIndex,

    wire_bus: WireBus,

    col_ptr: Vec<Felt>,
    col_op: Vec<Felt>,

    col_wire_out: WireColumn,
    col_wire_left: WireColumn,
    col_wire_right: WireColumn,

    num_read_rows: u32,
    num_eval_rows: u32,
}

impl CircuitEvaluation {
    /// Generates the nodes in the graph generated by evaluating the inputs and circuit
    /// located in a contiguous memory region.
    ///
    /// # Panics:
    /// This function panics if the number of rows for each section leads to more than
    /// [`MAX_NUM_ACE_WIRES`] wires.
    pub fn new(ctx: ContextId, clk: RowIndex, num_read_rows: u32, num_eval_rows: u32) -> Self {
        let num_wires = 2 * (num_read_rows as u64) + (num_eval_rows as u64);
        assert!(num_wires <= MAX_NUM_ACE_WIRES as u64, "too many wires");

        let num_rows = num_read_rows + num_eval_rows;

        Self {
            ctx,
            clk,
            wire_bus: WireBus::new(num_wires as u32),
            col_ptr: Vec::with_capacity(num_rows as usize),
            col_op: Vec::with_capacity(num_eval_rows as usize),
            col_wire_out: WireColumn::new(num_rows as usize),
            col_wire_left: WireColumn::new(num_rows as usize),
            col_wire_right: WireColumn::new(num_eval_rows as usize),
            num_read_rows,
            num_eval_rows,
        }
    }

    pub fn num_rows(&self) -> usize {
        (self.num_read_rows + self.num_eval_rows) as usize
    }

    pub fn clk(&self) -> u32 {
        self.clk.into()
    }

    pub fn ctx(&self) -> u32 {
        self.ctx.into()
    }

    pub fn num_read_rows(&self) -> u32 {
        self.num_read_rows
    }

    pub fn num_eval_rows(&self) -> u32 {
        self.num_eval_rows
    }

    /// Reads the word from memory at `ptr`, interpreting it as `[v_00, v_01, v_10, v_11]`, and
    /// adds wires with values `v_0 = QuadFelt(v_00, v_01)` and `v_1 = QuadFelt(v_10, v_11)`.
    pub fn do_read(&mut self, ptr: Felt, word: Word) {
        // Add first variable as QuadFelt to wire bus
        let v_0 = QuadFelt::from_basis_coefficients_fn(|i: usize| [word[0], word[1]][i]);
        let id_0 = self.wire_bus.insert(v_0);
        self.col_wire_out.push(id_0, v_0);

        // Add second variable as QuadFelt to wire bus
        let v_1 = QuadFelt::from_basis_coefficients_fn(|i: usize| [word[2], word[3]][i]);
        let id_1 = self.wire_bus.insert(v_1);
        self.col_wire_left.push(id_1, v_1);

        // Add pointer to trace
        self.col_ptr.push(ptr);
    }

    /// Reads the next instruction at `ptr`, requests the inputs from the wire bus
    /// and inserts a new wire with the result.
    pub fn do_eval(&mut self, ptr: Felt, instruction: Felt) -> Result<(), AceError> {
        // Decode instruction, ensuring it is valid
        let (id_l, id_r, op) = decode_instruction(instruction)
            .ok_or(AceError("failed to decode instruction".into()))?;

        // Read value of id_l from wire bus, increasing its multiplicity
        let v_l = self
            .wire_bus
            .read_value(id_l)
            .ok_or(AceError("failed to read from the wiring bus".into()))?;
        let id_l = Felt::from_u32(id_l);
        self.col_wire_left.push(id_l, v_l);

        // Read value of id_r from wire bus, increasing its multiplicity
        let v_r = self
            .wire_bus
            .read_value(id_r)
            .ok_or(AceError("failed to read from the wiring bus".into()))?;
        let id_r = Felt::from_u32(id_r);
        self.col_wire_right.push(id_r, v_r);

        // Compute v_out and insert it into the wire bus.
        let v_out = match op {
            Op::Sub => v_l - v_r,
            Op::Mul => v_l * v_r,
            Op::Add => v_l + v_r,
        };
        let id_out = self.wire_bus.insert(v_out);
        self.col_wire_out.push(id_out, v_out);

        // Add op to column
        let op_sub = -Felt::ONE;
        let op_mul = Felt::ZERO;
        let op_add = Felt::ONE;
        let op = match op {
            Op::Sub => op_sub,
            Op::Mul => op_mul,
            Op::Add => op_add,
        };
        self.col_op.push(op);

        // Add pointer to trace
        self.col_ptr.push(ptr);

        Ok(())
    }

    pub fn fill(&self, offset: usize, columns: &mut [Vec<Felt>]) {
        let num_read_rows = self.num_read_rows as usize;
        let num_eval_rows = self.num_eval_rows as usize;
        let num_rows = num_read_rows + num_eval_rows;
        let read_range = Range {
            start: offset,
            end: offset + num_read_rows,
        };
        let eval_range = Range {
            start: read_range.end,
            end: read_range.end + num_eval_rows,
        };

        // Fill start selector
        columns[SELECTOR_START_IDX][offset] = Felt::ONE;
        columns[SELECTOR_START_IDX][(offset + 1)..].fill(Felt::ZERO);

        // Block flag column
        let f_read = Felt::ZERO;
        let f_eval = Felt::ONE;
        columns[SELECTOR_BLOCK_IDX][read_range.clone()].fill(f_read);
        columns[SELECTOR_BLOCK_IDX][eval_range.clone()].fill(f_eval);

        // Fill ctx column which is constant across the section
        let ctx_felt = self.ctx.into();
        columns[CTX_IDX][offset..offset + num_rows].fill(ctx_felt);

        // Fill ptr column.
        columns[PTR_IDX][offset..offset + num_rows].copy_from_slice(&self.col_ptr);

        // Fill clk column which is constant across the section
        let clk_felt = self.clk.into();
        columns[CLK_IDX][offset..offset + num_rows].fill(clk_felt);

        // Fill n_eval which is constant across the read block
        let eval_section_first_idx = Felt::from_u32(self.num_eval_rows - 1);
        columns[READ_NUM_EVAL_IDX][read_range.clone()].fill(eval_section_first_idx);

        // Fill OP column for EVAL rows
        columns[EVAL_OP_IDX][eval_range.clone()].copy_from_slice(&self.col_op);

        // Fill wire 0 columns for all rows
        columns[ID_0_IDX][offset..offset + num_rows].copy_from_slice(&self.col_wire_out.id);
        columns[V_0_0_IDX][offset..offset + num_rows].copy_from_slice(&self.col_wire_out.v_0);
        columns[V_0_1_IDX][offset..offset + num_rows].copy_from_slice(&self.col_wire_out.v_1);

        // Fill wire 1 columns for all rows
        columns[ID_1_IDX][offset..offset + num_rows].copy_from_slice(&self.col_wire_left.id);
        columns[V_1_0_IDX][offset..offset + num_rows].copy_from_slice(&self.col_wire_left.v_0);
        columns[V_1_1_IDX][offset..offset + num_rows].copy_from_slice(&self.col_wire_left.v_1);

        // Fill wire 2 columns for EVAL rows
        columns[ID_2_IDX][eval_range.clone()].copy_from_slice(&self.col_wire_right.id);
        columns[V_2_0_IDX][eval_range.clone()].copy_from_slice(&self.col_wire_right.v_0);
        columns[V_2_1_IDX][eval_range.clone()].copy_from_slice(&self.col_wire_right.v_1);

        // Fill multiplicity 0 column for all rows
        let mut multiplicities_iter = self.wire_bus.wires.iter().map(|(_v, m)| Felt::from_u32(*m));
        // In the READ block, we inserted wires w_0 and w_1
        for row_index in read_range {
            let m_0 = multiplicities_iter
                .next()
                .expect("the m0 multiplicities were not constructed properly");
            let m_1 = multiplicities_iter
                .next()
                .expect("the m1 multiplicities were not constructed properly");
            columns[M_0_IDX][row_index] = m_0;
            columns[M_1_IDX][row_index] = m_1;
        }
        // In the EVAL block, we inserted wire w_0
        for row_index in eval_range {
            let m_0 = multiplicities_iter
                .next()
                .expect("the m0 multiplicities were not constructed properly");
            columns[M_0_IDX][row_index] = m_0;
        }

        debug_assert!(multiplicities_iter.next().is_none());
    }

    /// Returns the output value, if the circuit has finished evaluating.
    pub fn output_value(&self) -> Option<QuadFelt> {
        if !self.wire_bus.is_finalized() {
            return None;
        }
        self.wire_bus.wires.last().map(|(v, _m)| *v)
    }
}

/// Set of columns for a given wire containing `[id, v_0, v_1]`, where `id` is the wire identifier
/// and `v = (v_0, v_1)` is the quadratic extension field element value of the wire.
#[derive(Debug, Clone)]
struct WireColumn {
    id: Vec<Felt>,
    v_0: Vec<Felt>,
    v_1: Vec<Felt>,
}

impl WireColumn {
    fn new(num_rows: usize) -> Self {
        Self {
            id: Vec::with_capacity(num_rows),
            v_0: Vec::with_capacity(num_rows),
            v_1: Vec::with_capacity(num_rows),
        }
    }

    /// Pushes the wire `(id, v)` to the columns.
    fn push(&mut self, id: Felt, v: QuadFelt) {
        self.id.push(id);
        let v: &[Felt] = v.as_basis_coefficients_slice();
        self.v_0.push(v[0]);
        self.v_1.push(v[1]);
    }
}

/// A LogUp-based bus used for wiring the gates of the circuit.
///
/// Gates are fan-in 2 but can have fan-out up to the field characteristic which, given the bounds
/// on the execution trace length, means practically arbitrary fan-out.
/// The main idea, with some slight variations between the `READ` and `EVAL` sections, is, for each
/// gate, to "receive" the values of the input wires from the bus and to "send" the value of
/// the value of the output wire back with multiplicity equal to the fan-out of the respective gate.
/// Note that the messages include extra data in order to avoid collisions.
#[derive(Debug, Clone)]
struct WireBus {
    // Circuit ID as Felt of the next wire to be inserted
    id_next: Felt,
    // Pairs of values and multiplicities
    // The wire with index `id` is stored at `num_wires - 1 - id`
    wires: Vec<(QuadFelt, u32)>,
    // Total expected number of wires to be inserted.
    num_wires: u32,
}

impl WireBus {
    fn new(num_wires: u32) -> Self {
        Self {
            wires: Vec::with_capacity(num_wires as usize),
            num_wires,
            id_next: Felt::from_u32(num_wires - 1),
        }
    }

    /// Inserts a new value into the bus, and returns its expected id as `Felt`
    fn insert(&mut self, value: QuadFelt) -> Felt {
        debug_assert!(!self.is_finalized());
        self.wires.push((value, 0));
        let id = self.id_next;
        self.id_next -= Felt::ONE;
        id
    }

    /// Reads the value of a wire with given `id`, incrementing its multiplicity.
    /// Returns `None` if the requested wire has not been inserted yet.
    fn read_value(&mut self, id: u32) -> Option<QuadFelt> {
        // Ensures subtracting the id from num_wires results in a valid wire index
        let (v, m) = self
            .num_wires
            .checked_sub(id + 1)
            .and_then(|id| self.wires.get_mut(id as usize))?;
        *m += 1;
        Some(*v)
    }

    /// Return true if the expected number of wires have been inserted.
    fn is_finalized(&self) -> bool {
        self.wires.len() == self.num_wires as usize
    }
}