Skip to content

Commit 3b0a82a

Browse files
AntoineF4C5AntoineF4C5
authored
chore: refactor prefix checkpoints handling (#249)
* chore: refactor prefix checkpoints handling * chore: address comments * chore: further address comments
1 parent ce0a265 commit 3b0a82a

12 files changed

Lines changed: 93 additions & 43 deletions

File tree

‎joltworks/src/lookup_tables/prefixes/and.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use super::{PrefixCheckpoint, Prefixes, SparseDensePrefix};
1+
use super::{PrefixCheckpoint, PrefixCheckpoints, Prefixes, SparseDensePrefix};
22
use crate::{
33
field::{ChallengeFieldOps, FieldChallengeOps, JoltField},
44
utils::lookup_bits::LookupBits,
@@ -9,7 +9,7 @@ pub enum AndPrefix<const XLEN: usize> {}
99

1010
impl<const XLEN: usize, F: JoltField> SparseDensePrefix<F> for AndPrefix<XLEN> {
1111
fn prefix_mle<C>(
12-
checkpoints: &[PrefixCheckpoint<F>],
12+
checkpoints: &PrefixCheckpoints<F>,
1313
r_x: Option<C>,
1414
c: u32,
1515
mut b: LookupBits,
@@ -40,7 +40,7 @@ impl<const XLEN: usize, F: JoltField> SparseDensePrefix<F> for AndPrefix<XLEN> {
4040
}
4141

4242
fn update_prefix_checkpoint<C>(
43-
checkpoints: &[PrefixCheckpoint<F>],
43+
checkpoints: &PrefixCheckpoints<F>,
4444
r_x: C,
4545
r_y: C,
4646
j: usize,

‎joltworks/src/lookup_tables/prefixes/eq.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,14 +3,14 @@ use crate::{
33
utils::lookup_bits::LookupBits,
44
};
55

6-
use super::{PrefixCheckpoint, Prefixes, SparseDensePrefix};
6+
use super::{PrefixCheckpoint, PrefixCheckpoints, Prefixes, SparseDensePrefix};
77

88
/// Prefix component for equality comparison lookup table decomposition.
99
pub enum EqPrefix {}
1010

1111
impl<F: JoltField> SparseDensePrefix<F> for EqPrefix {
1212
fn prefix_mle<C>(
13-
checkpoints: &[PrefixCheckpoint<F>],
13+
checkpoints: &PrefixCheckpoints<F>,
1414
r_x: Option<C>,
1515
c: u32,
1616
mut b: LookupBits,
@@ -40,7 +40,7 @@ impl<F: JoltField> SparseDensePrefix<F> for EqPrefix {
4040
}
4141

4242
fn update_prefix_checkpoint<C>(
43-
checkpoints: &[PrefixCheckpoint<F>],
43+
checkpoints: &PrefixCheckpoints<F>,
4444
r_x: C,
4545
r_y: C,
4646
_: usize,

‎joltworks/src/lookup_tables/prefixes/less_than.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,14 +3,14 @@ use crate::{
33
utils::lookup_bits::LookupBits,
44
};
55

6-
use super::{PrefixCheckpoint, Prefixes, SparseDensePrefix};
6+
use super::{PrefixCheckpoint, PrefixCheckpoints, Prefixes, SparseDensePrefix};
77

88
/// Prefix component for less-than comparison lookup table decomposition.
99
pub enum LessThanPrefix {}
1010

1111
impl<F: JoltField> SparseDensePrefix<F> for LessThanPrefix {
1212
fn prefix_mle<C>(
13-
checkpoints: &[PrefixCheckpoint<F>],
13+
checkpoints: &PrefixCheckpoints<F>,
1414
r_x: Option<C>,
1515
c: u32,
1616
mut b: LookupBits,
@@ -46,7 +46,7 @@ impl<F: JoltField> SparseDensePrefix<F> for LessThanPrefix {
4646
}
4747

4848
fn update_prefix_checkpoint<C>(
49-
checkpoints: &[PrefixCheckpoint<F>],
49+
checkpoints: &PrefixCheckpoints<F>,
5050
r_x: C,
5151
r_y: C,
5252
_: usize,

‎joltworks/src/lookup_tables/prefixes/lower_word_no_msb.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use super::{PrefixCheckpoint, Prefixes, SparseDensePrefix};
1+
use super::{PrefixCheckpoint, PrefixCheckpoints, Prefixes, SparseDensePrefix};
22
use crate::{
33
field::{ChallengeFieldOps, FieldChallengeOps, JoltField},
44
utils::lookup_bits::LookupBits,
@@ -9,7 +9,7 @@ pub enum LowerWordNoMsbPrefix<const XLEN: usize> {}
99

1010
impl<const XLEN: usize, F: JoltField> SparseDensePrefix<F> for LowerWordNoMsbPrefix<XLEN> {
1111
fn prefix_mle<C>(
12-
checkpoints: &[PrefixCheckpoint<F>],
12+
checkpoints: &PrefixCheckpoints<F>,
1313
r_x: Option<C>,
1414
c: u32,
1515
mut b: LookupBits,
@@ -54,7 +54,7 @@ impl<const XLEN: usize, F: JoltField> SparseDensePrefix<F> for LowerWordNoMsbPre
5454
}
5555

5656
fn update_prefix_checkpoint<C>(
57-
checkpoints: &[PrefixCheckpoint<F>],
57+
checkpoints: &PrefixCheckpoints<F>,
5858
r_x: C,
5959
r_y: C,
6060
j: usize,

‎joltworks/src/lookup_tables/prefixes/mod.rs‎

Lines changed: 59 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,10 @@ use common::parallel::par_enabled;
1818
use num::FromPrimitive;
1919
use num_derive::FromPrimitive;
2020
use rayon::prelude::*;
21-
use std::{fmt::Display, ops::Index};
21+
use std::{
22+
fmt::Display,
23+
ops::{Index, IndexMut},
24+
};
2225
use strum::EnumCount;
2326
use strum_macros::{EnumCount as EnumCountMacro, EnumIter};
2427

@@ -62,7 +65,7 @@ pub trait SparseDensePrefix<F: JoltField>: 'static + Sync {
6265
/// over these variables as they range over the Boolean hypercube, so
6366
/// they can be represented by a single bitvector.
6467
fn prefix_mle<C>(
65-
checkpoints: &[PrefixCheckpoint<F>],
68+
checkpoints: &PrefixCheckpoints<F>,
6669
r_x: Option<C>,
6770
c: u32,
6871
b: LookupBits,
@@ -78,7 +81,7 @@ pub trait SparseDensePrefix<F: JoltField>: 'static + Sync {
7881
/// A checkpoint update may depend on the values of the other prefix checkpoints,
7982
/// so we pass in all such `checkpoints` to this function.
8083
fn update_prefix_checkpoint<C>(
81-
checkpoints: &[PrefixCheckpoint<F>],
84+
checkpoints: &PrefixCheckpoints<F>,
8285
r_x: C,
8386
r_y: C,
8487
j: usize,
@@ -119,6 +122,26 @@ pub struct PrefixEval<F>(F);
119122
/// Optional prefix evaluation cached after each pair of address-binding rounds (r_x, r_y).
120123
pub type PrefixCheckpoint<F: JoltField> = PrefixEval<Option<F>>;
121124

125+
#[derive(Clone)]
126+
// Stores the checkpoints for all prefixes, updated every two rounds of sumcheck.
127+
pub struct PrefixCheckpoints<F: JoltField>([PrefixCheckpoint<F>; Prefixes::COUNT]);
128+
129+
impl<F: JoltField> PrefixCheckpoints<F> {
130+
pub fn new() -> Self {
131+
Self(std::array::from_fn(|_| None.into()))
132+
}
133+
134+
pub fn len(&self) -> usize {
135+
Prefixes::COUNT
136+
}
137+
}
138+
139+
impl<F: JoltField> Default for PrefixCheckpoints<F> {
140+
fn default() -> Self {
141+
Self::new()
142+
}
143+
}
144+
122145
impl<F: JoltField> std::ops::Mul<F> for PrefixEval<F> {
123146
type Output = F;
124147

@@ -153,6 +176,34 @@ impl<F> PrefixCheckpoint<F> {
153176
}
154177
}
155178

179+
impl<F: JoltField> Index<usize> for PrefixCheckpoints<F> {
180+
type Output = Option<F>;
181+
182+
fn index(&self, index: usize) -> &Self::Output {
183+
&self.0[index].0
184+
}
185+
}
186+
187+
impl<F: JoltField> IndexMut<usize> for PrefixCheckpoints<F> {
188+
fn index_mut(&mut self, index: usize) -> &mut Self::Output {
189+
&mut self.0[index].0
190+
}
191+
}
192+
193+
impl<F: JoltField> Index<Prefixes> for PrefixCheckpoints<F> {
194+
type Output = Option<F>;
195+
196+
fn index(&self, prefix: Prefixes) -> &Self::Output {
197+
&self[prefix as usize]
198+
}
199+
}
200+
201+
impl<F: JoltField> IndexMut<Prefixes> for PrefixCheckpoints<F> {
202+
fn index_mut(&mut self, prefix: Prefixes) -> &mut Self::Output {
203+
&mut self[prefix as usize]
204+
}
205+
}
206+
156207
impl<F> Index<Prefixes> for &[PrefixEval<F>] {
157208
type Output = F;
158209

@@ -179,7 +230,7 @@ impl Prefixes {
179230
/// they can be represented by a single bitvector.
180231
pub fn prefix_mle<const XLEN: usize, F, C>(
181232
&self,
182-
checkpoints: &[PrefixCheckpoint<F>],
233+
checkpoints: &PrefixCheckpoints<F>,
183234
r_x: Option<C>,
184235
c: u32,
185236
b: LookupBits,
@@ -213,7 +264,7 @@ impl Prefixes {
213264
/// This function updates all the prefix checkpoints.
214265
#[tracing::instrument(skip_all)]
215266
pub fn update_checkpoints<const XLEN: usize, F, C>(
216-
checkpoints: &mut [PrefixCheckpoint<F>],
267+
checkpoints: &mut PrefixCheckpoints<F>,
217268
r_x: C,
218269
r_y: C,
219270
j: usize,
@@ -222,9 +273,9 @@ impl Prefixes {
222273
C: ChallengeFieldOps<F>,
223274
F: JoltField + FieldChallengeOps<C>,
224275
{
225-
debug_assert_eq!(checkpoints.len(), Self::COUNT);
226-
let previous_checkpoints = checkpoints.to_vec();
276+
let previous_checkpoints = checkpoints.clone();
227277
checkpoints
278+
.0
228279
.par_iter_mut()
229280
.with_min_len(par_enabled())
230281
.enumerate()
@@ -248,7 +299,7 @@ impl Prefixes {
248299
/// so we pass in all such `checkpoints` to this function.
249300
fn update_prefix_checkpoint<const XLEN: usize, F, C>(
250301
&self,
251-
checkpoints: &[PrefixCheckpoint<F>],
302+
checkpoints: &PrefixCheckpoints<F>,
252303
r_x: C,
253304
r_y: C,
254305
j: usize,

‎joltworks/src/lookup_tables/prefixes/msb.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use super::{PrefixCheckpoint, Prefixes, SparseDensePrefix};
1+
use super::{PrefixCheckpoint, PrefixCheckpoints, Prefixes, SparseDensePrefix};
22
use crate::{
33
field::{ChallengeFieldOps, FieldChallengeOps, JoltField},
44
utils::lookup_bits::LookupBits,
@@ -9,7 +9,7 @@ pub enum MsbPrefix<const XLEN: usize> {}
99

1010
impl<const XLEN: usize, F: JoltField> SparseDensePrefix<F> for MsbPrefix<XLEN> {
1111
fn prefix_mle<C>(
12-
checkpoints: &[PrefixCheckpoint<F>],
12+
checkpoints: &PrefixCheckpoints<F>,
1313
r_x: Option<C>,
1414
c: u32,
1515
mut _b: LookupBits,
@@ -32,7 +32,7 @@ impl<const XLEN: usize, F: JoltField> SparseDensePrefix<F> for MsbPrefix<XLEN> {
3232
}
3333

3434
fn update_prefix_checkpoint<C>(
35-
checkpoints: &[PrefixCheckpoint<F>],
35+
checkpoints: &PrefixCheckpoints<F>,
3636
r_x: C,
3737
_r_y: C,
3838
j: usize,

‎joltworks/src/lookup_tables/prefixes/nlw.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use super::{PrefixCheckpoint, Prefixes, SparseDensePrefix};
1+
use super::{PrefixCheckpoint, PrefixCheckpoints, Prefixes, SparseDensePrefix};
22
use crate::{
33
field::{ChallengeFieldOps, FieldChallengeOps, JoltField},
44
utils::lookup_bits::LookupBits,
@@ -12,7 +12,7 @@ pub enum NotLowerWordPrefix<const XLEN: usize> {}
1212

1313
impl<const XLEN: usize, F: JoltField> SparseDensePrefix<F> for NotLowerWordPrefix<XLEN> {
1414
fn prefix_mle<C>(
15-
checkpoints: &[PrefixCheckpoint<F>],
15+
checkpoints: &PrefixCheckpoints<F>,
1616
r_x: Option<C>,
1717
c: u32,
1818
mut b: LookupBits,
@@ -57,7 +57,7 @@ impl<const XLEN: usize, F: JoltField> SparseDensePrefix<F> for NotLowerWordPrefi
5757
}
5858

5959
fn update_prefix_checkpoint<C>(
60-
checkpoints: &[PrefixCheckpoint<F>],
60+
checkpoints: &PrefixCheckpoints<F>,
6161
r_x: C,
6262
r_y: C,
6363
j: usize,

‎joltworks/src/lookup_tables/prefixes/not_msb.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use super::{PrefixCheckpoint, Prefixes, SparseDensePrefix};
1+
use super::{PrefixCheckpoint, PrefixCheckpoints, Prefixes, SparseDensePrefix};
22
use crate::{
33
field::{ChallengeFieldOps, FieldChallengeOps, JoltField},
44
utils::lookup_bits::LookupBits,
@@ -9,7 +9,7 @@ pub enum NotMsbPrefix<const XLEN: usize> {}
99

1010
impl<const XLEN: usize, F: JoltField> SparseDensePrefix<F> for NotMsbPrefix<XLEN> {
1111
fn prefix_mle<C>(
12-
checkpoints: &[PrefixCheckpoint<F>],
12+
checkpoints: &PrefixCheckpoints<F>,
1313
r_x: Option<C>,
1414
c: u32,
1515
mut _b: LookupBits,
@@ -32,7 +32,7 @@ impl<const XLEN: usize, F: JoltField> SparseDensePrefix<F> for NotMsbPrefix<XLEN
3232
}
3333

3434
fn update_prefix_checkpoint<C>(
35-
checkpoints: &[PrefixCheckpoint<F>],
35+
checkpoints: &PrefixCheckpoints<F>,
3636
r_x: C,
3737
_r_y: C,
3838
j: usize,

‎joltworks/src/lookup_tables/prefixes/or.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use super::{PrefixCheckpoint, Prefixes, SparseDensePrefix};
1+
use super::{PrefixCheckpoint, PrefixCheckpoints, Prefixes, SparseDensePrefix};
22
use crate::{
33
field::{ChallengeFieldOps, FieldChallengeOps, JoltField},
44
utils::lookup_bits::LookupBits,
@@ -9,7 +9,7 @@ pub enum OrPrefix<const XLEN: usize> {}
99

1010
impl<const XLEN: usize, F: JoltField> SparseDensePrefix<F> for OrPrefix<XLEN> {
1111
fn prefix_mle<C>(
12-
checkpoints: &[PrefixCheckpoint<F>],
12+
checkpoints: &PrefixCheckpoints<F>,
1313
r_x: Option<C>,
1414
c: u32,
1515
mut b: LookupBits,
@@ -40,7 +40,7 @@ impl<const XLEN: usize, F: JoltField> SparseDensePrefix<F> for OrPrefix<XLEN> {
4040
}
4141

4242
fn update_prefix_checkpoint<C>(
43-
checkpoints: &[PrefixCheckpoint<F>],
43+
checkpoints: &PrefixCheckpoints<F>,
4444
r_x: C,
4545
r_y: C,
4646
j: usize,

‎joltworks/src/lookup_tables/prefixes/xor.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use super::{PrefixCheckpoint, Prefixes, SparseDensePrefix};
1+
use super::{PrefixCheckpoint, PrefixCheckpoints, Prefixes, SparseDensePrefix};
22
use crate::{
33
field::{ChallengeFieldOps, FieldChallengeOps, JoltField},
44
utils::lookup_bits::LookupBits,
@@ -9,7 +9,7 @@ pub enum XorPrefix<const XLEN: usize> {}
99

1010
impl<const XLEN: usize, F: JoltField> SparseDensePrefix<F> for XorPrefix<XLEN> {
1111
fn prefix_mle<C>(
12-
checkpoints: &[PrefixCheckpoint<F>],
12+
checkpoints: &PrefixCheckpoints<F>,
1313
r_x: Option<C>,
1414
c: u32,
1515
mut b: LookupBits,
@@ -41,7 +41,7 @@ impl<const XLEN: usize, F: JoltField> SparseDensePrefix<F> for XorPrefix<XLEN> {
4141
}
4242

4343
fn update_prefix_checkpoint<C>(
44-
checkpoints: &[PrefixCheckpoint<F>],
44+
checkpoints: &PrefixCheckpoints<F>,
4545
r_x: C,
4646
r_y: C,
4747
j: usize,

0 commit comments

Comments
 (0)