benches: fix divmod_128_64 distribution

jets-bench/benches/elements/main.rs

@@ -11,7 +11,7 @@ use simplicity::types::Final;
 use simplicity::Value;
 use simplicity_bench::input::{
     self, EqProduct, GenericProduct, InputSample, PrefixBit, Sha256Ctx, UniformBits,
-    UniformBitsExact,
+    DivMod12864Input,
 };
 use simplicity_bench::{
     genesis_pegin, BenchSample, EnvSampling, InputSampling, JetBuffer, JetParams, SimplicityCtx8,
@@ -508,7 +508,7 @@ fn bench(c: &mut Criterion) {
         (Elements::DivMod16, &EqProduct(UniformBits)),
         (Elements::DivMod32, &EqProduct(UniformBits)),
         (Elements::DivMod64, &EqProduct(UniformBits)),
-        (Elements::DivMod128_64, &GenericProduct(UniformBitsExact::<128>, UniformBitsExact::<64>)),
+        (Elements::DivMod128_64, &DivMod12864Input),
         // divide
         (Elements::Divide8, &EqProduct(UniformBits)),
         (Elements::Divide16, &EqProduct(UniformBits)),

jets-bench/src/input.rs

@@ -282,7 +282,7 @@ impl FlatValue {
     }
 
     /// Creates an iterator over the bits of the value
-    pub fn bit_iter(&self) -> FlatValueBitIter {
+    pub fn bit_iter(&self) -> FlatValueBitIter<'_> {
         FlatValueBitIter {
             inner: self,
             n_read: 0,
@@ -1322,6 +1322,66 @@ impl InputSampleExactSize for CheckSigSignature {
     }
 }
 
+pub struct DivMod12864Input;
+
+impl InputSample for DivMod12864Input {
+    fn n_distributions(&self) -> usize {
+        // divmod128_64 has specific preconditions (unlike the other divmod jets).
+        // If these are not met then it does a quick shortcut error return.
+        2
+    }
+
+    fn distribution_name(&self, dist: usize) -> String {
+        if dist == 0 {
+            "uniform".into()
+        } else {
+            "divmod12864precondition".into()
+        }
+    }
+
+    fn sample(&self, dist: usize, n_bits: usize) -> FlatValue {
+        assert_eq!(
+            n_bits, 192,
+            "attempt to sample divmod128_64 input with {} bits (expected 192)",
+            n_bits
+        );
+
+        if dist == 2 {
+            UniformBits.sample(0, 192)
+        } else if dist == 0 || dist == 1 {
+            // The input to div_mod_128_64 is a 128-bit integer `a` followed by a 64-bit integer `b`.
+            // It has two preconditions:
+            //     1. The high bit of b must be set.
+            //     2. ah (the high 64 bits of a) is strictly less than b.
+            //
+            // See the implementation of `simplicity_div_mod_128_64` in libsimplicity C/jets.c
+            // for how these are implemented.
+
+            // To implement this, we sample two random 64-bit strings, with one of them
+            // having its high bit forced to 1. The higher of the two (which will always
+            // have its high bit 1) is b, and the other one is ah.
+            let sample_1 = FlatValue::random_n_bits(63).prefix_one();
+            let sample_2 = FlatValue::random_n_bits(64);
+            for (bit1, bit2) in sample_1.bit_iter().zip(sample_2.bit_iter()) {
+                match (bit1, bit2) {
+                    (false, false) | (true, true) => {} // both equal
+                    (true, false) => return FlatValue::product(&[sample_2, UniformBits.sample(0, 64), sample_1]),
+                    (false, true) => return FlatValue::product(&[sample_1, UniformBits.sample(0, 64), sample_2]),
+                }
+            }
+            unreachable!("if we get here, two uniform 63-bit samples were exactly equal")
+        } else {
+            panic!("invalid distribution {} for div_mod_128_64", dist)
+        }
+    }
+}
+
+impl InputSampleExactSize for DivMod12864Input {
+    fn n_bits(&self) -> usize {
+        192
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;