feat: General optimizations

.gitignore

@@ -21,4 +21,7 @@ benchmark.csv
 # Build Artifacts
 recursion/gnark-ffi/build
 prover/build
-prover/*.tar.gz
+prover/*.tar.gz
+
+# IDE Conf
+.idea

Cargo.toml

@@ -74,7 +74,7 @@ ff = "0.13"
 futures = "0.3.30"
 futures-util = "0.3.14"
 getrandom = "=0.2.14" # 0.2.15 depends on yanked libc 0.2.154
-hashbrown = "0.14.5"
+hashbrown = { version = "0.14.5", features = ["serde"] }
 hex = "0.4.3"
 home = "0.5.9"
 hybrid-array = "0.2.0-rc"

core/Cargo.toml

@@ -42,6 +42,7 @@ bls12_381 = { workspace = true }
 cfg-if = { workspace = true }
 curve25519-dalek = { workspace = true }
 elliptic-curve = { workspace = true }
+hashbrown = { workspace = true }
 hex = { workspace = true }
 hybrid-array = { workspace = true }
 k256 = { workspace = true, features = ["expose-field"] }

core/benches/fibonacci.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use hashbrown::HashMap;
 
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
 use sphinx_core::{

core/src/alu/divrem/mod.rs

@@ -64,7 +64,7 @@ mod utils;
 
 use core::borrow::{Borrow, BorrowMut};
 use core::mem::size_of;
-use std::collections::HashMap;
+use hashbrown::HashMap;
 
 use p3_air::{Air, AirBuilder, BaseAir};
 use p3_field::AbstractField;
@@ -394,7 +394,7 @@ impl<F: PrimeField> MachineAir<F> for DivRemChip {
  }
  let mut alu_events = HashMap::new();
  alu_events.insert(Opcode::ADD, add_events);
- output.add_alu_events(&alu_events);
+ output.add_alu_events(&mut alu_events);
  }
 
  let mut lower_word = 0;

core/src/bytes/event.rs

@@ -1,12 +1,10 @@
-use std::collections::BTreeMap;
-
 use p3_field::PrimeField32;
 use serde::{Deserialize, Serialize};
 
 use super::ByteOpcode;
 
 /// A byte lookup event.
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
 pub struct ByteLookupEvent {
  /// The shard number, used for byte lookup table.
  pub shard: u32,
@@ -36,9 +34,10 @@ pub trait ByteRecord {
  fn add_byte_lookup_event(&mut self, blu_event: ByteLookupEvent);
 
  /// Adds a list of `ByteLookupEvent`s to the record.
+ #[inline]
  fn add_byte_lookup_events(&mut self, blu_events: Vec<ByteLookupEvent>) {
- for blu_event in blu_events.iter() {
- self.add_byte_lookup_event(*blu_event);
+ for blu_event in blu_events {
+ self.add_byte_lookup_event(blu_event);
  }
  }
 
@@ -121,6 +120,7 @@ pub trait ByteRecord {
 
 impl ByteLookupEvent {
  /// Creates a new `ByteLookupEvent`.
+ #[inline(always)]
  pub fn new(
  shard: u32,
  channel: u32,
@@ -147,13 +147,3 @@ impl ByteRecord for Vec<ByteLookupEvent> {
  self.push(blu_event);
  }
 }
-
-impl ByteRecord for BTreeMap<u32, BTreeMap<ByteLookupEvent, usize>> {
- fn add_byte_lookup_event(&mut self, blu_event: ByteLookupEvent) {
- *self
- .entry(blu_event.shard)
- .or_default()
- .entry(blu_event)
- .or_insert(0) += 1
- }
-}

core/src/bytes/mod.rs

@@ -5,7 +5,6 @@ pub mod opcode;
 pub mod trace;
 pub mod utils;
 
-use alloc::collections::BTreeMap;
 use core::borrow::BorrowMut;
 use std::marker::PhantomData;
 
@@ -35,19 +34,10 @@ pub const NUM_BYTE_LOOKUP_CHANNELS: u32 = 16;
 pub struct ByteChip<F>(PhantomData<F>);
 
 impl<F: Field> ByteChip<F> {
- /// Creates the preprocessed byte trace and event map.
+ /// Creates the preprocessed byte trace.
  ///
- /// This function returns a pair `(trace, map)`, where:
- /// - `trace` is a matrix containing all possible byte operations.
- /// - `map` is a map from a byte lookup to the corresponding row it appears in the table and
- /// the index of the result in the array of multiplicities.
- pub fn trace_and_map(
- shard: u32,
- ) -> (RowMajorMatrix<F>, BTreeMap<ByteLookupEvent, (usize, usize)>) {
- // A map from a byte lookup to its corresponding row in the table and index in the array of
- // multiplicities.
- let mut event_map = BTreeMap::new();
-
+ /// This function returns a `trace` which is a matrix containing all possible byte operations.
+ pub fn trace() -> RowMajorMatrix<F> {
  // The trace containing all values, with all multiplicities set to zero.
  let mut initial_trace = RowMajorMatrix::new(
  vec![F::zero(); NUM_ROWS * NUM_BYTE_PREPROCESSED_COLS],
@@ -65,10 +55,11 @@ impl<F: Field> ByteChip<F> {
  col.b = F::from_canonical_u8(b);
  col.c = F::from_canonical_u8(c);
 
+ let shard = 0;
  // Iterate over all operations for results and updating the table map.
  for channel in 0..NUM_BYTE_LOOKUP_CHANNELS {
- for (i, opcode) in opcodes.iter().enumerate() {
-let event = match opcode {
+ for opcode in opcodes.iter() {
+ match opcode {
  ByteOpcode::AND => {
  let and = b & c;
  col.and = F::from_canonical_u8(and);
@@ -176,11 +167,25 @@ impl<F: Field> ByteChip<F> {
  ByteLookupEvent::new(shard, channel, *opcode, v, 0, 0, 0)
  }
  };
- event_map.insert(event, (row_index, i));
  }
  }
  }
 
- (initial_trace, event_map)
+ initial_trace
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use p3_baby_bear::BabyBear;
+ use std::time::Instant;
+
+ use super::*;
+
+ #[test]
+ fn test_trace_and_map() {
+ let start = Instant::now();
+ ByteChip::<BabyBear>::trace();
+ println!("trace and map: {:?}", start.elapsed());
  }
 }

core/src/bytes/trace.rs

@@ -1,5 +1,6 @@
-use std::{borrow::BorrowMut, collections::BTreeMap};
+use std::borrow::BorrowMut;
 
+use hashbrown::HashMap;
 use p3_field::Field;
 use p3_matrix::dense::RowMajorMatrix;
 
@@ -9,14 +10,15 @@ use super::{
 };
 use crate::{
  air::{EventLens, MachineAir, WithEvents},
+ bytes::ByteOpcode,
  runtime::{ExecutionRecord, Program},
 };
 
 pub const NUM_ROWS: usize = 1 << 16;
 
 impl<'a, F: Field> WithEvents<'a> for ByteChip<F> {
  // the byte lookups
- type Events = &'a BTreeMap<u32, BTreeMap<ByteLookupEvent, usize>>;
+ type Events = &'a HashMap<u32, HashMap<ByteLookupEvent, usize>>;
 }
 
 impl<F: Field> MachineAir<F> for ByteChip<F> {
@@ -33,10 +35,7 @@ impl<F: Field> MachineAir<F> for ByteChip<F> {
  }
 
  fn generate_preprocessed_trace(&self, _program: &Self::Program) -> Option<RowMajorMatrix<F>> {
- // OPT: We should be able to make this a constant. Also, trace / map should be separate.
- // Since we only need the trace and not the map, we can just pass 0 as the shard.
- let (trace, _) = Self::trace_and_map(0);
-
+ let trace = Self::trace();
  Some(trace)
  }
 
@@ -53,23 +52,23 @@ impl<F: Field> MachineAir<F> for ByteChip<F> {
  input: &EL,
  _output: &mut ExecutionRecord,
  ) -> RowMajorMatrix<F> {
- let shard = input.index();
- let (_, event_map) = Self::trace_and_map(shard);
-
  let mut trace = RowMajorMatrix::new(
  vec![F::zero(); NUM_BYTE_MULT_COLS * NUM_ROWS],
  NUM_BYTE_MULT_COLS,
  );
 
+ let shard = input.index();
  for (lookup, mult) in input.events()[&shard].iter() {
- let (row, index) = event_map[lookup];
+ let row = if lookup.opcode != ByteOpcode::U16Range {
+ ((lookup.b << 8) + lookup.c) as usize
+ } else {
+ lookup.a1 as usize
+ };
+ let index = lookup.opcode as usize;
  let channel = lookup.channel as usize;
- let cols: &mut ByteMultCols<F> = trace.row_mut(row).borrow_mut();
 
-// Update the trace multiplicity
+let cols: &mut ByteMultCols<F> = trace.row_mut(row).borrow_mut();
  cols.mult_channels[channel].multiplicities[index] += F::from_canonical_usize(*mult);
-
- // Set the shard column as the current shard.
  cols.shard = F::from_canonical_u32(shard);
  }
 

core/src/cpu/columns/channel.rs

@@ -13,6 +13,7 @@ pub struct ChannelSelectorCols<T> {
 }
 
 impl<F: Field> ChannelSelectorCols<F> {
+ #[inline(always)]
  pub fn populate(&mut self, channel: u32) {
  self.channel_selectors = [F::zero(); NUM_BYTE_LOOKUP_CHANNELS as usize];
  self.channel_selectors[channel as usize] = F::one();

.gitignore

@@ -21,4 +21,7 @@ benchmark.csv
 # Build Artifacts
 recursion/gnark-ffi/build
 prover/build
-prover/*.tar.gz
+prover/*.tar.gz
+
+# IDE Conf
+.idea

Cargo.toml

@@ -74,7 +74,7 @@ ff = "0.13"
 futures = "0.3.30"
 futures-util = "0.3.14"
 getrandom = "=0.2.14" # 0.2.15 depends on yanked libc 0.2.154
-hashbrown = "0.14.5"
+hashbrown = { version = "0.14.5", features = ["serde"] }
 hex = "0.4.3"
 home = "0.5.9"
 hybrid-array = "0.2.0-rc"

core/Cargo.toml

@@ -42,6 +42,7 @@ bls12_381 = { workspace = true }
 cfg-if = { workspace = true }
 curve25519-dalek = { workspace = true }
 elliptic-curve = { workspace = true }
+hashbrown = { workspace = true }
 hex = { workspace = true }
 hybrid-array = { workspace = true }
 k256 = { workspace = true, features = ["expose-field"] }

core/benches/fibonacci.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use hashbrown::HashMap;
 
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
 use sphinx_core::{

core/src/alu/divrem/mod.rs

@@ -64,7 +64,7 @@ mod utils;
 
 use core::borrow::{Borrow, BorrowMut};
 use core::mem::size_of;
-use std::collections::HashMap;
+use hashbrown::HashMap;
 
 use p3_air::{Air, AirBuilder, BaseAir};
 use p3_field::AbstractField;
@@ -394,7 +394,7 @@ impl<F: PrimeField> MachineAir<F> for DivRemChip {
  }
  let mut alu_events = HashMap::new();
  alu_events.insert(Opcode::ADD, add_events);
- output.add_alu_events(&alu_events);
+ output.add_alu_events(&mut alu_events);
  }
 
  let mut lower_word = 0;

core/src/bytes/event.rs

@@ -1,12 +1,10 @@
-use std::collections::BTreeMap;
-
 use p3_field::PrimeField32;
 use serde::{Deserialize, Serialize};
 
 use super::ByteOpcode;
 
 /// A byte lookup event.
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
 pub struct ByteLookupEvent {
  /// The shard number, used for byte lookup table.
  pub shard: u32,
@@ -36,9 +34,10 @@ pub trait ByteRecord {
  fn add_byte_lookup_event(&mut self, blu_event: ByteLookupEvent);
 
  /// Adds a list of `ByteLookupEvent`s to the record.
+ #[inline]
  fn add_byte_lookup_events(&mut self, blu_events: Vec<ByteLookupEvent>) {
- for blu_event in blu_events.iter() {
- self.add_byte_lookup_event(*blu_event);
+ for blu_event in blu_events {
+ self.add_byte_lookup_event(blu_event);
  }
  }
 
@@ -121,6 +120,7 @@ pub trait ByteRecord {
 
 impl ByteLookupEvent {
  /// Creates a new `ByteLookupEvent`.
+ #[inline(always)]
  pub fn new(
  shard: u32,
  channel: u32,
@@ -147,13 +147,3 @@ impl ByteRecord for Vec<ByteLookupEvent> {
  self.push(blu_event);
  }
 }
-
-impl ByteRecord for BTreeMap<u32, BTreeMap<ByteLookupEvent, usize>> {
- fn add_byte_lookup_event(&mut self, blu_event: ByteLookupEvent) {
- *self
- .entry(blu_event.shard)
- .or_default()
- .entry(blu_event)
- .or_insert(0) += 1
- }
-}

core/src/bytes/mod.rs

@@ -5,7 +5,6 @@ pub mod opcode;
 pub mod trace;
 pub mod utils;
 
-use alloc::collections::BTreeMap;
 use core::borrow::BorrowMut;
 use std::marker::PhantomData;
 
@@ -35,19 +34,10 @@ pub const NUM_BYTE_LOOKUP_CHANNELS: u32 = 16;
 pub struct ByteChip<F>(PhantomData<F>);
 
 impl<F: Field> ByteChip<F> {
- /// Creates the preprocessed byte trace and event map.
+ /// Creates the preprocessed byte trace.
  ///
- /// This function returns a pair `(trace, map)`, where:
- /// - `trace` is a matrix containing all possible byte operations.
- /// - `map` is a map from a byte lookup to the corresponding row it appears in the table and
- /// the index of the result in the array of multiplicities.
- pub fn trace_and_map(
- shard: u32,
- ) -> (RowMajorMatrix<F>, BTreeMap<ByteLookupEvent, (usize, usize)>) {
- // A map from a byte lookup to its corresponding row in the table and index in the array of
- // multiplicities.
- let mut event_map = BTreeMap::new();
-
+ /// This function returns a `trace` which is a matrix containing all possible byte operations.
+ pub fn trace() -> RowMajorMatrix<F> {
  // The trace containing all values, with all multiplicities set to zero.
  let mut initial_trace = RowMajorMatrix::new(
  vec![F::zero(); NUM_ROWS * NUM_BYTE_PREPROCESSED_COLS],
@@ -65,10 +55,11 @@ impl<F: Field> ByteChip<F> {
  col.b = F::from_canonical_u8(b);
  col.c = F::from_canonical_u8(c);
 
+ let shard = 0;
  // Iterate over all operations for results and updating the table map.
  for channel in 0..NUM_BYTE_LOOKUP_CHANNELS {
- for (i, opcode) in opcodes.iter().enumerate() {
-let event = match opcode {
+ for opcode in opcodes.iter() {
+ match opcode {
  ByteOpcode::AND => {
  let and = b & c;
  col.and = F::from_canonical_u8(and);
@@ -176,11 +167,25 @@ impl<F: Field> ByteChip<F> {
  ByteLookupEvent::new(shard, channel, *opcode, v, 0, 0, 0)
  }
  };
- event_map.insert(event, (row_index, i));
  }
  }
  }
 
- (initial_trace, event_map)
+ initial_trace
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use p3_baby_bear::BabyBear;
+ use std::time::Instant;
+
+ use super::*;
+
+ #[test]
+ fn test_trace_and_map() {
+ let start = Instant::now();
+ ByteChip::<BabyBear>::trace();
+ println!("trace and map: {:?}", start.elapsed());
  }
 }

core/src/bytes/trace.rs

@@ -1,5 +1,6 @@
-use std::{borrow::BorrowMut, collections::BTreeMap};
+use std::borrow::BorrowMut;
 
+use hashbrown::HashMap;
 use p3_field::Field;
 use p3_matrix::dense::RowMajorMatrix;
 
@@ -9,14 +10,15 @@ use super::{
 };
 use crate::{
  air::{EventLens, MachineAir, WithEvents},
+ bytes::ByteOpcode,
  runtime::{ExecutionRecord, Program},
 };
 
 pub const NUM_ROWS: usize = 1 << 16;
 
 impl<'a, F: Field> WithEvents<'a> for ByteChip<F> {
  // the byte lookups
- type Events = &'a BTreeMap<u32, BTreeMap<ByteLookupEvent, usize>>;
+ type Events = &'a HashMap<u32, HashMap<ByteLookupEvent, usize>>;
 }
 
 impl<F: Field> MachineAir<F> for ByteChip<F> {
@@ -33,10 +35,7 @@ impl<F: Field> MachineAir<F> for ByteChip<F> {
  }
 
  fn generate_preprocessed_trace(&self, _program: &Self::Program) -> Option<RowMajorMatrix<F>> {
- // OPT: We should be able to make this a constant. Also, trace / map should be separate.
- // Since we only need the trace and not the map, we can just pass 0 as the shard.
- let (trace, _) = Self::trace_and_map(0);
-
+ let trace = Self::trace();
  Some(trace)
  }
 
@@ -53,23 +52,23 @@ impl<F: Field> MachineAir<F> for ByteChip<F> {
  input: &EL,
  _output: &mut ExecutionRecord,
  ) -> RowMajorMatrix<F> {
- let shard = input.index();
- let (_, event_map) = Self::trace_and_map(shard);
-
  let mut trace = RowMajorMatrix::new(
  vec![F::zero(); NUM_BYTE_MULT_COLS * NUM_ROWS],
  NUM_BYTE_MULT_COLS,
  );
 
+ let shard = input.index();
  for (lookup, mult) in input.events()[&shard].iter() {
- let (row, index) = event_map[lookup];
+ let row = if lookup.opcode != ByteOpcode::U16Range {
+ ((lookup.b << 8) + lookup.c) as usize
+ } else {
+ lookup.a1 as usize
+ };
+ let index = lookup.opcode as usize;
  let channel = lookup.channel as usize;
- let cols: &mut ByteMultCols<F> = trace.row_mut(row).borrow_mut();
 
-// Update the trace multiplicity
+let cols: &mut ByteMultCols<F> = trace.row_mut(row).borrow_mut();
  cols.mult_channels[channel].multiplicities[index] += F::from_canonical_usize(*mult);
-
- // Set the shard column as the current shard.
  cols.shard = F::from_canonical_u32(shard);
  }
 

core/src/cpu/columns/channel.rs

@@ -13,6 +13,7 @@ pub struct ChannelSelectorCols<T> {
 }
 
 impl<F: Field> ChannelSelectorCols<F> {
+ #[inline(always)]
  pub fn populate(&mut self, channel: u32) {
  self.channel_selectors = [F::zero(); NUM_BYTE_LOOKUP_CHANNELS as usize];
  self.channel_selectors[channel as usize] = F::one();

.gitignore

@@ -21,4 +21,7 @@ benchmark.csv
 # Build Artifacts
 recursion/gnark-ffi/build
 prover/build
-prover/*.tar.gz
+prover/*.tar.gz
+
+# IDE Conf
+.idea

Cargo.toml

@@ -74,7 +74,7 @@ ff = "0.13"
 futures = "0.3.30"
 futures-util = "0.3.14"
 getrandom = "=0.2.14" # 0.2.15 depends on yanked libc 0.2.154
-hashbrown = "0.14.5"
+hashbrown = { version = "0.14.5", features = ["serde"] }
 hex = "0.4.3"
 home = "0.5.9"
 hybrid-array = "0.2.0-rc"

core/Cargo.toml

@@ -42,6 +42,7 @@ bls12_381 = { workspace = true }
 cfg-if = { workspace = true }
 curve25519-dalek = { workspace = true }
 elliptic-curve = { workspace = true }
+hashbrown = { workspace = true }
 hex = { workspace = true }
 hybrid-array = { workspace = true }
 k256 = { workspace = true, features = ["expose-field"] }

core/benches/fibonacci.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use hashbrown::HashMap;
 
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
 use sphinx_core::{

core/src/alu/divrem/mod.rs

@@ -64,7 +64,7 @@ mod utils;
 
 use core::borrow::{Borrow, BorrowMut};
 use core::mem::size_of;
-use std::collections::HashMap;
+use hashbrown::HashMap;
 
 use p3_air::{Air, AirBuilder, BaseAir};
 use p3_field::AbstractField;
@@ -394,7 +394,7 @@ impl<F: PrimeField> MachineAir<F> for DivRemChip {
  }
  let mut alu_events = HashMap::new();
  alu_events.insert(Opcode::ADD, add_events);
- output.add_alu_events(&alu_events);
+ output.add_alu_events(&mut alu_events);
  }
 
  let mut lower_word = 0;

core/src/bytes/event.rs

@@ -1,12 +1,10 @@
-use std::collections::BTreeMap;
-
 use p3_field::PrimeField32;
 use serde::{Deserialize, Serialize};
 
 use super::ByteOpcode;
 
 /// A byte lookup event.
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
 pub struct ByteLookupEvent {
  /// The shard number, used for byte lookup table.
  pub shard: u32,
@@ -36,9 +34,10 @@ pub trait ByteRecord {
  fn add_byte_lookup_event(&mut self, blu_event: ByteLookupEvent);
 
  /// Adds a list of `ByteLookupEvent`s to the record.
+ #[inline]
  fn add_byte_lookup_events(&mut self, blu_events: Vec<ByteLookupEvent>) {
- for blu_event in blu_events.iter() {
- self.add_byte_lookup_event(*blu_event);
+ for blu_event in blu_events {
+ self.add_byte_lookup_event(blu_event);
  }
  }
 
@@ -121,6 +120,7 @@ pub trait ByteRecord {
 
 impl ByteLookupEvent {
  /// Creates a new `ByteLookupEvent`.
+ #[inline(always)]
  pub fn new(
  shard: u32,
  channel: u32,
@@ -147,13 +147,3 @@ impl ByteRecord for Vec<ByteLookupEvent> {
  self.push(blu_event);
  }
 }
-
-impl ByteRecord for BTreeMap<u32, BTreeMap<ByteLookupEvent, usize>> {
- fn add_byte_lookup_event(&mut self, blu_event: ByteLookupEvent) {
- *self
- .entry(blu_event.shard)
- .or_default()
- .entry(blu_event)
- .or_insert(0) += 1
- }
-}

core/src/bytes/mod.rs

@@ -5,7 +5,6 @@ pub mod opcode;
 pub mod trace;
 pub mod utils;
 
-use alloc::collections::BTreeMap;
 use core::borrow::BorrowMut;
 use std::marker::PhantomData;
 
@@ -35,19 +34,10 @@ pub const NUM_BYTE_LOOKUP_CHANNELS: u32 = 16;
 pub struct ByteChip<F>(PhantomData<F>);
 
 impl<F: Field> ByteChip<F> {
- /// Creates the preprocessed byte trace and event map.
+ /// Creates the preprocessed byte trace.
  ///
- /// This function returns a pair `(trace, map)`, where:
- /// - `trace` is a matrix containing all possible byte operations.
- /// - `map` is a map from a byte lookup to the corresponding row it appears in the table and
- /// the index of the result in the array of multiplicities.
- pub fn trace_and_map(
- shard: u32,
- ) -> (RowMajorMatrix<F>, BTreeMap<ByteLookupEvent, (usize, usize)>) {
- // A map from a byte lookup to its corresponding row in the table and index in the array of
- // multiplicities.
- let mut event_map = BTreeMap::new();
-
+ /// This function returns a `trace` which is a matrix containing all possible byte operations.
+ pub fn trace() -> RowMajorMatrix<F> {
  // The trace containing all values, with all multiplicities set to zero.
  let mut initial_trace = RowMajorMatrix::new(
  vec![F::zero(); NUM_ROWS * NUM_BYTE_PREPROCESSED_COLS],
@@ -65,10 +55,11 @@ impl<F: Field> ByteChip<F> {
  col.b = F::from_canonical_u8(b);
  col.c = F::from_canonical_u8(c);
 
+ let shard = 0;
  // Iterate over all operations for results and updating the table map.
  for channel in 0..NUM_BYTE_LOOKUP_CHANNELS {
- for (i, opcode) in opcodes.iter().enumerate() {
-let event = match opcode {
+ for opcode in opcodes.iter() {
+ match opcode {
  ByteOpcode::AND => {
  let and = b & c;
  col.and = F::from_canonical_u8(and);
@@ -176,11 +167,25 @@ impl<F: Field> ByteChip<F> {
  ByteLookupEvent::new(shard, channel, *opcode, v, 0, 0, 0)
  }
  };
- event_map.insert(event, (row_index, i));
  }
  }
  }
 
- (initial_trace, event_map)
+ initial_trace
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use p3_baby_bear::BabyBear;
+ use std::time::Instant;
+
+ use super::*;
+
+ #[test]
+ fn test_trace_and_map() {
+ let start = Instant::now();
+ ByteChip::<BabyBear>::trace();
+ println!("trace and map: {:?}", start.elapsed());
  }
 }

core/src/bytes/trace.rs

@@ -1,5 +1,6 @@
-use std::{borrow::BorrowMut, collections::BTreeMap};
+use std::borrow::BorrowMut;
 
+use hashbrown::HashMap;
 use p3_field::Field;
 use p3_matrix::dense::RowMajorMatrix;
 
@@ -9,14 +10,15 @@ use super::{
 };
 use crate::{
  air::{EventLens, MachineAir, WithEvents},
+ bytes::ByteOpcode,
  runtime::{ExecutionRecord, Program},
 };
 
 pub const NUM_ROWS: usize = 1 << 16;
 
 impl<'a, F: Field> WithEvents<'a> for ByteChip<F> {
  // the byte lookups
- type Events = &'a BTreeMap<u32, BTreeMap<ByteLookupEvent, usize>>;
+ type Events = &'a HashMap<u32, HashMap<ByteLookupEvent, usize>>;
 }
 
 impl<F: Field> MachineAir<F> for ByteChip<F> {
@@ -33,10 +35,7 @@ impl<F: Field> MachineAir<F> for ByteChip<F> {
  }
 
  fn generate_preprocessed_trace(&self, _program: &Self::Program) -> Option<RowMajorMatrix<F>> {
- // OPT: We should be able to make this a constant. Also, trace / map should be separate.
- // Since we only need the trace and not the map, we can just pass 0 as the shard.
- let (trace, _) = Self::trace_and_map(0);
-
+ let trace = Self::trace();
  Some(trace)
  }
 
@@ -53,23 +52,23 @@ impl<F: Field> MachineAir<F> for ByteChip<F> {
  input: &EL,
  _output: &mut ExecutionRecord,
  ) -> RowMajorMatrix<F> {
- let shard = input.index();
- let (_, event_map) = Self::trace_and_map(shard);
-
  let mut trace = RowMajorMatrix::new(
  vec![F::zero(); NUM_BYTE_MULT_COLS * NUM_ROWS],
  NUM_BYTE_MULT_COLS,
  );
 
+ let shard = input.index();
  for (lookup, mult) in input.events()[&shard].iter() {
- let (row, index) = event_map[lookup];
+ let row = if lookup.opcode != ByteOpcode::U16Range {
+ ((lookup.b << 8) + lookup.c) as usize
+ } else {
+ lookup.a1 as usize
+ };
+ let index = lookup.opcode as usize;
  let channel = lookup.channel as usize;
- let cols: &mut ByteMultCols<F> = trace.row_mut(row).borrow_mut();
 
-// Update the trace multiplicity
+let cols: &mut ByteMultCols<F> = trace.row_mut(row).borrow_mut();
  cols.mult_channels[channel].multiplicities[index] += F::from_canonical_usize(*mult);
-
- // Set the shard column as the current shard.
  cols.shard = F::from_canonical_u32(shard);
  }
 

core/src/cpu/columns/channel.rs

@@ -13,6 +13,7 @@ pub struct ChannelSelectorCols<T> {
 }
 
 impl<F: Field> ChannelSelectorCols<F> {
+ #[inline(always)]
  pub fn populate(&mut self, channel: u32) {
  self.channel_selectors = [F::zero(); NUM_BYTE_LOOKUP_CHANNELS as usize];
  self.channel_selectors[channel as usize] = F::one();

.gitignore

@@ -21,4 +21,7 @@ benchmark.csv
 # Build Artifacts
 recursion/gnark-ffi/build
 prover/build
-prover/*.tar.gz
+prover/*.tar.gz
+
+# IDE Conf
+.idea

Cargo.toml

@@ -74,7 +74,7 @@ ff = "0.13"
 futures = "0.3.30"
 futures-util = "0.3.14"
 getrandom = "=0.2.14" # 0.2.15 depends on yanked libc 0.2.154
-hashbrown = "0.14.5"
+hashbrown = { version = "0.14.5", features = ["serde"] }
 hex = "0.4.3"
 home = "0.5.9"
 hybrid-array = "0.2.0-rc"

core/Cargo.toml

@@ -42,6 +42,7 @@ bls12_381 = { workspace = true }
 cfg-if = { workspace = true }
 curve25519-dalek = { workspace = true }
 elliptic-curve = { workspace = true }
+hashbrown = { workspace = true }
 hex = { workspace = true }
 hybrid-array = { workspace = true }
 k256 = { workspace = true, features = ["expose-field"] }

core/benches/fibonacci.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use hashbrown::HashMap;
 
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
 use sphinx_core::{

core/src/alu/divrem/mod.rs

@@ -64,7 +64,7 @@ mod utils;
 
 use core::borrow::{Borrow, BorrowMut};
 use core::mem::size_of;
-use std::collections::HashMap;
+use hashbrown::HashMap;
 
 use p3_air::{Air, AirBuilder, BaseAir};
 use p3_field::AbstractField;
@@ -394,7 +394,7 @@ impl<F: PrimeField> MachineAir<F> for DivRemChip {
  }
  let mut alu_events = HashMap::new();
  alu_events.insert(Opcode::ADD, add_events);
- output.add_alu_events(&alu_events);
+ output.add_alu_events(&mut alu_events);
  }
 
  let mut lower_word = 0;

core/src/bytes/event.rs

@@ -1,12 +1,10 @@
-use std::collections::BTreeMap;
-
 use p3_field::PrimeField32;
 use serde::{Deserialize, Serialize};
 
 use super::ByteOpcode;
 
 /// A byte lookup event.
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
 pub struct ByteLookupEvent {
  /// The shard number, used for byte lookup table.
  pub shard: u32,
@@ -36,9 +34,10 @@ pub trait ByteRecord {
  fn add_byte_lookup_event(&mut self, blu_event: ByteLookupEvent);
 
  /// Adds a list of `ByteLookupEvent`s to the record.
+ #[inline]
  fn add_byte_lookup_events(&mut self, blu_events: Vec<ByteLookupEvent>) {
- for blu_event in blu_events.iter() {
- self.add_byte_lookup_event(*blu_event);
+ for blu_event in blu_events {
+ self.add_byte_lookup_event(blu_event);
  }
  }
 
@@ -121,6 +120,7 @@ pub trait ByteRecord {
 
 impl ByteLookupEvent {
  /// Creates a new `ByteLookupEvent`.
+ #[inline(always)]
  pub fn new(
  shard: u32,
  channel: u32,
@@ -147,13 +147,3 @@ impl ByteRecord for Vec<ByteLookupEvent> {
  self.push(blu_event);
  }
 }
-
-impl ByteRecord for BTreeMap<u32, BTreeMap<ByteLookupEvent, usize>> {
- fn add_byte_lookup_event(&mut self, blu_event: ByteLookupEvent) {
- *self
- .entry(blu_event.shard)
- .or_default()
- .entry(blu_event)
- .or_insert(0) += 1
- }
-}

core/src/bytes/mod.rs

@@ -5,7 +5,6 @@ pub mod opcode;
 pub mod trace;
 pub mod utils;
 
-use alloc::collections::BTreeMap;
 use core::borrow::BorrowMut;
 use std::marker::PhantomData;
 
@@ -35,19 +34,10 @@ pub const NUM_BYTE_LOOKUP_CHANNELS: u32 = 16;
 pub struct ByteChip<F>(PhantomData<F>);
 
 impl<F: Field> ByteChip<F> {
- /// Creates the preprocessed byte trace and event map.
+ /// Creates the preprocessed byte trace.
  ///
- /// This function returns a pair `(trace, map)`, where:
- /// - `trace` is a matrix containing all possible byte operations.
- /// - `map` is a map from a byte lookup to the corresponding row it appears in the table and
- /// the index of the result in the array of multiplicities.
- pub fn trace_and_map(
- shard: u32,
- ) -> (RowMajorMatrix<F>, BTreeMap<ByteLookupEvent, (usize, usize)>) {
- // A map from a byte lookup to its corresponding row in the table and index in the array of
- // multiplicities.
- let mut event_map = BTreeMap::new();
-
+ /// This function returns a `trace` which is a matrix containing all possible byte operations.
+ pub fn trace() -> RowMajorMatrix<F> {
  // The trace containing all values, with all multiplicities set to zero.
  let mut initial_trace = RowMajorMatrix::new(
  vec![F::zero(); NUM_ROWS * NUM_BYTE_PREPROCESSED_COLS],
@@ -65,10 +55,11 @@ impl<F: Field> ByteChip<F> {
  col.b = F::from_canonical_u8(b);
  col.c = F::from_canonical_u8(c);
 
+ let shard = 0;
  // Iterate over all operations for results and updating the table map.
  for channel in 0..NUM_BYTE_LOOKUP_CHANNELS {
- for (i, opcode) in opcodes.iter().enumerate() {
-let event = match opcode {
+ for opcode in opcodes.iter() {
+ match opcode {
  ByteOpcode::AND => {
  let and = b & c;
  col.and = F::from_canonical_u8(and);
@@ -176,11 +167,25 @@ impl<F: Field> ByteChip<F> {
  ByteLookupEvent::new(shard, channel, *opcode, v, 0, 0, 0)
  }
  };
- event_map.insert(event, (row_index, i));
  }
  }
  }
 
- (initial_trace, event_map)
+ initial_trace
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use p3_baby_bear::BabyBear;
+ use std::time::Instant;
+
+ use super::*;
+
+ #[test]
+ fn test_trace_and_map() {
+ let start = Instant::now();
+ ByteChip::<BabyBear>::trace();
+ println!("trace and map: {:?}", start.elapsed());
  }
 }

core/src/bytes/trace.rs

@@ -1,5 +1,6 @@
-use std::{borrow::BorrowMut, collections::BTreeMap};
+use std::borrow::BorrowMut;
 
+use hashbrown::HashMap;
 use p3_field::Field;
 use p3_matrix::dense::RowMajorMatrix;
 
@@ -9,14 +10,15 @@ use super::{
 };
 use crate::{
  air::{EventLens, MachineAir, WithEvents},
+ bytes::ByteOpcode,
  runtime::{ExecutionRecord, Program},
 };
 
 pub const NUM_ROWS: usize = 1 << 16;
 
 impl<'a, F: Field> WithEvents<'a> for ByteChip<F> {
  // the byte lookups
- type Events = &'a BTreeMap<u32, BTreeMap<ByteLookupEvent, usize>>;
+ type Events = &'a HashMap<u32, HashMap<ByteLookupEvent, usize>>;
 }
 
 impl<F: Field> MachineAir<F> for ByteChip<F> {
@@ -33,10 +35,7 @@ impl<F: Field> MachineAir<F> for ByteChip<F> {
  }
 
  fn generate_preprocessed_trace(&self, _program: &Self::Program) -> Option<RowMajorMatrix<F>> {
- // OPT: We should be able to make this a constant. Also, trace / map should be separate.
- // Since we only need the trace and not the map, we can just pass 0 as the shard.
- let (trace, _) = Self::trace_and_map(0);
-
+ let trace = Self::trace();
  Some(trace)
  }
 
@@ -53,23 +52,23 @@ impl<F: Field> MachineAir<F> for ByteChip<F> {
  input: &EL,
  _output: &mut ExecutionRecord,
  ) -> RowMajorMatrix<F> {
- let shard = input.index();
- let (_, event_map) = Self::trace_and_map(shard);
-
  let mut trace = RowMajorMatrix::new(
  vec![F::zero(); NUM_BYTE_MULT_COLS * NUM_ROWS],
  NUM_BYTE_MULT_COLS,
  );
 
+ let shard = input.index();
  for (lookup, mult) in input.events()[&shard].iter() {
- let (row, index) = event_map[lookup];
+ let row = if lookup.opcode != ByteOpcode::U16Range {
+ ((lookup.b << 8) + lookup.c) as usize
+ } else {
+ lookup.a1 as usize
+ };
+ let index = lookup.opcode as usize;
  let channel = lookup.channel as usize;
- let cols: &mut ByteMultCols<F> = trace.row_mut(row).borrow_mut();
 
-// Update the trace multiplicity
+let cols: &mut ByteMultCols<F> = trace.row_mut(row).borrow_mut();
  cols.mult_channels[channel].multiplicities[index] += F::from_canonical_usize(*mult);
-
- // Set the shard column as the current shard.
  cols.shard = F::from_canonical_u32(shard);
  }
 

core/src/cpu/columns/channel.rs

@@ -13,6 +13,7 @@ pub struct ChannelSelectorCols<T> {
 }
 
 impl<F: Field> ChannelSelectorCols<F> {
+ #[inline(always)]
  pub fn populate(&mut self, channel: u32) {
  self.channel_selectors = [F::zero(); NUM_BYTE_LOOKUP_CHANNELS as usize];
  self.channel_selectors[channel as usize] = F::one();

.gitignore

@@ -21,4 +21,7 @@ benchmark.csv
 # Build Artifacts
 recursion/gnark-ffi/build
 prover/build
-prover/*.tar.gz
+prover/*.tar.gz
+
+# IDE Conf
+.idea

Cargo.toml

@@ -74,7 +74,7 @@ ff = "0.13"
 futures = "0.3.30"
 futures-util = "0.3.14"
 getrandom = "=0.2.14" # 0.2.15 depends on yanked libc 0.2.154
-hashbrown = "0.14.5"
+hashbrown = { version = "0.14.5", features = ["serde"] }
 hex = "0.4.3"
 home = "0.5.9"
 hybrid-array = "0.2.0-rc"

core/Cargo.toml

@@ -42,6 +42,7 @@ bls12_381 = { workspace = true }
 cfg-if = { workspace = true }
 curve25519-dalek = { workspace = true }
 elliptic-curve = { workspace = true }
+hashbrown = { workspace = true }
 hex = { workspace = true }
 hybrid-array = { workspace = true }
 k256 = { workspace = true, features = ["expose-field"] }

core/benches/fibonacci.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use hashbrown::HashMap;
 
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
 use sphinx_core::{

core/src/alu/divrem/mod.rs

@@ -64,7 +64,7 @@ mod utils;
 
 use core::borrow::{Borrow, BorrowMut};
 use core::mem::size_of;
-use std::collections::HashMap;
+use hashbrown::HashMap;
 
 use p3_air::{Air, AirBuilder, BaseAir};
 use p3_field::AbstractField;
@@ -394,7 +394,7 @@ impl<F: PrimeField> MachineAir<F> for DivRemChip {
  }
  let mut alu_events = HashMap::new();
  alu_events.insert(Opcode::ADD, add_events);
- output.add_alu_events(&alu_events);
+ output.add_alu_events(&mut alu_events);
  }
 
  let mut lower_word = 0;

core/src/bytes/event.rs

@@ -1,12 +1,10 @@
-use std::collections::BTreeMap;
-
 use p3_field::PrimeField32;
 use serde::{Deserialize, Serialize};
 
 use super::ByteOpcode;
 
 /// A byte lookup event.
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
 pub struct ByteLookupEvent {
  /// The shard number, used for byte lookup table.
  pub shard: u32,
@@ -36,9 +34,10 @@ pub trait ByteRecord {
  fn add_byte_lookup_event(&mut self, blu_event: ByteLookupEvent);
 
  /// Adds a list of `ByteLookupEvent`s to the record.
+ #[inline]
  fn add_byte_lookup_events(&mut self, blu_events: Vec<ByteLookupEvent>) {
- for blu_event in blu_events.iter() {
- self.add_byte_lookup_event(*blu_event);
+ for blu_event in blu_events {
+ self.add_byte_lookup_event(blu_event);
  }
  }
 
@@ -121,6 +120,7 @@ pub trait ByteRecord {
 
 impl ByteLookupEvent {
  /// Creates a new `ByteLookupEvent`.
+ #[inline(always)]
  pub fn new(
  shard: u32,
  channel: u32,
@@ -147,13 +147,3 @@ impl ByteRecord for Vec<ByteLookupEvent> {
  self.push(blu_event);
  }
 }
-
-impl ByteRecord for BTreeMap<u32, BTreeMap<ByteLookupEvent, usize>> {
- fn add_byte_lookup_event(&mut self, blu_event: ByteLookupEvent) {
- *self
- .entry(blu_event.shard)
- .or_default()
- .entry(blu_event)
- .or_insert(0) += 1
- }
-}

core/src/bytes/mod.rs

@@ -5,7 +5,6 @@ pub mod opcode;
 pub mod trace;
 pub mod utils;
 
-use alloc::collections::BTreeMap;
 use core::borrow::BorrowMut;
 use std::marker::PhantomData;
 
@@ -35,19 +34,10 @@ pub const NUM_BYTE_LOOKUP_CHANNELS: u32 = 16;
 pub struct ByteChip<F>(PhantomData<F>);
 
 impl<F: Field> ByteChip<F> {
- /// Creates the preprocessed byte trace and event map.
+ /// Creates the preprocessed byte trace.
  ///
- /// This function returns a pair `(trace, map)`, where:
- /// - `trace` is a matrix containing all possible byte operations.
- /// - `map` is a map from a byte lookup to the corresponding row it appears in the table and
- /// the index of the result in the array of multiplicities.
- pub fn trace_and_map(
- shard: u32,
- ) -> (RowMajorMatrix<F>, BTreeMap<ByteLookupEvent, (usize, usize)>) {
- // A map from a byte lookup to its corresponding row in the table and index in the array of
- // multiplicities.
- let mut event_map = BTreeMap::new();
-
+ /// This function returns a `trace` which is a matrix containing all possible byte operations.
+ pub fn trace() -> RowMajorMatrix<F> {
  // The trace containing all values, with all multiplicities set to zero.
  let mut initial_trace = RowMajorMatrix::new(
  vec![F::zero(); NUM_ROWS * NUM_BYTE_PREPROCESSED_COLS],
@@ -65,10 +55,11 @@ impl<F: Field> ByteChip<F> {
  col.b = F::from_canonical_u8(b);
  col.c = F::from_canonical_u8(c);
 
+ let shard = 0;
  // Iterate over all operations for results and updating the table map.
  for channel in 0..NUM_BYTE_LOOKUP_CHANNELS {
- for (i, opcode) in opcodes.iter().enumerate() {
-let event = match opcode {
+ for opcode in opcodes.iter() {
+ match opcode {
  ByteOpcode::AND => {
  let and = b & c;
  col.and = F::from_canonical_u8(and);
@@ -176,11 +167,25 @@ impl<F: Field> ByteChip<F> {
  ByteLookupEvent::new(shard, channel, *opcode, v, 0, 0, 0)
  }
  };
- event_map.insert(event, (row_index, i));
  }
  }
  }
 
- (initial_trace, event_map)
+ initial_trace
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use p3_baby_bear::BabyBear;
+ use std::time::Instant;
+
+ use super::*;
+
+ #[test]
+ fn test_trace_and_map() {
+ let start = Instant::now();
+ ByteChip::<BabyBear>::trace();
+ println!("trace and map: {:?}", start.elapsed());
  }
 }

core/src/bytes/trace.rs

@@ -1,5 +1,6 @@
-use std::{borrow::BorrowMut, collections::BTreeMap};
+use std::borrow::BorrowMut;
 
+use hashbrown::HashMap;
 use p3_field::Field;
 use p3_matrix::dense::RowMajorMatrix;
 
@@ -9,14 +10,15 @@ use super::{
 };
 use crate::{
  air::{EventLens, MachineAir, WithEvents},
+ bytes::ByteOpcode,
  runtime::{ExecutionRecord, Program},
 };
 
 pub const NUM_ROWS: usize = 1 << 16;
 
 impl<'a, F: Field> WithEvents<'a> for ByteChip<F> {
  // the byte lookups
- type Events = &'a BTreeMap<u32, BTreeMap<ByteLookupEvent, usize>>;
+ type Events = &'a HashMap<u32, HashMap<ByteLookupEvent, usize>>;
 }
 
 impl<F: Field> MachineAir<F> for ByteChip<F> {
@@ -33,10 +35,7 @@ impl<F: Field> MachineAir<F> for ByteChip<F> {
  }
 
  fn generate_preprocessed_trace(&self, _program: &Self::Program) -> Option<RowMajorMatrix<F>> {
- // OPT: We should be able to make this a constant. Also, trace / map should be separate.
- // Since we only need the trace and not the map, we can just pass 0 as the shard.
- let (trace, _) = Self::trace_and_map(0);
-
+ let trace = Self::trace();
  Some(trace)
  }
 
@@ -53,23 +52,23 @@ impl<F: Field> MachineAir<F> for ByteChip<F> {
  input: &EL,
  _output: &mut ExecutionRecord,
  ) -> RowMajorMatrix<F> {
- let shard = input.index();
- let (_, event_map) = Self::trace_and_map(shard);
-
  let mut trace = RowMajorMatrix::new(
  vec![F::zero(); NUM_BYTE_MULT_COLS * NUM_ROWS],
  NUM_BYTE_MULT_COLS,
  );
 
+ let shard = input.index();
  for (lookup, mult) in input.events()[&shard].iter() {
- let (row, index) = event_map[lookup];
+ let row = if lookup.opcode != ByteOpcode::U16Range {
+ ((lookup.b << 8) + lookup.c) as usize
+ } else {
+ lookup.a1 as usize
+ };
+ let index = lookup.opcode as usize;
  let channel = lookup.channel as usize;
- let cols: &mut ByteMultCols<F> = trace.row_mut(row).borrow_mut();
 
-// Update the trace multiplicity
+let cols: &mut ByteMultCols<F> = trace.row_mut(row).borrow_mut();
  cols.mult_channels[channel].multiplicities[index] += F::from_canonical_usize(*mult);
-
- // Set the shard column as the current shard.
  cols.shard = F::from_canonical_u32(shard);
  }
 

core/src/cpu/columns/channel.rs

@@ -13,6 +13,7 @@ pub struct ChannelSelectorCols<T> {
 }
 
 impl<F: Field> ChannelSelectorCols<F> {
+ #[inline(always)]
  pub fn populate(&mut self, channel: u32) {
  self.channel_selectors = [F::zero(); NUM_BYTE_LOOKUP_CHANNELS as usize];
  self.channel_selectors[channel as usize] = F::one();

.gitignore

@@ -21,4 +21,7 @@ benchmark.csv
 # Build Artifacts
 recursion/gnark-ffi/build
 prover/build
-prover/*.tar.gz
+prover/*.tar.gz
+
+# IDE Conf
+.idea

Cargo.toml

@@ -74,7 +74,7 @@ ff = "0.13"
 futures = "0.3.30"
 futures-util = "0.3.14"
 getrandom = "=0.2.14" # 0.2.15 depends on yanked libc 0.2.154
-hashbrown = "0.14.5"
+hashbrown = { version = "0.14.5", features = ["serde"] }
 hex = "0.4.3"
 home = "0.5.9"
 hybrid-array = "0.2.0-rc"

core/Cargo.toml

@@ -42,6 +42,7 @@ bls12_381 = { workspace = true }
 cfg-if = { workspace = true }
 curve25519-dalek = { workspace = true }
 elliptic-curve = { workspace = true }
+hashbrown = { workspace = true }
 hex = { workspace = true }
 hybrid-array = { workspace = true }
 k256 = { workspace = true, features = ["expose-field"] }

core/benches/fibonacci.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use hashbrown::HashMap;
 
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
 use sphinx_core::{

core/src/alu/divrem/mod.rs

@@ -64,7 +64,7 @@ mod utils;
 
 use core::borrow::{Borrow, BorrowMut};
 use core::mem::size_of;
-use std::collections::HashMap;
+use hashbrown::HashMap;
 
 use p3_air::{Air, AirBuilder, BaseAir};
 use p3_field::AbstractField;
@@ -394,7 +394,7 @@ impl<F: PrimeField> MachineAir<F> for DivRemChip {
  }
  let mut alu_events = HashMap::new();
  alu_events.insert(Opcode::ADD, add_events);
- output.add_alu_events(&alu_events);
+ output.add_alu_events(&mut alu_events);
  }
 
  let mut lower_word = 0;

core/src/bytes/event.rs

@@ -1,12 +1,10 @@
-use std::collections::BTreeMap;
-
 use p3_field::PrimeField32;
 use serde::{Deserialize, Serialize};
 
 use super::ByteOpcode;
 
 /// A byte lookup event.
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
 pub struct ByteLookupEvent {
  /// The shard number, used for byte lookup table.
  pub shard: u32,
@@ -36,9 +34,10 @@ pub trait ByteRecord {
  fn add_byte_lookup_event(&mut self, blu_event: ByteLookupEvent);
 
  /// Adds a list of `ByteLookupEvent`s to the record.
+ #[inline]
  fn add_byte_lookup_events(&mut self, blu_events: Vec<ByteLookupEvent>) {
- for blu_event in blu_events.iter() {
- self.add_byte_lookup_event(*blu_event);
+ for blu_event in blu_events {
+ self.add_byte_lookup_event(blu_event);
  }
  }
 
@@ -121,6 +120,7 @@ pub trait ByteRecord {
 
 impl ByteLookupEvent {
  /// Creates a new `ByteLookupEvent`.
+ #[inline(always)]
  pub fn new(
  shard: u32,
  channel: u32,
@@ -147,13 +147,3 @@ impl ByteRecord for Vec<ByteLookupEvent> {
  self.push(blu_event);
  }
 }
-
-impl ByteRecord for BTreeMap<u32, BTreeMap<ByteLookupEvent, usize>> {
- fn add_byte_lookup_event(&mut self, blu_event: ByteLookupEvent) {
- *self
- .entry(blu_event.shard)
- .or_default()
- .entry(blu_event)
- .or_insert(0) += 1
- }
-}

core/src/bytes/mod.rs

@@ -5,7 +5,6 @@ pub mod opcode;
 pub mod trace;
 pub mod utils;
 
-use alloc::collections::BTreeMap;
 use core::borrow::BorrowMut;
 use std::marker::PhantomData;
 
@@ -35,19 +34,10 @@ pub const NUM_BYTE_LOOKUP_CHANNELS: u32 = 16;
 pub struct ByteChip<F>(PhantomData<F>);
 
 impl<F: Field> ByteChip<F> {
- /// Creates the preprocessed byte trace and event map.
+ /// Creates the preprocessed byte trace.
  ///
- /// This function returns a pair `(trace, map)`, where:
- /// - `trace` is a matrix containing all possible byte operations.
- /// - `map` is a map from a byte lookup to the corresponding row it appears in the table and
- /// the index of the result in the array of multiplicities.
- pub fn trace_and_map(
- shard: u32,
- ) -> (RowMajorMatrix<F>, BTreeMap<ByteLookupEvent, (usize, usize)>) {
- // A map from a byte lookup to its corresponding row in the table and index in the array of
- // multiplicities.
- let mut event_map = BTreeMap::new();
-
+ /// This function returns a `trace` which is a matrix containing all possible byte operations.
+ pub fn trace() -> RowMajorMatrix<F> {
  // The trace containing all values, with all multiplicities set to zero.
  let mut initial_trace = RowMajorMatrix::new(
  vec![F::zero(); NUM_ROWS * NUM_BYTE_PREPROCESSED_COLS],
@@ -65,10 +55,11 @@ impl<F: Field> ByteChip<F> {
  col.b = F::from_canonical_u8(b);
  col.c = F::from_canonical_u8(c);
 
+ let shard = 0;
  // Iterate over all operations for results and updating the table map.
  for channel in 0..NUM_BYTE_LOOKUP_CHANNELS {
- for (i, opcode) in opcodes.iter().enumerate() {
-let event = match opcode {
+ for opcode in opcodes.iter() {
+ match opcode {
  ByteOpcode::AND => {
  let and = b & c;
  col.and = F::from_canonical_u8(and);
@@ -176,11 +167,25 @@ impl<F: Field> ByteChip<F> {
  ByteLookupEvent::new(shard, channel, *opcode, v, 0, 0, 0)
  }
  };
- event_map.insert(event, (row_index, i));
  }
  }
  }
 
- (initial_trace, event_map)
+ initial_trace
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use p3_baby_bear::BabyBear;
+ use std::time::Instant;
+
+ use super::*;
+
+ #[test]
+ fn test_trace_and_map() {
+ let start = Instant::now();
+ ByteChip::<BabyBear>::trace();
+ println!("trace and map: {:?}", start.elapsed());
  }
 }

core/src/bytes/trace.rs

@@ -1,5 +1,6 @@
-use std::{borrow::BorrowMut, collections::BTreeMap};
+use std::borrow::BorrowMut;
 
+use hashbrown::HashMap;
 use p3_field::Field;
 use p3_matrix::dense::RowMajorMatrix;
 
@@ -9,14 +10,15 @@ use super::{
 };
 use crate::{
  air::{EventLens, MachineAir, WithEvents},
+ bytes::ByteOpcode,
  runtime::{ExecutionRecord, Program},
 };
 
 pub const NUM_ROWS: usize = 1 << 16;
 
 impl<'a, F: Field> WithEvents<'a> for ByteChip<F> {
  // the byte lookups
- type Events = &'a BTreeMap<u32, BTreeMap<ByteLookupEvent, usize>>;
+ type Events = &'a HashMap<u32, HashMap<ByteLookupEvent, usize>>;
 }
 
 impl<F: Field> MachineAir<F> for ByteChip<F> {
@@ -33,10 +35,7 @@ impl<F: Field> MachineAir<F> for ByteChip<F> {
  }
 
  fn generate_preprocessed_trace(&self, _program: &Self::Program) -> Option<RowMajorMatrix<F>> {
- // OPT: We should be able to make this a constant. Also, trace / map should be separate.
- // Since we only need the trace and not the map, we can just pass 0 as the shard.
- let (trace, _) = Self::trace_and_map(0);
-
+ let trace = Self::trace();
  Some(trace)
  }
 
@@ -53,23 +52,23 @@ impl<F: Field> MachineAir<F> for ByteChip<F> {
  input: &EL,
  _output: &mut ExecutionRecord,
  ) -> RowMajorMatrix<F> {
- let shard = input.index();
- let (_, event_map) = Self::trace_and_map(shard);
-
  let mut trace = RowMajorMatrix::new(
  vec![F::zero(); NUM_BYTE_MULT_COLS * NUM_ROWS],
  NUM_BYTE_MULT_COLS,
  );
 
+ let shard = input.index();
  for (lookup, mult) in input.events()[&shard].iter() {
- let (row, index) = event_map[lookup];
+ let row = if lookup.opcode != ByteOpcode::U16Range {
+ ((lookup.b << 8) + lookup.c) as usize
+ } else {
+ lookup.a1 as usize
+ };
+ let index = lookup.opcode as usize;
  let channel = lookup.channel as usize;
- let cols: &mut ByteMultCols<F> = trace.row_mut(row).borrow_mut();
 
-// Update the trace multiplicity
+let cols: &mut ByteMultCols<F> = trace.row_mut(row).borrow_mut();
  cols.mult_channels[channel].multiplicities[index] += F::from_canonical_usize(*mult);
-
- // Set the shard column as the current shard.
  cols.shard = F::from_canonical_u32(shard);
  }
 

core/src/cpu/columns/channel.rs

@@ -13,6 +13,7 @@ pub struct ChannelSelectorCols<T> {
 }
 
 impl<F: Field> ChannelSelectorCols<F> {
+ #[inline(always)]
  pub fn populate(&mut self, channel: u32) {
  self.channel_selectors = [F::zero(); NUM_BYTE_LOOKUP_CHANNELS as usize];
  self.channel_selectors[channel as usize] = F::one();

.gitignore

@@ -21,4 +21,7 @@ benchmark.csv
 # Build Artifacts
 recursion/gnark-ffi/build
 prover/build
-prover/*.tar.gz
+prover/*.tar.gz
+
+# IDE Conf
+.idea

Cargo.toml

@@ -74,7 +74,7 @@ ff = "0.13"
 futures = "0.3.30"
 futures-util = "0.3.14"
 getrandom = "=0.2.14" # 0.2.15 depends on yanked libc 0.2.154
-hashbrown = "0.14.5"
+hashbrown = { version = "0.14.5", features = ["serde"] }
 hex = "0.4.3"
 home = "0.5.9"
 hybrid-array = "0.2.0-rc"

core/Cargo.toml

@@ -42,6 +42,7 @@ bls12_381 = { workspace = true }
 cfg-if = { workspace = true }
 curve25519-dalek = { workspace = true }
 elliptic-curve = { workspace = true }
+hashbrown = { workspace = true }
 hex = { workspace = true }
 hybrid-array = { workspace = true }
 k256 = { workspace = true, features = ["expose-field"] }

core/benches/fibonacci.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use hashbrown::HashMap;
 
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
 use sphinx_core::{

core/src/alu/divrem/mod.rs

@@ -64,7 +64,7 @@ mod utils;
 
 use core::borrow::{Borrow, BorrowMut};
 use core::mem::size_of;
-use std::collections::HashMap;
+use hashbrown::HashMap;
 
 use p3_air::{Air, AirBuilder, BaseAir};
 use p3_field::AbstractField;
@@ -394,7 +394,7 @@ impl<F: PrimeField> MachineAir<F> for DivRemChip {
  }
  let mut alu_events = HashMap::new();
  alu_events.insert(Opcode::ADD, add_events);
- output.add_alu_events(&alu_events);
+ output.add_alu_events(&mut alu_events);
  }
 
  let mut lower_word = 0;

core/src/bytes/event.rs

@@ -1,12 +1,10 @@
-use std::collections::BTreeMap;
-
 use p3_field::PrimeField32;
 use serde::{Deserialize, Serialize};
 
 use super::ByteOpcode;
 
 /// A byte lookup event.
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
 pub struct ByteLookupEvent {
  /// The shard number, used for byte lookup table.
  pub shard: u32,
@@ -36,9 +34,10 @@ pub trait ByteRecord {
  fn add_byte_lookup_event(&mut self, blu_event: ByteLookupEvent);
 
  /// Adds a list of `ByteLookupEvent`s to the record.
+ #[inline]
  fn add_byte_lookup_events(&mut self, blu_events: Vec<ByteLookupEvent>) {
- for blu_event in blu_events.iter() {
- self.add_byte_lookup_event(*blu_event);
+ for blu_event in blu_events {
+ self.add_byte_lookup_event(blu_event);
  }
  }
 
@@ -121,6 +120,7 @@ pub trait ByteRecord {
 
 impl ByteLookupEvent {
  /// Creates a new `ByteLookupEvent`.
+ #[inline(always)]
  pub fn new(
  shard: u32,
  channel: u32,
@@ -147,13 +147,3 @@ impl ByteRecord for Vec<ByteLookupEvent> {
  self.push(blu_event);
  }
 }
-
-impl ByteRecord for BTreeMap<u32, BTreeMap<ByteLookupEvent, usize>> {
- fn add_byte_lookup_event(&mut self, blu_event: ByteLookupEvent) {
- *self
- .entry(blu_event.shard)
- .or_default()
- .entry(blu_event)
- .or_insert(0) += 1
- }
-}

core/src/bytes/mod.rs

@@ -5,7 +5,6 @@ pub mod opcode;
 pub mod trace;
 pub mod utils;
 
-use alloc::collections::BTreeMap;
 use core::borrow::BorrowMut;
 use std::marker::PhantomData;
 
@@ -35,19 +34,10 @@ pub const NUM_BYTE_LOOKUP_CHANNELS: u32 = 16;
 pub struct ByteChip<F>(PhantomData<F>);
 
 impl<F: Field> ByteChip<F> {
- /// Creates the preprocessed byte trace and event map.
+ /// Creates the preprocessed byte trace.
  ///
- /// This function returns a pair `(trace, map)`, where:
- /// - `trace` is a matrix containing all possible byte operations.
- /// - `map` is a map from a byte lookup to the corresponding row it appears in the table and
- /// the index of the result in the array of multiplicities.
- pub fn trace_and_map(
- shard: u32,
- ) -> (RowMajorMatrix<F>, BTreeMap<ByteLookupEvent, (usize, usize)>) {
- // A map from a byte lookup to its corresponding row in the table and index in the array of
- // multiplicities.
- let mut event_map = BTreeMap::new();
-
+ /// This function returns a `trace` which is a matrix containing all possible byte operations.
+ pub fn trace() -> RowMajorMatrix<F> {
  // The trace containing all values, with all multiplicities set to zero.
  let mut initial_trace = RowMajorMatrix::new(
  vec![F::zero(); NUM_ROWS * NUM_BYTE_PREPROCESSED_COLS],
@@ -65,10 +55,11 @@ impl<F: Field> ByteChip<F> {
  col.b = F::from_canonical_u8(b);
  col.c = F::from_canonical_u8(c);
 
+ let shard = 0;
  // Iterate over all operations for results and updating the table map.
  for channel in 0..NUM_BYTE_LOOKUP_CHANNELS {
- for (i, opcode) in opcodes.iter().enumerate() {
-let event = match opcode {
+ for opcode in opcodes.iter() {
+ match opcode {
  ByteOpcode::AND => {
  let and = b & c;
  col.and = F::from_canonical_u8(and);
@@ -176,11 +167,25 @@ impl<F: Field> ByteChip<F> {
  ByteLookupEvent::new(shard, channel, *opcode, v, 0, 0, 0)
  }
  };
- event_map.insert(event, (row_index, i));
  }
  }
  }
 
- (initial_trace, event_map)
+ initial_trace
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use p3_baby_bear::BabyBear;
+ use std::time::Instant;
+
+ use super::*;
+
+ #[test]
+ fn test_trace_and_map() {
+ let start = Instant::now();
+ ByteChip::<BabyBear>::trace();
+ println!("trace and map: {:?}", start.elapsed());
  }
 }

core/src/bytes/trace.rs

@@ -1,5 +1,6 @@
-use std::{borrow::BorrowMut, collections::BTreeMap};
+use std::borrow::BorrowMut;
 
+use hashbrown::HashMap;
 use p3_field::Field;
 use p3_matrix::dense::RowMajorMatrix;
 
@@ -9,14 +10,15 @@ use super::{
 };
 use crate::{
  air::{EventLens, MachineAir, WithEvents},
+ bytes::ByteOpcode,
  runtime::{ExecutionRecord, Program},
 };
 
 pub const NUM_ROWS: usize = 1 << 16;
 
 impl<'a, F: Field> WithEvents<'a> for ByteChip<F> {
  // the byte lookups
- type Events = &'a BTreeMap<u32, BTreeMap<ByteLookupEvent, usize>>;
+ type Events = &'a HashMap<u32, HashMap<ByteLookupEvent, usize>>;
 }
 
 impl<F: Field> MachineAir<F> for ByteChip<F> {
@@ -33,10 +35,7 @@ impl<F: Field> MachineAir<F> for ByteChip<F> {
  }
 
  fn generate_preprocessed_trace(&self, _program: &Self::Program) -> Option<RowMajorMatrix<F>> {
- // OPT: We should be able to make this a constant. Also, trace / map should be separate.
- // Since we only need the trace and not the map, we can just pass 0 as the shard.
- let (trace, _) = Self::trace_and_map(0);
-
+ let trace = Self::trace();
  Some(trace)
  }
 
@@ -53,23 +52,23 @@ impl<F: Field> MachineAir<F> for ByteChip<F> {
  input: &EL,
  _output: &mut ExecutionRecord,
  ) -> RowMajorMatrix<F> {
- let shard = input.index();
- let (_, event_map) = Self::trace_and_map(shard);
-
  let mut trace = RowMajorMatrix::new(
  vec![F::zero(); NUM_BYTE_MULT_COLS * NUM_ROWS],
  NUM_BYTE_MULT_COLS,
  );
 
+ let shard = input.index();
  for (lookup, mult) in input.events()[&shard].iter() {
- let (row, index) = event_map[lookup];
+ let row = if lookup.opcode != ByteOpcode::U16Range {
+ ((lookup.b << 8) + lookup.c) as usize
+ } else {
+ lookup.a1 as usize
+ };
+ let index = lookup.opcode as usize;
  let channel = lookup.channel as usize;
- let cols: &mut ByteMultCols<F> = trace.row_mut(row).borrow_mut();
 
-// Update the trace multiplicity
+let cols: &mut ByteMultCols<F> = trace.row_mut(row).borrow_mut();
  cols.mult_channels[channel].multiplicities[index] += F::from_canonical_usize(*mult);
-
- // Set the shard column as the current shard.
  cols.shard = F::from_canonical_u32(shard);
  }
 

core/src/cpu/columns/channel.rs

@@ -13,6 +13,7 @@ pub struct ChannelSelectorCols<T> {
 }
 
 impl<F: Field> ChannelSelectorCols<F> {
+ #[inline(always)]
  pub fn populate(&mut self, channel: u32) {
  self.channel_selectors = [F::zero(); NUM_BYTE_LOOKUP_CHANNELS as usize];
  self.channel_selectors[channel as usize] = F::one();

.gitignore

@@ -21,4 +21,7 @@ benchmark.csv
 # Build Artifacts
 recursion/gnark-ffi/build
 prover/build
-prover/*.tar.gz
+prover/*.tar.gz
+
+# IDE Conf
+.idea

Cargo.toml

@@ -74,7 +74,7 @@ ff = "0.13"
 futures = "0.3.30"
 futures-util = "0.3.14"
 getrandom = "=0.2.14" # 0.2.15 depends on yanked libc 0.2.154
-hashbrown = "0.14.5"
+hashbrown = { version = "0.14.5", features = ["serde"] }
 hex = "0.4.3"
 home = "0.5.9"
 hybrid-array = "0.2.0-rc"

core/Cargo.toml

@@ -42,6 +42,7 @@ bls12_381 = { workspace = true }
 cfg-if = { workspace = true }
 curve25519-dalek = { workspace = true }
 elliptic-curve = { workspace = true }
+hashbrown = { workspace = true }
 hex = { workspace = true }
 hybrid-array = { workspace = true }
 k256 = { workspace = true, features = ["expose-field"] }

core/benches/fibonacci.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use hashbrown::HashMap;
 
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
 use sphinx_core::{

core/src/alu/divrem/mod.rs

@@ -64,7 +64,7 @@ mod utils;
 
 use core::borrow::{Borrow, BorrowMut};
 use core::mem::size_of;
-use std::collections::HashMap;
+use hashbrown::HashMap;
 
 use p3_air::{Air, AirBuilder, BaseAir};
 use p3_field::AbstractField;
@@ -394,7 +394,7 @@ impl<F: PrimeField> MachineAir<F> for DivRemChip {
  }
  let mut alu_events = HashMap::new();
  alu_events.insert(Opcode::ADD, add_events);
- output.add_alu_events(&alu_events);
+ output.add_alu_events(&mut alu_events);
  }
 
  let mut lower_word = 0;

core/src/bytes/event.rs

@@ -1,12 +1,10 @@
-use std::collections::BTreeMap;
-
 use p3_field::PrimeField32;
 use serde::{Deserialize, Serialize};
 
 use super::ByteOpcode;
 
 /// A byte lookup event.
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
 pub struct ByteLookupEvent {
  /// The shard number, used for byte lookup table.
  pub shard: u32,
@@ -36,9 +34,10 @@ pub trait ByteRecord {
  fn add_byte_lookup_event(&mut self, blu_event: ByteLookupEvent);
 
  /// Adds a list of `ByteLookupEvent`s to the record.
+ #[inline]
  fn add_byte_lookup_events(&mut self, blu_events: Vec<ByteLookupEvent>) {
- for blu_event in blu_events.iter() {
- self.add_byte_lookup_event(*blu_event);
+ for blu_event in blu_events {
+ self.add_byte_lookup_event(blu_event);
  }
  }
 
@@ -121,6 +120,7 @@ pub trait ByteRecord {
 
 impl ByteLookupEvent {
  /// Creates a new `ByteLookupEvent`.
+ #[inline(always)]
  pub fn new(
  shard: u32,
  channel: u32,
@@ -147,13 +147,3 @@ impl ByteRecord for Vec<ByteLookupEvent> {
  self.push(blu_event);
  }
 }
-
-impl ByteRecord for BTreeMap<u32, BTreeMap<ByteLookupEvent, usize>> {
- fn add_byte_lookup_event(&mut self, blu_event: ByteLookupEvent) {
- *self
- .entry(blu_event.shard)
- .or_default()
- .entry(blu_event)
- .or_insert(0) += 1
- }
-}

core/src/bytes/mod.rs

@@ -5,7 +5,6 @@ pub mod opcode;
 pub mod trace;
 pub mod utils;
 
-use alloc::collections::BTreeMap;
 use core::borrow::BorrowMut;
 use std::marker::PhantomData;
 
@@ -35,19 +34,10 @@ pub const NUM_BYTE_LOOKUP_CHANNELS: u32 = 16;
 pub struct ByteChip<F>(PhantomData<F>);
 
 impl<F: Field> ByteChip<F> {
- /// Creates the preprocessed byte trace and event map.
+ /// Creates the preprocessed byte trace.
  ///
- /// This function returns a pair `(trace, map)`, where:
- /// - `trace` is a matrix containing all possible byte operations.
- /// - `map` is a map from a byte lookup to the corresponding row it appears in the table and
- /// the index of the result in the array of multiplicities.
- pub fn trace_and_map(
- shard: u32,
- ) -> (RowMajorMatrix<F>, BTreeMap<ByteLookupEvent, (usize, usize)>) {
- // A map from a byte lookup to its corresponding row in the table and index in the array of
- // multiplicities.
- let mut event_map = BTreeMap::new();
-
+ /// This function returns a `trace` which is a matrix containing all possible byte operations.
+ pub fn trace() -> RowMajorMatrix<F> {
  // The trace containing all values, with all multiplicities set to zero.
  let mut initial_trace = RowMajorMatrix::new(
  vec![F::zero(); NUM_ROWS * NUM_BYTE_PREPROCESSED_COLS],
@@ -65,10 +55,11 @@ impl<F: Field> ByteChip<F> {
  col.b = F::from_canonical_u8(b);
  col.c = F::from_canonical_u8(c);
 
+ let shard = 0;
  // Iterate over all operations for results and updating the table map.
  for channel in 0..NUM_BYTE_LOOKUP_CHANNELS {
- for (i, opcode) in opcodes.iter().enumerate() {
-let event = match opcode {
+ for opcode in opcodes.iter() {
+ match opcode {
  ByteOpcode::AND => {
  let and = b & c;
  col.and = F::from_canonical_u8(and);
@@ -176,11 +167,25 @@ impl<F: Field> ByteChip<F> {
  ByteLookupEvent::new(shard, channel, *opcode, v, 0, 0, 0)
  }
  };
- event_map.insert(event, (row_index, i));
  }
  }
  }
 
- (initial_trace, event_map)
+ initial_trace
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use p3_baby_bear::BabyBear;
+ use std::time::Instant;
+
+ use super::*;
+
+ #[test]
+ fn test_trace_and_map() {
+ let start = Instant::now();
+ ByteChip::<BabyBear>::trace();
+ println!("trace and map: {:?}", start.elapsed());
  }
 }

core/src/bytes/trace.rs

@@ -1,5 +1,6 @@
-use std::{borrow::BorrowMut, collections::BTreeMap};
+use std::borrow::BorrowMut;
 
+use hashbrown::HashMap;
 use p3_field::Field;
 use p3_matrix::dense::RowMajorMatrix;
 
@@ -9,14 +10,15 @@ use super::{
 };
 use crate::{
  air::{EventLens, MachineAir, WithEvents},
+ bytes::ByteOpcode,
  runtime::{ExecutionRecord, Program},
 };
 
 pub const NUM_ROWS: usize = 1 << 16;
 
 impl<'a, F: Field> WithEvents<'a> for ByteChip<F> {
  // the byte lookups
- type Events = &'a BTreeMap<u32, BTreeMap<ByteLookupEvent, usize>>;
+ type Events = &'a HashMap<u32, HashMap<ByteLookupEvent, usize>>;
 }
 
 impl<F: Field> MachineAir<F> for ByteChip<F> {
@@ -33,10 +35,7 @@ impl<F: Field> MachineAir<F> for ByteChip<F> {
  }
 
  fn generate_preprocessed_trace(&self, _program: &Self::Program) -> Option<RowMajorMatrix<F>> {
- // OPT: We should be able to make this a constant. Also, trace / map should be separate.
- // Since we only need the trace and not the map, we can just pass 0 as the shard.
- let (trace, _) = Self::trace_and_map(0);
-
+ let trace = Self::trace();
  Some(trace)
  }
 
@@ -53,23 +52,23 @@ impl<F: Field> MachineAir<F> for ByteChip<F> {
  input: &EL,
  _output: &mut ExecutionRecord,
  ) -> RowMajorMatrix<F> {
- let shard = input.index();
- let (_, event_map) = Self::trace_and_map(shard);
-
  let mut trace = RowMajorMatrix::new(
  vec![F::zero(); NUM_BYTE_MULT_COLS * NUM_ROWS],
  NUM_BYTE_MULT_COLS,
  );
 
+ let shard = input.index();
  for (lookup, mult) in input.events()[&shard].iter() {
- let (row, index) = event_map[lookup];
+ let row = if lookup.opcode != ByteOpcode::U16Range {
+ ((lookup.b << 8) + lookup.c) as usize
+ } else {
+ lookup.a1 as usize
+ };
+ let index = lookup.opcode as usize;
  let channel = lookup.channel as usize;
- let cols: &mut ByteMultCols<F> = trace.row_mut(row).borrow_mut();
 
-// Update the trace multiplicity
+let cols: &mut ByteMultCols<F> = trace.row_mut(row).borrow_mut();
  cols.mult_channels[channel].multiplicities[index] += F::from_canonical_usize(*mult);
-
- // Set the shard column as the current shard.
  cols.shard = F::from_canonical_u32(shard);
  }
 

core/src/cpu/columns/channel.rs

@@ -13,6 +13,7 @@ pub struct ChannelSelectorCols<T> {
 }
 
 impl<F: Field> ChannelSelectorCols<F> {
+ #[inline(always)]
  pub fn populate(&mut self, channel: u32) {
  self.channel_selectors = [F::zero(); NUM_BYTE_LOOKUP_CHANNELS as usize];
  self.channel_selectors[channel as usize] = F::one();