openvm-circuit 2.0.1

use std::{
    array,
    borrow::{Borrow, BorrowMut},
    iter,
};

use openvm_circuit_primitives::{ColumnsAir, StructReflection, StructReflectionHelper};
use openvm_circuit_primitives_derive::AlignedBorrow;
use openvm_cpu_backend::CpuBackend;
use openvm_stark_backend::{
    interaction::{InteractionBuilder, PermutationCheckBus},
    p3_air::{Air, AirBuilder, BaseAir},
    p3_field::{PrimeCharacteristicRing, PrimeField32},
    p3_matrix::{dense::RowMajorMatrix, Matrix},
    p3_maybe_rayon::prelude::*,
    prover::AirProvingContext,
    BaseAirWithPublicValues, PartitionedBaseAir, StarkProtocolConfig, Val,
};
use tracing::instrument;

use super::{merkle::SerialReceiver, online::INITIAL_TIMESTAMP};
use crate::{
    arch::{hasher::Hasher, ADDR_SPACE_OFFSET, DEFAULT_BLOCK_SIZE},
    primitives::Chip,
    system::memory::{
        controller::CHUNK, offline_checker::MemoryBus, MemoryAddress, MemoryImage,
        TimestampedEquipartition,
    },
};

/// Number of DEFAULT_BLOCK_SIZE blocks per CHUNK (e.g., 2 for 8/4).
/// Blocks are on the same row only for Merkle tree hashing (8 bytes at a time).
/// Memory bus interactions use per-block timestamps.
pub const BLOCKS_PER_CHUNK: usize = CHUNK / DEFAULT_BLOCK_SIZE;

/// The values describe aligned chunk of memory of size `CHUNK`---the data together with the last
/// accessed timestamp---in either the initial or final memory state.
#[repr(C)]
#[derive(Debug, AlignedBorrow, StructReflection)]
pub struct PersistentBoundaryCols<T, const CHUNK: usize> {
    // `expand_direction` =  1 corresponds to initial memory state
    // `expand_direction` = -1 corresponds to final memory state
    // `expand_direction` =  0 corresponds to irrelevant row (all interactions multiplicity 0)
    pub expand_direction: T,
    pub address_space: T,
    pub leaf_label: T,
    pub values: [T; CHUNK],
    pub hash: [T; CHUNK],
    /// Per-block timestamps. Each DEFAULT_BLOCK_SIZE block within the chunk has its own timestamp.
    /// For untouched blocks, timestamp stays at 0 (balances: boundary sends at t=0 init, receives
    /// at t=0 final).
    pub timestamps: [T; BLOCKS_PER_CHUNK],
}

/// Imposes the following constraints:
/// - `expand_direction` should be -1, 0, 1
///
/// Sends the following interactions:
/// - if `expand_direction` is 1, sends `[0, 0, address_space_label, leaf_label]` to `merkle_bus`.
/// - if `expand_direction` is -1, receives `[1, 0, address_space_label, leaf_label]` from
///   `merkle_bus`.
#[derive(Clone, Debug, ColumnsAir)]
#[columns_via(PersistentBoundaryCols<u8, CHUNK>)]
pub struct PersistentBoundaryAir<const CHUNK: usize> {
    pub memory_bus: MemoryBus,
    pub merkle_bus: PermutationCheckBus,
    pub compression_bus: PermutationCheckBus,
}

impl<const CHUNK: usize, F> BaseAir<F> for PersistentBoundaryAir<CHUNK> {
    fn width(&self) -> usize {
        PersistentBoundaryCols::<F, CHUNK>::width()
    }
}

impl<const CHUNK: usize, F> BaseAirWithPublicValues<F> for PersistentBoundaryAir<CHUNK> {}
impl<const CHUNK: usize, F> PartitionedBaseAir<F> for PersistentBoundaryAir<CHUNK> {}

impl<const CHUNK: usize, AB: InteractionBuilder> Air<AB> for PersistentBoundaryAir<CHUNK> {
    fn eval(&self, builder: &mut AB) {
        let main = builder.main();
        let local = main.row_slice(0).expect("window should have two elements");
        let local: &PersistentBoundaryCols<AB::Var, CHUNK> = (*local).borrow();

        // `direction` should be -1, 0, 1
        builder.assert_eq(
            local.expand_direction,
            local.expand_direction * local.expand_direction * local.expand_direction,
        );

        // Constrain that an "initial" row has all timestamp zero.
        // Since `direction` is constrained to be in {-1, 0, 1}, we can select `direction == 1`
        // with the constraint below.
        let mut when_initial =
            builder.when(local.expand_direction * (local.expand_direction + AB::F::ONE));
        for i in 0..BLOCKS_PER_CHUNK {
            when_initial.assert_zero(local.timestamps[i]);
        }

        let mut expand_fields = vec![
            // direction =  1 => is_final = 0
            // direction = -1 => is_final = 1
            local.expand_direction.into(),
            AB::Expr::ZERO,
            local.address_space - AB::F::from_u32(ADDR_SPACE_OFFSET),
            local.leaf_label.into(),
        ];
        expand_fields.extend(local.hash.map(Into::into));
        self.merkle_bus
            .interact(builder, expand_fields, local.expand_direction.into());

        self.compression_bus.interact(
            builder,
            iter::empty()
                .chain(local.values.map(Into::into))
                .chain(iter::repeat_n(AB::Expr::ZERO, CHUNK))
                .chain(local.hash.map(Into::into)),
            local.expand_direction * local.expand_direction,
        );

        let chunk_size_f = AB::F::from_usize(CHUNK);
        for block_idx in 0..BLOCKS_PER_CHUNK {
            let offset = AB::F::from_usize(block_idx * DEFAULT_BLOCK_SIZE);
            // Split the 1xCHUNK leaf into DEFAULT_BLOCK_SIZE-sized bus messages.
            // Each block uses its own timestamp - untouched blocks stay at t=0.
            self.memory_bus
                .send(
                    MemoryAddress::new(
                        local.address_space,
                        local.leaf_label * chunk_size_f + offset,
                    ),
                    local.values
                        [block_idx * DEFAULT_BLOCK_SIZE..(block_idx + 1) * DEFAULT_BLOCK_SIZE]
                        .to_vec(),
                    local.timestamps[block_idx],
                )
                .eval(builder, local.expand_direction);
        }
    }
}

pub struct PersistentBoundaryChip<F, const CHUNK: usize> {
    pub air: PersistentBoundaryAir<CHUNK>,
    touched_labels: Option<Vec<FinalTouchedLabel<F, CHUNK>>>,
    overridden_height: Option<usize>,
}

#[derive(Debug)]
pub struct FinalTouchedLabel<F, const CHUNK: usize> {
    address_space: u32,
    label: u32,
    init_values: [F; CHUNK],
    final_values: [F; CHUNK],
    init_hash: [F; CHUNK],
    final_hash: [F; CHUNK],
    /// Per-block timestamps. Each DEFAULT_BLOCK_SIZE block has its own timestamp.
    final_timestamps: [u32; BLOCKS_PER_CHUNK],
}

type BlockInfo<F> = (usize, u32, [F; DEFAULT_BLOCK_SIZE]); // (block_idx, timestamp, values)
type EnrichedEntry<F> = ((u32, u32), BlockInfo<F>); // (chunk_key, block_info)
pub(crate) type ChunkedTouchedMemory<F> = Vec<((u32, u32), Vec<BlockInfo<F>>)>;

pub(crate) fn group_touched_memory_by_chunk<F: Copy + Send + Sync>(
    final_memory: &TimestampedEquipartition<F, DEFAULT_BLOCK_SIZE>,
) -> ChunkedTouchedMemory<F> {
    let mut enriched: Vec<EnrichedEntry<F>> = final_memory
        .par_iter()
        .map(|&((addr_space, ptr), ts_values)| {
            let chunk_label = ptr / CHUNK as u32;
            let block_idx = ((ptr % CHUNK as u32) / DEFAULT_BLOCK_SIZE as u32) as usize;
            let key = (addr_space, chunk_label);
            let block_info = (block_idx, ts_values.timestamp, ts_values.values);
            (key, block_info)
        })
        .collect();
    enriched.sort_unstable_by_key(|(key, _)| *key);

    enriched
        .chunk_by(|a, b| a.0 == b.0)
        .map(|group| {
            let key = group[0].0;
            let blocks = group.iter().map(|&(_, info)| info).collect();
            (key, blocks)
        })
        .collect()
}

impl<const CHUNK: usize, F: PrimeField32> PersistentBoundaryChip<F, CHUNK> {
    pub fn new(
        memory_bus: MemoryBus,
        merkle_bus: PermutationCheckBus,
        compression_bus: PermutationCheckBus,
    ) -> Self {
        Self {
            air: PersistentBoundaryAir {
                memory_bus,
                merkle_bus,
                compression_bus,
            },
            touched_labels: None,
            overridden_height: None,
        }
    }

    pub fn set_overridden_height(&mut self, overridden_height: usize) {
        self.overridden_height = Some(overridden_height);
    }

    /// Finalize the boundary chip with per-block timestamped memory.
    ///
    /// `final_memory` is at DEFAULT_BLOCK_SIZE granularity (4 bytes per entry, single timestamp
    /// each). This function rechunks into CHUNK-sized (8 bytes) groups with per-block
    /// timestamps. Untouched blocks within a touched chunk get values from initial_memory and
    /// timestamp 0.
    #[instrument(name = "boundary_finalize", level = "debug", skip_all)]
    pub(crate) fn finalize<H>(
        &mut self,
        initial_memory: &MemoryImage,
        // Touched stuff at DEFAULT_BLOCK_SIZE granularity
        final_memory: &TimestampedEquipartition<F, DEFAULT_BLOCK_SIZE>,
        hasher: &H,
    ) where
        H: Hasher<CHUNK, F> + Sync + for<'a> SerialReceiver<&'a [F]>,
    {
        let final_touched_labels: Vec<_> = group_touched_memory_by_chunk(final_memory)
            .into_par_iter()
            .map(|((addr_space, chunk_label), blocks)| {
                let chunk_ptr = chunk_label * CHUNK as u32;
                // SAFETY: addr_space from `final_memory` are all in bounds
                let init_values: [F; CHUNK] = array::from_fn(|i| unsafe {
                    initial_memory.get_f::<F>(addr_space, chunk_ptr + i as u32)
                });

                let mut final_values = init_values;
                let mut timestamps = [0u32; BLOCKS_PER_CHUNK];

                for (block_idx, ts, values) in blocks {
                    timestamps[block_idx] = ts;
                    for (i, &val) in values.iter().enumerate() {
                        final_values[block_idx * DEFAULT_BLOCK_SIZE + i] = val;
                    }
                }

                let initial_hash = hasher.hash(&init_values);
                let final_hash = hasher.hash(&final_values);
                FinalTouchedLabel {
                    address_space: addr_space,
                    label: chunk_label,
                    init_values,
                    final_values,
                    init_hash: initial_hash,
                    final_hash,
                    final_timestamps: timestamps,
                }
            })
            .collect();
        for l in &final_touched_labels {
            hasher.receive(&l.init_values);
            hasher.receive(&l.final_values);
        }
        self.touched_labels = Some(final_touched_labels);
    }
}

impl<const CHUNK: usize, RA, SC> Chip<RA, CpuBackend<SC>> for PersistentBoundaryChip<Val<SC>, CHUNK>
where
    SC: StarkProtocolConfig,
    Val<SC>: PrimeField32,
{
    fn generate_proving_ctx(&self, _: RA) -> AirProvingContext<CpuBackend<SC>> {
        let trace = {
            let touched_labels = self
                .touched_labels
                .as_ref()
                .expect("Cannot generate trace before finalization");
            let width = PersistentBoundaryCols::<Val<SC>, CHUNK>::width();
            // Boundary AIR should always present in order to fix the AIR ID of merkle AIR.
            let mut height = (2 * touched_labels.len()).next_power_of_two();
            if let Some(mut oh) = self.overridden_height {
                oh = oh.next_power_of_two();
                assert!(
                    oh >= height,
                    "Overridden height is less than the required height"
                );
                height = oh;
            }
            let mut rows = Val::<SC>::zero_vec(height * width);

            rows.par_chunks_mut(2 * width)
                .zip(touched_labels.par_iter())
                .for_each(|(row, touched_label)| {
                    let (initial_row, final_row) = row.split_at_mut(width);
                    *initial_row.borrow_mut() = PersistentBoundaryCols {
                        expand_direction: Val::<SC>::ONE,
                        address_space: Val::<SC>::from_u32(touched_label.address_space),
                        leaf_label: Val::<SC>::from_u32(touched_label.label),
                        values: touched_label.init_values,
                        hash: touched_label.init_hash,
                        timestamps: [Val::<SC>::from_u32(INITIAL_TIMESTAMP); BLOCKS_PER_CHUNK],
                    };

                    *final_row.borrow_mut() = PersistentBoundaryCols {
                        expand_direction: Val::<SC>::NEG_ONE,
                        address_space: Val::<SC>::from_u32(touched_label.address_space),
                        leaf_label: Val::<SC>::from_u32(touched_label.label),
                        values: touched_label.final_values,
                        hash: touched_label.final_hash,
                        timestamps: touched_label.final_timestamps.map(Val::<SC>::from_u32),
                    };
                });
            RowMajorMatrix::new(rows, width)
        };
        AirProvingContext::simple_no_pis(trace)
    }
}