From 3908e496e2ea047174c1f5f5e9b59b2e718834d6 Mon Sep 17 00:00:00 2001
From: Michael Zhu <mchl.zhu.96@gmail.com>
Date: Mon, 20 Nov 2023 12:16:47 -0500
Subject: [PATCH] Add operand_chunks method to JoltInstruction; add
 chunk_operand function to instruction_utils

---
 jolt-core/src/jolt/instruction/mod.rs    |  1 +
 jolt-core/src/utils/instruction_utils.rs | 10 ++++++++++
 2 files changed, 11 insertions(+)

diff --git a/jolt-core/src/jolt/instruction/mod.rs b/jolt-core/src/jolt/instruction/mod.rs
index 918a65553..94e80dcfc 100644
--- a/jolt-core/src/jolt/instruction/mod.rs
+++ b/jolt-core/src/jolt/instruction/mod.rs
@@ -9,6 +9,7 @@ pub trait JoltInstruction {
   fn g_poly_degree(&self, C: usize) -> usize;
   fn subtables<F: PrimeField>(&self, C: usize) -> Vec<Box<dyn LassoSubtable<F>>>;
   fn to_indices(&self, C: usize, log_M: usize) -> Vec<usize>;
+  fn operand_chunks(&self, C: usize, log_M: usize) -> (Vec<u64>, Vec<u64>);
 }
 
 pub trait Opcode {
diff --git a/jolt-core/src/utils/instruction_utils.rs b/jolt-core/src/utils/instruction_utils.rs
index 4a580378d..c6a975b3d 100644
--- a/jolt-core/src/utils/instruction_utils.rs
+++ b/jolt-core/src/utils/instruction_utils.rs
@@ -14,6 +14,16 @@ pub fn concatenate_lookups<F: PrimeField>(vals: &[F], C: usize, shift_bits: usiz
   sum
 }
 
+pub fn chunk_operand(x: u64, num_chunks: usize, bits_per_chunk: usize) -> Vec<u64> {
+  let bit_mask = (1 << bits_per_chunk) - 1;
+  (0..num_chunks)
+    .map(|i| {
+      let shift = ((num_chunks - i - 1) * bits_per_chunk) as u32;
+      x.checked_shr(shift).unwrap_or(0) & bit_mask
+    })
+    .collect()
+}
+
 pub fn chunk_and_concatenate_operands(x: u64, y: u64, C: usize, log_M: usize) -> Vec<usize> {
   let operand_bits: usize = log_M / 2;
   let operand_bit_mask: usize = (1 << operand_bits) - 1;