Introduce RoundIdx enum instead of using i32.

This is a more secure way to distinguish the initial round from
the main rounds.

examples/sha256/table16/compression.rs

@@ -16,6 +16,7 @@ mod subregion_initial;
 mod subregion_main;
 
 use compression_gates::CompressionGate;
+use compression_util::RoundIdx;
 
 /// A variable that represents the `[A,B,C,D]` words of the SHA-256 internal state.
 ///
@@ -29,7 +30,7 @@ use compression_gates::CompressionGate;
 ///   are needed.
 #[derive(Clone, Debug)]
 pub struct AbcdVar {
-    idx: i32,
+    round_idx: RoundIdx,
     val: Option<u32>,
     a: SpreadVar,
     b: SpreadVar,
@@ -51,7 +52,7 @@ pub struct AbcdVar {
 ///   are needed.
 #[derive(Clone, Debug)]
 pub struct EfghVar {
-    idx: i32,
+    round_idx: RoundIdx,
     val: Option<u32>,
     a_lo: SpreadVar,
     a_hi: SpreadVar,
@@ -711,13 +712,8 @@ impl CompressionConfig {
             || "compress",
             |mut region| {
                 state = initialized_state.clone();
-                for idx in 0..64 {
+                for (idx, w_halves) in w_halves.iter().enumerate() {
-                    state = self.assign_round(
+                    state = self.assign_round(&mut region, idx.into(), state.clone(), &w_halves)?;
-                        &mut region,
-                        idx,
-                        state.clone(),
-                        &w_halves[idx as usize],
-                    )?;
                 }
                 Ok(())
             },

examples/sha256/table16/compression/compression_util.rs

@@ -39,101 +39,155 @@ pub const SUBREGION_MAIN_WORD: usize =
     DECOMPOSE_ABCD + SIGMA_0_ROWS + DECOMPOSE_EFGH + SIGMA_1_ROWS + CH_ROWS + MAJ_ROWS;
 pub const SUBREGION_MAIN_ROWS: usize = SUBREGION_MAIN_LEN * SUBREGION_MAIN_WORD;
 
-/// Returns starting row number of a compression round
+/// Round index.
-pub fn get_round_row(round_idx: i32) -> usize {
+#[derive(Debug, Copy, Clone)]
-    assert!(round_idx >= -1);
+pub enum RoundIdx {
-    assert!(round_idx < 64);
+    Init,
-    if round_idx == -1 {
+    Main(usize),
-        // Init subregion
+}
-        0
+
-    } else {
+impl RoundIdx {
-        // Main subregion
+    pub(crate) fn as_usize(&self) -> usize {
-        (round_idx as usize) * SUBREGION_MAIN_WORD
+        match self {
+            Self::Main(idx) => *idx,
+            _ => panic!(),
+        }
     }
 }
 
-pub fn get_decompose_e_row(round_idx: i32) -> usize {
+impl From<usize> for RoundIdx {
+    fn from(idx: usize) -> Self {
+        Self::Main(idx)
+    }
+}
+
+impl std::ops::Add<usize> for RoundIdx {
+    type Output = Self;
+
+    fn add(self, rhs: usize) -> Self::Output {
+        match self {
+            Self::Main(idx) => Self::Main(idx + rhs),
+            _ => panic!(),
+        }
+    }
+}
+
+impl Ord for RoundIdx {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        match (self, other) {
+            (Self::Main(idx_0), Self::Main(idx_1)) => idx_0.cmp(idx_1),
+            _ => panic!(),
+        }
+    }
+}
+
+impl PartialOrd for RoundIdx {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl PartialEq for RoundIdx {
+    fn eq(&self, other: &Self) -> bool {
+        match (self, other) {
+            (Self::Main(idx_0), Self::Main(idx_1)) => idx_0 == idx_1,
+            _ => panic!(),
+        }
+    }
+}
+
+impl Eq for RoundIdx {}
+
+/// Returns starting row number of a compression round
+pub fn get_round_row(round_idx: RoundIdx) -> usize {
+    match round_idx {
+        RoundIdx::Init => 0,
+        RoundIdx::Main(idx) => {
+            assert!(idx < 64);
+            (idx as usize) * SUBREGION_MAIN_WORD
+        }
+    }
+}
+
+pub fn get_decompose_e_row(round_idx: RoundIdx) -> usize {
     get_round_row(round_idx)
 }
 
-pub fn get_decompose_f_row(round_idx: i32) -> usize {
+pub fn get_decompose_f_row(round_idx: RoundIdx) -> usize {
-    assert_eq!(round_idx, -1);
+    assert!(matches!(round_idx, RoundIdx::Init));
     get_decompose_e_row(round_idx) + DECOMPOSE_EFGH
 }
 
-pub fn get_decompose_g_row(round_idx: i32) -> usize {
+pub fn get_decompose_g_row(round_idx: RoundIdx) -> usize {
     get_decompose_f_row(round_idx) + DECOMPOSE_EFGH
 }
 
-pub fn get_upper_sigma_1_row(round_idx: i32) -> usize {
+pub fn get_upper_sigma_1_row(round_idx: RoundIdx) -> usize {
-    assert!(round_idx >= 0);
+    assert!(matches!(round_idx, RoundIdx::Main(_)));
     get_decompose_e_row(round_idx) + DECOMPOSE_EFGH + 1
 }
 
-pub fn get_ch_row(round_idx: i32) -> usize {
+pub fn get_ch_row(round_idx: RoundIdx) -> usize {
-    assert!(round_idx >= 0);
+    assert!(matches!(round_idx, RoundIdx::Main(_)));
     get_decompose_e_row(round_idx) + DECOMPOSE_EFGH + SIGMA_1_ROWS + 1
 }
 
-pub fn get_ch_neg_row(round_idx: i32) -> usize {
+pub fn get_ch_neg_row(round_idx: RoundIdx) -> usize {
     get_ch_row(round_idx) + CH_ROWS / 2
 }
 
-pub fn get_decompose_a_row(round_idx: i32) -> usize {
+pub fn get_decompose_a_row(round_idx: RoundIdx) -> usize {
-    if round_idx == -1 {
+    match round_idx {
-        get_h_row(round_idx) + DECOMPOSE_EFGH
+        RoundIdx::Init => get_h_row(round_idx) + DECOMPOSE_EFGH,
-    } else {
+        _ => get_ch_neg_row(round_idx) - 1 + CH_ROWS / 2,
-        get_ch_neg_row(round_idx) - 1 + CH_ROWS / 2
     }
 }
 
-pub fn get_upper_sigma_0_row(round_idx: i32) -> usize {
+pub fn get_upper_sigma_0_row(round_idx: RoundIdx) -> usize {
-    assert!(round_idx >= 0);
+    assert!(matches!(round_idx, RoundIdx::Main(_)));
     get_decompose_a_row(round_idx) + DECOMPOSE_ABCD + 1
 }
 
-pub fn get_decompose_b_row(round_idx: i32) -> usize {
+pub fn get_decompose_b_row(round_idx: RoundIdx) -> usize {
-    assert_eq!(round_idx, -1);
+    assert!(matches!(round_idx, RoundIdx::Init));
     get_decompose_a_row(round_idx) + DECOMPOSE_ABCD
 }
 
-pub fn get_decompose_c_row(round_idx: i32) -> usize {
+pub fn get_decompose_c_row(round_idx: RoundIdx) -> usize {
     get_decompose_b_row(round_idx) + DECOMPOSE_ABCD
 }
 
-pub fn get_maj_row(round_idx: i32) -> usize {
+pub fn get_maj_row(round_idx: RoundIdx) -> usize {
-    assert!(round_idx >= 0);
+    assert!(matches!(round_idx, RoundIdx::Main(_)));
     get_upper_sigma_0_row(round_idx) + SIGMA_0_ROWS
 }
 
 // Get state word rows
-pub fn get_h_row(round_idx: i32) -> usize {
+pub fn get_h_row(round_idx: RoundIdx) -> usize {
-    if round_idx == -1 {
+    match round_idx {
-        get_decompose_g_row(round_idx) + DECOMPOSE_EFGH
+        RoundIdx::Init => get_decompose_g_row(round_idx) + DECOMPOSE_EFGH,
-    } else {
+        _ => get_ch_row(round_idx) - 1,
-        get_ch_row(round_idx) - 1
     }
 }
 
-pub fn get_h_prime_row(round_idx: i32) -> usize {
+pub fn get_h_prime_row(round_idx: RoundIdx) -> usize {
-    assert!(round_idx >= 0);
+    assert!(matches!(round_idx, RoundIdx::Main(_)));
     get_ch_row(round_idx)
 }
 
-pub fn get_d_row(round_idx: i32) -> usize {
+pub fn get_d_row(round_idx: RoundIdx) -> usize {
-    if round_idx == -1 {
+    match round_idx {
-        get_decompose_c_row(round_idx) + DECOMPOSE_ABCD
+        RoundIdx::Init => get_decompose_c_row(round_idx) + DECOMPOSE_ABCD,
-    } else {
+        _ => get_ch_row(round_idx) + 2,
-        get_ch_row(round_idx) + 2
     }
 }
 
-pub fn get_e_new_row(round_idx: i32) -> usize {
+pub fn get_e_new_row(round_idx: RoundIdx) -> usize {
-    assert!(round_idx >= 0);
+    assert!(matches!(round_idx, RoundIdx::Main(_)));
     get_d_row(round_idx)
 }
 
-pub fn get_a_new_row(round_idx: i32) -> usize {
+pub fn get_a_new_row(round_idx: RoundIdx) -> usize {
     get_maj_row(round_idx)
 }
 
@@ -232,15 +286,15 @@ impl CompressionConfig {
     pub(super) fn decompose_a(
         &self,
         region: &mut Region<'_, pallas::Base>,
-        idx: i32,
+        round_idx: RoundIdx,
         a_val: Option<u32>,
     ) -> Result<RoundWordA, Error> {
-        let row = get_decompose_a_row(idx);
+        let row = get_decompose_a_row(round_idx);
 
         let (dense_halves, spread_halves) = self.assign_word_halves(region, row, a_val)?;
         let (a, b, c_lo, c_mid, c_hi, d) = self.decompose_abcd(region, row, a_val)?;
         let a_pieces = AbcdVar {
-            idx,
+            round_idx,
             val: a_val,
             a,
             b,
@@ -255,15 +309,15 @@ impl CompressionConfig {
     pub(super) fn decompose_e(
         &self,
         region: &mut Region<'_, pallas::Base>,
-        idx: i32,
+        round_idx: RoundIdx,
         e_val: Option<u32>,
     ) -> Result<RoundWordE, Error> {
-        let row = get_decompose_e_row(idx);
+        let row = get_decompose_e_row(round_idx);
 
         let (dense_halves, spread_halves) = self.assign_word_halves(region, row, e_val)?;
         let (a_lo, a_hi, b_lo, b_hi, c, d) = self.decompose_efgh(region, row, e_val)?;
         let e_pieces = EfghVar {
-            idx,
+            round_idx,
             val: e_val,
             a_lo,
             a_hi,
@@ -278,7 +332,7 @@ impl CompressionConfig {
     pub(super) fn assign_upper_sigma_0(
         &self,
         region: &mut Region<'_, pallas::Base>,
-        idx: i32,
+        round_idx: RoundIdx,
         word: AbcdVar,
     ) -> Result<(CellValue16, CellValue16), Error> {
         // Rename these here for ease of matching the gates to the specification.
@@ -286,7 +340,7 @@ impl CompressionConfig {
         let a_4 = self.extras[1];
         let a_5 = self.message_schedule;
 
-        let row = get_upper_sigma_0_row(idx);
+        let row = get_upper_sigma_0_row(round_idx);
 
         self.s_upper_sigma_0.enable(region, row)?;
 
@@ -363,7 +417,7 @@ impl CompressionConfig {
     pub(super) fn assign_upper_sigma_1(
         &self,
         region: &mut Region<'_, pallas::Base>,
-        idx: i32,
+        round_idx: RoundIdx,
         word: EfghVar,
     ) -> Result<(CellValue16, CellValue16), Error> {
         // Rename these here for ease of matching the gates to the specification.
@@ -371,7 +425,7 @@ impl CompressionConfig {
         let a_4 = self.extras[1];
         let a_5 = self.message_schedule;
 
-        let row = get_upper_sigma_1_row(idx);
+        let row = get_upper_sigma_1_row(round_idx);
 
         self.s_upper_sigma_1.enable(region, row)?;
 
@@ -473,14 +527,14 @@ impl CompressionConfig {
     pub(super) fn assign_ch(
         &self,
         region: &mut Region<'_, pallas::Base>,
-        idx: i32,
+        round_idx: RoundIdx,
         spread_halves_e: (CellValue32, CellValue32),
         spread_halves_f: (CellValue32, CellValue32),
     ) -> Result<(CellValue16, CellValue16), Error> {
         let a_3 = self.extras[0];
         let a_4 = self.extras[1];
 
-        let row = get_ch_row(idx);
+        let row = get_ch_row(round_idx);
 
         self.s_ch.enable(region, row)?;
 
@@ -521,11 +575,11 @@ impl CompressionConfig {
     pub(super) fn assign_ch_neg(
         &self,
         region: &mut Region<'_, pallas::Base>,
-        idx: i32,
+        round_idx: RoundIdx,
         spread_halves_e: (CellValue32, CellValue32),
         spread_halves_g: (CellValue32, CellValue32),
     ) -> Result<(CellValue16, CellValue16), Error> {
-        let row = get_ch_neg_row(idx);
+        let row = get_ch_neg_row(round_idx);
 
         self.s_ch_neg.enable(region, row)?;
 
@@ -626,7 +680,7 @@ impl CompressionConfig {
     pub(super) fn assign_maj(
         &self,
         region: &mut Region<'_, pallas::Base>,
-        idx: i32,
+        round_idx: RoundIdx,
         spread_halves_a: (CellValue32, CellValue32),
         spread_halves_b: (CellValue32, CellValue32),
         spread_halves_c: (CellValue32, CellValue32),
@@ -634,7 +688,7 @@ impl CompressionConfig {
         let a_4 = self.extras[1];
         let a_5 = self.message_schedule;
 
-        let row = get_maj_row(idx);
+        let row = get_maj_row(round_idx);
 
         self.s_maj.enable(region, row)?;
 
@@ -687,7 +741,7 @@ impl CompressionConfig {
     pub(super) fn assign_h_prime(
         &self,
         region: &mut Region<'_, pallas::Base>,
-        idx: i32,
+        round_idx: RoundIdx,
         h: (CellValue16, CellValue16),
         ch: (CellValue16, CellValue16),
         ch_neg: (CellValue16, CellValue16),
@@ -695,7 +749,7 @@ impl CompressionConfig {
         k: u32,
         w: &(CellValue16, CellValue16),
     ) -> Result<(CellValue16, CellValue16), Error> {
-        let row = get_h_prime_row(idx);
+        let row = get_h_prime_row(round_idx);
         self.s_h_prime.enable(region, row)?;
 
         let a_4 = self.extras[1];
@@ -777,11 +831,11 @@ impl CompressionConfig {
     pub(super) fn assign_e_new(
         &self,
         region: &mut Region<'_, pallas::Base>,
-        idx: i32,
+        round_idx: RoundIdx,
         d: &(CellValue16, CellValue16),
         h_prime: &(CellValue16, CellValue16),
     ) -> Result<(CellValue16, CellValue16), Error> {
-        let row = get_e_new_row(idx);
+        let row = get_e_new_row(round_idx);
 
         self.s_e_new.enable(region, row)?;
 
@@ -818,12 +872,12 @@ impl CompressionConfig {
     pub(super) fn assign_a_new(
         &self,
         region: &mut Region<'_, pallas::Base>,
-        idx: i32,
+        round_idx: RoundIdx,
         maj: (CellValue16, CellValue16),
         sigma_0: (CellValue16, CellValue16),
         h_prime: (CellValue16, CellValue16),
     ) -> Result<(CellValue16, CellValue16), Error> {
-        let row = get_a_new_row(idx);
+        let row = get_a_new_row(round_idx);
 
         self.s_a_new.enable(region, row)?;
 

examples/sha256/table16/compression/subregion_initial.rs

@@ -11,30 +11,28 @@ impl CompressionConfig {
     ) -> Result<State, Error> {
         let a_7 = self.extras[3];
 
-        let idx = -1;
-
         // Decompose E into (6, 5, 14, 7)-bit chunks
-        let e = self.decompose_e(region, idx, Some(iv[4]))?;
+        let e = self.decompose_e(region, RoundIdx::Init, Some(iv[4]))?;
 
         // Decompose F, G
-        let f = self.decompose_f(region, idx, Some(iv[5]))?;
+        let f = self.decompose_f(region, RoundIdx::Init, Some(iv[5]))?;
-        let g = self.decompose_g(region, idx, Some(iv[6]))?;
+        let g = self.decompose_g(region, RoundIdx::Init, Some(iv[6]))?;
 
         // Assign H
-        let h_row = get_h_row(idx);
+        let h_row = get_h_row(RoundIdx::Init);
         let h_dense =
             self.assign_word_halves_dense(region, h_row, a_7, h_row + 1, a_7, Some(iv[7]))?;
         let h = RoundWordDense::new(h_dense);
 
         // Decompose A into (2, 11, 9, 10)-bit chunks
-        let a = self.decompose_a(region, idx, Some(iv[0]))?;
+        let a = self.decompose_a(region, RoundIdx::Init, Some(iv[0]))?;
 
         // Decompose B, C
-        let b = self.decompose_b(region, idx, Some(iv[1]))?;
+        let b = self.decompose_b(region, RoundIdx::Init, Some(iv[1]))?;
-        let c = self.decompose_c(region, idx, Some(iv[2]))?;
+        let c = self.decompose_c(region, RoundIdx::Init, Some(iv[2]))?;
 
         // Assign D
-        let d_row = get_d_row(idx);
+        let d_row = get_d_row(RoundIdx::Init);
         let d_dense =
             self.assign_word_halves_dense(region, d_row, a_7, d_row + 1, a_7, Some(iv[3]))?;
         let d = RoundWordDense::new(d_dense);
@@ -60,37 +58,35 @@ impl CompressionConfig {
         let a_7 = self.extras[3];
         let (a, b, c, d, e, f, g, h) = match_state(state);
 
-        let idx = -1;
-
         // Decompose E into (6, 5, 14, 7)-bit chunks
         let e = val_from_dense_halves(&e.dense_halves);
-        let e = self.decompose_e(region, idx, e)?;
+        let e = self.decompose_e(region, RoundIdx::Init, e)?;
 
         // Decompose F, G
         let f = val_from_dense_halves(&f.dense_halves);
-        let f = self.decompose_f(region, idx, f)?;
+        let f = self.decompose_f(region, RoundIdx::Init, f)?;
         let g = val_from_dense_halves(&g.dense_halves);
-        let g = self.decompose_g(region, idx, g)?;
+        let g = self.decompose_g(region, RoundIdx::Init, g)?;
 
         // Assign H
         let h = val_from_dense_halves(&h.dense_halves);
-        let h_row = get_h_row(idx);
+        let h_row = get_h_row(RoundIdx::Init);
         let h_dense = self.assign_word_halves_dense(region, h_row, a_7, h_row + 1, a_7, h)?;
         let h = RoundWordDense::new(h_dense);
 
         // Decompose A into (2, 11, 9, 10)-bit chunks
         let a = val_from_dense_halves(&a.dense_halves);
-        let a = self.decompose_a(region, idx, a)?;
+        let a = self.decompose_a(region, RoundIdx::Init, a)?;
 
         // Decompose B, C
         let b = val_from_dense_halves(&b.dense_halves);
-        let b = self.decompose_b(region, idx, b)?;
+        let b = self.decompose_b(region, RoundIdx::Init, b)?;
         let c = val_from_dense_halves(&c.dense_halves);
-        let c = self.decompose_c(region, idx, c)?;
+        let c = self.decompose_c(region, RoundIdx::Init, c)?;
 
         // Assign D
         let d = val_from_dense_halves(&d.dense_halves);
-        let d_row = get_d_row(idx);
+        let d_row = get_d_row(RoundIdx::Init);
         let d_dense = self.assign_word_halves_dense(region, d_row, a_7, d_row + 1, a_7, d)?;
         let d = RoundWordDense::new(d_dense);
 
@@ -109,10 +105,10 @@ impl CompressionConfig {
     fn decompose_b(
         &self,
         region: &mut Region<'_, pallas::Base>,
-        idx: i32,
+        round_idx: RoundIdx,
         b_val: Option<u32>,
     ) -> Result<RoundWordSpread, Error> {
-        let row = get_decompose_b_row(idx);
+        let row = get_decompose_b_row(round_idx);
 
         let (dense_halves, spread_halves) = self.assign_word_halves(region, row, b_val)?;
         self.decompose_abcd(region, row, b_val)?;
@@ -122,10 +118,10 @@ impl CompressionConfig {
     fn decompose_c(
         &self,
         region: &mut Region<'_, pallas::Base>,
-        idx: i32,
+        round_idx: RoundIdx,
         c_val: Option<u32>,
     ) -> Result<RoundWordSpread, Error> {
-        let row = get_decompose_c_row(idx);
+        let row = get_decompose_c_row(round_idx);
 
         let (dense_halves, spread_halves) = self.assign_word_halves(region, row, c_val)?;
         self.decompose_abcd(region, row, c_val)?;
@@ -135,10 +131,10 @@ impl CompressionConfig {
     fn decompose_f(
         &self,
         region: &mut Region<'_, pallas::Base>,
-        idx: i32,
+        round_idx: RoundIdx,
         f_val: Option<u32>,
     ) -> Result<RoundWordSpread, Error> {
-        let row = get_decompose_f_row(idx);
+        let row = get_decompose_f_row(round_idx);
 
         let (dense_halves, spread_halves) = self.assign_word_halves(region, row, f_val)?;
         self.decompose_efgh(region, row, f_val)?;
@@ -148,10 +144,10 @@ impl CompressionConfig {
     fn decompose_g(
         &self,
         region: &mut Region<'_, pallas::Base>,
-        idx: i32,
+        round_idx: RoundIdx,
         g_val: Option<u32>,
     ) -> Result<RoundWordSpread, Error> {
-        let row = get_decompose_g_row(idx);
+        let row = get_decompose_g_row(round_idx);
 
         let (dense_halves, spread_halves) = self.assign_word_halves(region, row, g_val)?;
         self.decompose_efgh(region, row, g_val)?;

examples/sha256/table16/compression/subregion_main.rs

@@ -7,10 +7,12 @@ impl CompressionConfig {
     pub fn assign_round(
         &self,
         region: &mut Region<'_, pallas::Base>,
-        idx: i32,
+        round_idx: RoundIdx,
         state: State,
         schedule_word: &(CellValue16, CellValue16),
     ) -> Result<State, Error> {
+        assert!(matches!(round_idx, RoundIdx::Main(_)));
+
         let a_3 = self.extras[0];
         let a_4 = self.extras[1];
         let a_7 = self.extras[3];
@@ -18,29 +20,29 @@ impl CompressionConfig {
         let (a, b, c, d, e, f, g, h) = match_state(state);
 
         // s_upper_sigma_1(E)
-        let sigma_1 = self.assign_upper_sigma_1(region, idx, e.pieces.clone().unwrap())?;
+        let sigma_1 = self.assign_upper_sigma_1(region, round_idx, e.pieces.clone().unwrap())?;
 
         // Ch(E, F, G)
         let ch = self.assign_ch(
             region,
-            idx,
+            round_idx,
             e.spread_halves.clone().unwrap(),
             f.spread_halves.clone(),
         )?;
         let ch_neg = self.assign_ch_neg(
             region,
-            idx,
+            round_idx,
             e.spread_halves.clone().unwrap(),
             g.spread_halves.clone(),
         )?;
 
         // s_upper_sigma_0(A)
-        let sigma_0 = self.assign_upper_sigma_0(region, idx, a.pieces.clone().unwrap())?;
+        let sigma_0 = self.assign_upper_sigma_0(region, round_idx, a.pieces.clone().unwrap())?;
 
         // Maj(A, B, C)
         let maj = self.assign_maj(
             region,
-            idx,
+            round_idx,
             a.spread_halves.clone().unwrap(),
             b.spread_halves.clone(),
             c.spread_halves.clone(),
@@ -49,26 +51,26 @@ impl CompressionConfig {
         // H' = H + Ch(E, F, G) + s_upper_sigma_1(E) + K + W
         let h_prime = self.assign_h_prime(
             region,
-            idx,
+            round_idx,
             h.dense_halves,
             ch,
             ch_neg,
             sigma_1,
-            ROUND_CONSTANTS[idx as usize],
+            ROUND_CONSTANTS[round_idx.as_usize()],
             schedule_word,
         )?;
 
         // E_new = H' + D
-        let e_new_dense = self.assign_e_new(region, idx, &d.dense_halves, &h_prime)?;
+        let e_new_dense = self.assign_e_new(region, round_idx, &d.dense_halves, &h_prime)?;
         let e_new_val = val_from_dense_halves(&e_new_dense);
 
         // A_new = H' + Maj(A, B, C) + sigma_0(A)
-        let a_new_dense = self.assign_a_new(region, idx, maj, sigma_0, h_prime)?;
+        let a_new_dense = self.assign_a_new(region, round_idx, maj, sigma_0, h_prime)?;
         let a_new_val = val_from_dense_halves(&a_new_dense);
 
-        if idx < 63 {
+        if round_idx < 63.into() {
             // Assign and copy A_new
-            let a_new_row = get_decompose_a_row(idx + 1);
+            let a_new_row = get_decompose_a_row(round_idx + 1);
             a_new_dense
                 .0
                 .copy_advice(|| "a_new_lo", region, a_7, a_new_row)?;
@@ -77,7 +79,7 @@ impl CompressionConfig {
                 .copy_advice(|| "a_new_hi", region, a_7, a_new_row + 1)?;
 
             // Assign and copy E_new
-            let e_new_row = get_decompose_e_row(idx + 1);
+            let e_new_row = get_decompose_e_row(round_idx + 1);
             e_new_dense
                 .0
                 .copy_advice(|| "e_new_lo", region, a_7, e_new_row)?;
@@ -86,10 +88,10 @@ impl CompressionConfig {
                 .copy_advice(|| "e_new_hi", region, a_7, e_new_row + 1)?;
 
             // Decompose A into (2, 11, 9, 10)-bit chunks
-            let a_new = self.decompose_a(region, idx + 1, a_new_val)?;
+            let a_new = self.decompose_a(region, round_idx + 1, a_new_val)?;
 
             // Decompose E into (6, 5, 14, 7)-bit chunks
-            let e_new = self.decompose_e(region, idx + 1, e_new_val)?;
+            let e_new = self.decompose_e(region, round_idx + 1, e_new_val)?;
 
             Ok(State::new(
                 StateWord::A(a_new),