miden-processor 0.22.1

use alloc::{collections::BTreeMap, vec::Vec};
use core::mem::MaybeUninit;

use miden_air::trace::RANGE_CHECK_TRACE_WIDTH;

use super::RowIndex;
use crate::{
    Felt, ZERO,
    utils::{assume_init_vec, uninit_vector},
};

mod aux_trace;
pub use aux_trace::AuxTraceBuilder;

#[cfg(test)]
mod tests;

// TRACE
// ================================================================================================

pub struct RangeCheckTrace {
    pub(crate) trace: [Vec<Felt>; RANGE_CHECK_TRACE_WIDTH],
    pub(crate) aux_builder: AuxTraceBuilder,
}

// RANGE CHECKER
// ================================================================================================

/// Range checker for the VM.
///
/// This component is responsible for building an execution trace for all 16-bit range checks
/// performed by the VM. Thus, the [RangeChecker] doesn't actually check if a given value fits
/// into 16-bits, but rather keeps track of all 16-bit range checks performed by the VM.
///
/// ## Execution trace
/// The execution trace generated by the range checker consists of 2 columns. Conceptually, the
/// table starts with value 0 and ends with value 65535.
///
/// The layout is illustrated below.
///
///    m     v
/// ├─────┴─────┤
///
/// In the above, the meaning of the columns is as follows:
/// - Column `v` contains the value being range-checked where `v` must be a 16-bit value. The values
///   must be in increasing order and the jump allowed between two values should be a power of 3
///   less than or equal to 3^7, and duplicates are allowed.
/// - Column `m` specifies the lookup multiplicity, which is how many lookups are to be included for
///   a given value.
///
/// Thus, for example, if a value was range-checked just once, we'll need to add a single row to
/// the table with (m, v) set to (1, v), where v is the value. If the value was range-checked 5
/// times, we'll need to specify the row (5, v).
pub struct RangeChecker {
    /// Tracks lookup count for each checked value.
    lookups: BTreeMap<u16, usize>,
    /// Range check lookups performed by all user operations, grouped and sorted by clock cycle.
    /// Each cycle is mapped to a vector of the range checks requested at that cycle, which can
    /// come from the stack, memory, or both.
    cycle_lookups: BTreeMap<RowIndex, Vec<u16>>,
}

impl RangeChecker {
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Returns a new [RangeChecker] instantiated with an empty lookup table.
    pub fn new() -> Self {
        let mut lookups = BTreeMap::new();
        // we need to make sure that the first row after the padded section and the last row of the
        // range checker table are initialized. this simplifies trace table building later on.
        lookups.insert(0, 0);
        lookups.insert(u16::MAX, 0);
        Self { lookups, cycle_lookups: BTreeMap::new() }
    }

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

    /// Adds the specified value to the trace of this range checker's lookups.
    pub fn add_value(&mut self, value: u16) {
        self.lookups.entry(value).and_modify(|v| *v += 1).or_insert(1);
    }

    /// Adds range check lookups from the stack or memory to this [RangeChecker] instance.
    pub fn add_range_checks(&mut self, clk: RowIndex, values: &[u16]) {
        // range checks requests only come from memory or from the stack, which always request 2 or
        // 4 lookups respectively.
        debug_assert!(values.len() == 2 || values.len() == 4);

        for value in values.iter() {
            // add the specified value to the trace of this range checker's lookups.
            self.add_value(*value);
        }

        // track the range check requests at each cycle
        // TODO: optimize this to use a struct instead of vectors, e.g. (#2793):
        // struct MemoryLookupValues {
        //   num_lookups: u8,
        //   lookup_values: [u16; 6],
        // }
        self.cycle_lookups
            .entry(clk)
            .and_modify(|entry| entry.append(&mut values.to_vec()))
            .or_insert_with(|| values.to_vec());
    }

    // EXECUTION TRACE GENERATION (INTERNAL)
    // --------------------------------------------------------------------------------------------

    /// Converts this [RangeChecker] into an execution trace with 2 columns and the number of rows
    /// specified by the `target_len` parameter.
    ///
    /// If the number of rows needed to represent execution trace of this range checker is smaller
    /// than `target_len` parameter, the trace is padded with extra rows.
    ///
    /// # Panics
    /// Panics if `target_len` is not a power of two or is smaller than the trace length needed
    /// to represent all lookups in this range checker.
    pub fn into_trace_with_table(self, trace_len: usize, target_len: usize) -> RangeCheckTrace {
        assert!(target_len.is_power_of_two(), "target trace length is not a power of two");

        // determine the length of the trace required to support all the lookups in this range
        // checker, and make sure this length is smaller than or equal to the target trace length.
        assert!(trace_len <= target_len, "target trace length too small");

        // allocate memory for the trace
        let mut trace = [uninit_vector(target_len), uninit_vector(target_len)];

        // determine the number of padding rows needed to get to target trace length and pad the
        // table with the required number of rows.
        let num_padding_rows = target_len - trace_len;
        trace[0][..num_padding_rows].fill(MaybeUninit::new(ZERO));
        trace[1][..num_padding_rows].fill(MaybeUninit::new(ZERO));

        // build the trace table
        let mut i = num_padding_rows;
        let mut prev_value = 0u16;
        for (&value, &num_lookups) in self.lookups.iter() {
            write_rows(&mut trace, &mut i, num_lookups, value, prev_value);
            prev_value = value;
        }

        // pad the trace with an extra row of 0 lookups for u16::MAX so that when b_range is built
        // there is space for the inclusion of u16::MAX range check lookups before the trace ends.
        // (When there is data at the end of the main trace, auxiliary bus columns always need to be
        // one row longer than the main trace, since values in the bus column are based on data from
        // the "current" row of the main trace but placed into the "next" row of the bus column.)
        write_trace_row(&mut trace, &mut i, 0, (u16::MAX).into());

        assert_eq!(i, target_len, "range checker trace not fully initialized; trace_len mismatch");

        // all elements are now initialized
        let [t0, t1] = trace;
        let trace = unsafe { [assume_init_vec(t0), assume_init_vec(t1)] };

        RangeCheckTrace {
            trace,
            aux_builder: AuxTraceBuilder::new(
                self.lookups.keys().cloned().collect(),
                self.cycle_lookups,
                num_padding_rows,
            ),
        }
    }

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

    /// Returns the number of rows needed to support all 16-bit lookups requested by the VM.
    pub fn get_number_range_checker_rows(&self) -> usize {
        // pad the trace length by one, to account for an extra row of the u16::MAX value at the end
        // of the trace, required for building the `b_range` column.
        let mut num_rows = 1;

        let mut prev_value = 0u16;
        for value in self.lookups.keys() {
            // add one row for each value in the range checker table
            num_rows += 1;
            // determine the delta between this and the previous value. we need to know this delta
            // to determine if we need to insert any "bridge" rows to the  table, this is needed
            // since the gap between two values in the range checker can only be a power of 3 less
            // than or equal to 3^7.
            let delta = value - prev_value;
            num_rows += get_num_bridge_rows(delta);
            prev_value = *value;
        }
        num_rows
    }

    // TEST HELPERS
    // --------------------------------------------------------------------------------------------

    /// Returns length of execution trace required to describe all 16-bit range checks performed
    /// by the VM.
    #[cfg(test)]
    pub fn trace_len(&self) -> usize {
        self.get_number_range_checker_rows()
    }

    /// Converts this [RangeChecker] into an execution trace with 2 columns and the number of rows
    /// specified by the `target_len` parameter.
    ///
    /// Wrapper for [`RangeChecker::into_trace_with_table`].
    #[cfg(test)]
    pub fn into_trace(self, target_len: usize) -> RangeCheckTrace {
        let table_len = self.get_number_range_checker_rows();
        self.into_trace_with_table(table_len, target_len)
    }
}

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

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

/// Calculates the number of bridge rows that are need to be added to the trace between two values
/// to be range checked.
pub fn get_num_bridge_rows(delta: u16) -> usize {
    let mut gap = delta;
    let mut bridge_rows = 0_usize;
    let mut stride = 3_u16.pow(7);
    while gap != stride {
        if gap > stride {
            bridge_rows += 1;
            gap -= stride;
        } else {
            stride /= 3;
        }
    }
    bridge_rows
}

/// Adds a row for the values to be range checked. In case the difference between the current and
/// next value is not a power of 3, this function will add additional bridge rows to the trace.
fn write_rows(
    trace: &mut [Vec<MaybeUninit<Felt>>],
    step: &mut usize,
    num_lookups: usize,
    value: u16,
    prev_value: u16,
) {
    let mut gap = value - prev_value;
    let mut prev_val = prev_value;
    let mut stride = 3_u16.pow(7);
    while gap != stride {
        if gap > stride {
            gap -= stride;
            prev_val += stride;
            write_trace_row(trace, step, 0, prev_val as u64);
        } else {
            stride /= 3;
        }
    }
    write_trace_row(trace, step, num_lookups, value as u64);
}

/// Populates a single row at the specified step in the trace table.
fn write_trace_row(
    trace: &mut [Vec<MaybeUninit<Felt>>],
    step: &mut usize,
    num_lookups: usize,
    value: u64,
) {
    trace[0][*step].write(Felt::new(num_lookups as u64));
    trace[1][*step].write(Felt::new(value));
    *step += 1;
}