Feat/#482 unify signed bit extraction

[KimiWu123]

To close #482

[matthiasgoergens]

As far as I can tell, every single time you use SignedExtendConfig::construct_limb you do .limbs.iter().last().unwrap().expr()). Perhaps just make a helper function that does this for you?

(That's more or less exactly what the motivation behind my is_negative was.)

[matthiasgoergens]

You can let the Rust compiler handle the type inference here:

diff --git a/ceno_zkvm/src/gadgets/is_lt.rs b/ceno_zkvm/src/gadgets/is_lt.rs
index 1dd11659..12da88cd 100644
--- a/ceno_zkvm/src/gadgets/is_lt.rs
+++ b/ceno_zkvm/src/gadgets/is_lt.rs
@@ -368,12 +368,12 @@ impl InnerSignedLtConfig {
     ) -> Result<(), ZKVMError> {
         let lhs_value = Value::new_unchecked(lhs as Word);
         let rhs_value = Value::new_unchecked(rhs as Word);
-        self.is_lhs_neg.assign_instance::<E>(
+        self.is_lhs_neg.assign_instance(
             instance,
             lkm,
             *lhs_value.as_u16_limbs().last().unwrap() as u64,
         )?;
-        self.is_rhs_neg.assign_instance::<E>(
+        self.is_rhs_neg.assign_instance(
             instance,
             lkm,
             *rhs_value.as_u16_limbs().last().unwrap() as u64,
diff --git a/ceno_zkvm/src/gadgets/signed_ext.rs b/ceno_zkvm/src/gadgets/signed_ext.rs
index 162a2745..1252532c 100644
--- a/ceno_zkvm/src/gadgets/signed_ext.rs
+++ b/ceno_zkvm/src/gadgets/signed_ext.rs
@@ -6,6 +6,7 @@ use crate::{
     set_val,
     witness::LkMultiplicity,
 };
+use ff::PrimeField;
 use ff_ext::ExtensionField;
 use std::mem::MaybeUninit;
 
@@ -76,9 +77,9 @@ impl SignedExtendConfig {
         UInt::from_exprs_unchecked(vec![limb0, self.msb.expr() * 0xffff])
     }
 
-    pub fn assign_instance<E: ExtensionField>(
+    pub fn assign_instance<BaseField: PrimeField>(
         &self,
-        instance: &mut [MaybeUninit<E::BaseField>],
+        instance: &mut [MaybeUninit<BaseField>],
         lk_multiplicity: &mut LkMultiplicity,
         val: u64,
     ) -> Result<(), ZKVMError> {
@@ -91,7 +92,7 @@ impl SignedExtendConfig {
         };
 
         assert_ux(lk_multiplicity, 2 * val - (msb << self.n_bits));
-        set_val!(instance, self.msb, E::BaseField::from(msb));
+        set_val!(instance, self.msb, BaseField::from(msb));
 
         Ok(())
     }
diff --git a/ceno_zkvm/src/instructions/riscv/memory/load.rs b/ceno_zkvm/src/instructions/riscv/memory/load.rs
index 1bc76f31..1a10234f 100644
--- a/ceno_zkvm/src/instructions/riscv/memory/load.rs
+++ b/ceno_zkvm/src/instructions/riscv/memory/load.rs
@@ -249,7 +249,7 @@ impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for LoadInstruction<E,
             _ => 0,
         };
         if let Some(signed_ext_config) = config.signed_extend_config.as_ref() {
-            signed_ext_config.assign_instance::<E>(instance, lk_multiplicity, val)?;
+            signed_ext_config.assign_instance(instance, lk_multiplicity, val)?;
         }
 
         Ok(())
diff --git a/ceno_zkvm/src/instructions/riscv/mulh.rs b/ceno_zkvm/src/instructions/riscv/mulh.rs
index 19f92a32..91c987e0 100644
--- a/ceno_zkvm/src/instructions/riscv/mulh.rs
+++ b/ceno_zkvm/src/instructions/riscv/mulh.rs
@@ -284,7 +284,7 @@ impl<E: ExtensionField> Signed<E> {
         lkm: &mut LkMultiplicity,
         val: &Value<u32>,
     ) -> Result<i32, ZKVMError> {
-        self.is_negative.assign_instance::<E>(
+        self.is_negative.assign_instance(
             instance,
             lkm,
             *val.as_u16_limbs().last().unwrap() as u64,
diff --git a/ceno_zkvm/src/instructions/riscv/shift_imm.rs b/ceno_zkvm/src/instructions/riscv/shift_imm.rs
index 8b8f46b2..3541fcf0 100644
--- a/ceno_zkvm/src/instructions/riscv/shift_imm.rs
+++ b/ceno_zkvm/src/instructions/riscv/shift_imm.rs
@@ -139,7 +139,7 @@ impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for ShiftImmInstructio
             InsnKind::SLLI => (rs1_read.as_u64() * imm as u64) >> UInt::<E>::TOTAL_BITS,
             InsnKind::SRAI | InsnKind::SRLI => {
                 if I::INST_KIND == InsnKind::SRAI {
-                    config.is_lt_config.as_ref().unwrap().assign_instance::<E>(
+                    config.is_lt_config.as_ref().unwrap().assign_instance(
                         instance,
                         lk_multiplicity,
                         *rs1_read.as_u16_limbs().last().unwrap() as u64,
diff --git a/ceno_zkvm/src/instructions/riscv/slti.rs b/ceno_zkvm/src/instructions/riscv/slti.rs
index 85efba19..557524c2 100644
--- a/ceno_zkvm/src/instructions/riscv/slti.rs
+++ b/ceno_zkvm/src/instructions/riscv/slti.rs
@@ -116,7 +116,7 @@ impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for SetLessThanImmInst
                     .assign_instance(instance, lkm, rs1 as u64, imm as u64)?;
             }
             InsnKind::SLTI => {
-                config.is_rs1_neg.as_ref().unwrap().assign_instance::<E>(
+                config.is_rs1_neg.as_ref().unwrap().assign_instance(
                     instance,
                     lkm,
                     *rs1_value.as_u16_limbs().last().unwrap() as u64,

[KimiWu123]

I added a PhantomData<E> in SignedExtendConfig in 1a8f40a. The reason why I added that is because I want to make our generic type more consistent (aligned to ExtensionField instead of BaseField). What do you think?

ceno/ceno_zkvm/src/gadgets/signed_ext.rs

Line 19 in 1a8f40a

_marker: PhantomData<E>,

[matthiasgoergens]

I think that can also work, yes. Let me double check.

[matthiasgoergens]

@KimiWu123 I extracted the phantom data change to #518 because this PR was getting a bit too long for me to wrap my head around.

So we can get #518 in first, and that should make this PR easier to read.

[matthiasgoergens]

The wording here is a bit confusing. Let me think of a better one.

[KimiWu123]

We don't need this comment now so I removed it.

[matthiasgoergens]

It's very satisfying that we can remove all this code. Nice!

[matthiasgoergens]

Looks like a nice improvement.

Please have a look at my comments (especially the one about type inference) before merging, but I'm otherwise really happy with this in principle.