zisk-sm-binary 1.3.0-alpha

Binary operations state machine for the ZisK zkVM
//! The `BinaryCounter` module defines a device for tracking and processing binary-related operations
//! sent over the data bus. It serves a purpose:
//! - Counting different types of binary operations, to decide if uses specific add instances or not.
//!
//! This module implements the `Metrics` and `BusDevice` traits, enabling seamless integration with
//! the system bus for both monitoring and input generation.

use crate::{
    add_family_kind, BinaryBasicFrops, BinaryExtensionFrops, KIND_ADD_FULL, KIND_ADD_HI,
    KIND_SH3ADD_ADD, KIND_SH3ADD_HI,
};
use zisk_common::{BusDevice, BusId, Counter, Metrics, A, B, OP, OPERATION_BUS_ID, OP_TYPE};
use zisk_core::ZiskOperationType;

/// The `BinaryCounter` struct represents a counter that monitors and measures
/// binary-related operations on the data bus.
///
/// It tracks specific operations and types and updates differents counters for each
/// accepted operation whenever data is processed on the bus.
///
/// The buckets are **disjoint**: every binary / binary-extension operation on the bus lands in
/// exactly one of them, so their sum is the total number of operations. Each bucket corresponds to
/// the air (or set of airs) able to prove that operation, which is what lets the planner size the
/// instances — see [`crate::add_shape`] for the split.
#[derive(Default)]
pub struct BinaryCounter {
    /// Counter for binary add operations needing the full 64-bit add (only add, no addw).
    /// Proven by `BinaryAdd`, or by `Binary` when no dedicated add air is used.
    pub counter_add: Counter,

    /// Counter for add operations whose result fits in the low limb ([`AddShape::Hi`] and
    /// [`AddShape::HiNeg`]). `BinaryAddHi` packs these, LANES_X_ROW per row, in any of its slots.
    pub counter_add_hi: Counter,

    /// Counter for basic binary operations, but not considering add operations
    pub counter_basic_wo_add: Counter,

    /// SH3ADD whose whole result fits in the low limb, so the packed airs prove it too.
    pub counter_sh3add_hi: Counter,

    /// SH3ADD that only the full 64-bit add can take.
    pub counter_sh3add_add: Counter,

    /// Counter for binary extension operations. Both extension airs are instantiated `full`, so
    /// they all belong to one bucket.
    pub counter_extension: Counter,
}

impl BinaryCounter {
    /// Creates a new instance of `BinaryCounter`.
    ///
    /// # Arguments
    /// * `mode` - The mode of the bus device.
    ///
    /// # Returns
    /// A new `BinaryCounter` instance.
    pub fn new() -> Self {
        Self::default()
    }

    /// Processes data received on the bus, updating counters and generating inputs when applicable.
    ///
    /// # Arguments
    /// * `bus_id` - The ID of the bus sending the data.
    /// * `data` - The data received from the bus.
    /// * `pending` – A queue of pending bus operations used to send derived inputs.
    ///
    /// # Returns
    /// A boolean indicating whether the program should continue execution or terminate.
    /// Returns `true` to continue execution, `false` to stop.
    #[inline(always)]
    pub fn process_data(&mut self, bus_id: &BusId, data: &[u64]) -> bool {
        debug_assert!(*bus_id == OPERATION_BUS_ID);

        self.measure(data);

        true
    }
}

impl Metrics for BinaryCounter {
    /// Tracks activity on the connected bus and updates counters for recognized operations.
    ///
    /// # Arguments
    /// * `data` - The data received from the bus.
    ///
    /// # Returns
    /// An empty vector, as this implementation does not produce any derived inputs for the bus.
    #[inline(always)]
    fn measure(&mut self, data: &[u64]) {
        // Precomputed constants to avoid casting each time
        const BINARY: u64 = ZiskOperationType::Binary as u64;
        const BINARY_E: u64 = ZiskOperationType::BinaryE as u64;

        let op_type = data[OP_TYPE];
        if op_type == BINARY {
            // Always read the OP index (assume well-formed trace)
            let op = data[OP] as u8;

            // One classifier for the whole family, shared with every collector, so the sizing here
            // and the collection later can never disagree about where an operation belongs.
            let counter = match add_family_kind(op, data[A], data[B]) {
                KIND_ADD_HI => &mut self.counter_add_hi,
                KIND_ADD_FULL => &mut self.counter_add,
                KIND_SH3ADD_HI => &mut self.counter_sh3add_hi,
                KIND_SH3ADD_ADD => &mut self.counter_sh3add_add,
                _ => &mut self.counter_basic_wo_add,
            };

            if BinaryBasicFrops::is_frequent_op(op, data[A], data[B]) {
                counter.update_frops(1);
            } else {
                counter.update(1);
            }
        } else if op_type == BINARY_E {
            if BinaryExtensionFrops::is_frequent_op(data[OP] as u8, data[A], data[B]) {
                self.counter_extension.update_frops(1);
            } else {
                self.counter_extension.update(1);
            }
        }
    }

    /// Provides a dynamic reference for downcasting purposes.
    ///
    /// # Returns
    /// A reference to `self` as `dyn std::any::Any`.
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }
}

impl BusDevice<u64> for BinaryCounter {
    /// Provides a dynamic reference for downcasting purposes.
    fn as_any(self: Box<Self>) -> Box<dyn std::any::Any> {
        self
    }
}