Perform bit decomposition as assignment if offset is known. (#2612)

If we detect a bit decomposition pattern like `X + Y * 256 + Z * 65536 =
T`, the witgen inference returns a bit decomposition effect. But if the
value of `T` is a known number, we can actually extract the concrete
values for `X`, `Y` and `Z` at compile-time and continue solving with
this additional information.

Author: chriseth

executor/src/witgen/jit/affine_symbolic_expression.rs

@@ -6,7 +6,7 @@ use std::{
 
 use itertools::Itertools;
 use num_traits::Zero;
-use powdr_number::{log2_exact, FieldElement};
+use powdr_number::{log2_exact, FieldElement, LargeInt};
 
 use crate::witgen::jit::effect::Assertion;
 
@@ -243,66 +243,93 @@ impl<T: FieldElement, V: Ord + Clone + Display> AffineSymbolicExpression<T, V> {
             .coefficients
             .iter()
             .map(|(var, coeff)| {
-                let c = coeff.try_to_number()?;
+                let coeff = coeff.try_to_number()?;
                 let rc = self.range_constraints.get(var)?;
-                Some((var.clone(), c, rc))
+                let is_negative = !coeff.is_in_lower_half();
+                let coeff_abs = if is_negative { -coeff } else { coeff };
+                // We could work with non-powers of two, but it would require
+                // division instead of shifts.
+                let exponent = log2_exact(coeff_abs.to_arbitrary_integer())?;
+                // We negate here because we are solving
+                // c_1 * x_1 + c_2 * x_2 + ... + offset = 0,
+                // instead of
+                // c_1 * x_1 + c_2 * x_2 + ... = offset.
+                Some((var.clone(), rc, !is_negative, coeff_abs, exponent))
             })
             .collect::<Option<Vec<_>>>();
         let Some(constrained_coefficients) = constrained_coefficients else {
             return Ok(ProcessResult::empty());
         };
 
+        // If the offset is a known number, we gradually remove the
+        // components from this number.
+        let mut offset = self.offset.try_to_number();
+        let mut concrete_assignments = vec![];
+
         // Check if they are mutually exclusive and compute assignments.
         let mut covered_bits: <T as FieldElement>::Integer = 0.into();
         let mut components = vec![];
-        for (variable, coeff, constraint) in constrained_coefficients {
-            let is_negative = !coeff.is_in_lower_half();
-            let coeff_abs = if is_negative { -coeff } else { coeff };
-            let Some(exponent) = log2_exact(coeff_abs.to_arbitrary_integer()) else {
-                // We could work with non-powers of two, but it would require
-                // division instead of shifts.
-                return Ok(ProcessResult::empty());
-            };
+        for (variable, constraint, is_negative, coeff_abs, exponent) in constrained_coefficients
+            .into_iter()
+            .sorted_by_key(|(_, _, _, _, exponent)| *exponent)
+        {
             let bit_mask = *constraint.multiple(coeff_abs).mask();
             if !(bit_mask & covered_bits).is_zero() {
                 // Overlapping range constraints.
                 return Ok(ProcessResult::empty());
             } else {
                 covered_bits |= bit_mask;
             }
-            components.push(BitDecompositionComponent {
-                variable,
-                // We negate here because we are solving
-                // c_1 * x_1 + c_2 * x_2 + ... + offset = 0,
-                // instead of
-                // c_1 * x_1 + c_2 * x_2 + ... = offset.
-                is_negative: !is_negative,
-                exponent: exponent as u64,
-                bit_mask,
-            });
+
+            // If the offset is a known number, we create concrete assignments and modify the offset.
+            // if it is not known, we return a BitDecomposition effect.
+            if let Some(offset) = &mut offset {
+                let mut component = if is_negative { -*offset } else { *offset }.to_integer();
+                if component > (T::modulus() - 1.into()) >> 1 {
+                    // Convert a signed finite field element into two's complement.
+                    // a regular subtraction would underflow, so we do this.
+                    // We add the difference between negative numbers in the field
+                    // and negative numbers in two's complement.
+                    component += T::Integer::MAX - T::modulus() + 1.into();
+                };
+                component &= bit_mask;
+                concrete_assignments.push(Effect::Assignment(
+                    variable.clone(),
+                    T::from(component >> exponent).into(),
+                ));
+                if is_negative {
+                    *offset += T::from(component);
+                } else {
+                    *offset -= T::from(component);
+                }
+            } else {
+                components.push(BitDecompositionComponent {
+                    variable,
+                    is_negative,
+                    exponent: exponent as u64,
+                    bit_mask,
+                });
+            }
         }
 
         if covered_bits >= T::modulus() {
             return Ok(ProcessResult::empty());
         }
 
-        if !components.iter().any(|c| c.is_negative) {
-            // If all coefficients are positive and the offset is known, we can check
-            // that all bits are covered. If not, then there is no way to extract
-            // the components and thus we have a conflict.
-            if let Some(offset) = self.offset.try_to_number() {
-                if offset.to_integer() & !covered_bits != 0.into() {
-                    return Err(Error::ConflictingRangeConstraints);
-                }
+        if let Some(offset) = offset {
+            if offset != 0.into() {
+                return Err(Error::ConstraintUnsatisfiable);
             }
+            assert_eq!(concrete_assignments.len(), self.coefficients.len());
+            Ok(ProcessResult::complete(concrete_assignments))
+        } else {
+            Ok(ProcessResult::complete(vec![Effect::BitDecomposition(
+                BitDecomposition {
+                    value: self.offset.clone(),
+                    components,
+                },
+            )]))
         }
-
-        Ok(ProcessResult::complete(vec![Effect::BitDecomposition(
-            BitDecomposition {
-                value: self.offset.clone(),
-                components,
-            },
-        )]))
     }
 
     fn transfer_constraints(&self) -> Option<Effect<T, V>> {

executor/src/witgen/jit/block_machine_processor.rs

@@ -520,19 +520,19 @@ main_binary::operation_id_next[0] = main_binary::operation_id[1];
 call_var(9, 0, 0) = main_binary::operation_id_next[0];
 main_binary::operation_id_next[1] = main_binary::operation_id[2];
 call_var(9, 1, 0) = main_binary::operation_id_next[1];
-2**24 * main_binary::A_byte[2] + 2**0 * main_binary::A[2] := main_binary::A[3];
+2**0 * main_binary::A[2] + 2**24 * main_binary::A_byte[2] := main_binary::A[3];
 call_var(9, 2, 1) = main_binary::A_byte[2];
-2**16 * main_binary::A_byte[1] + 2**0 * main_binary::A[1] := main_binary::A[2];
+2**0 * main_binary::A[1] + 2**16 * main_binary::A_byte[1] := main_binary::A[2];
 call_var(9, 1, 1) = main_binary::A_byte[1];
-2**8 * main_binary::A_byte[0] + 2**0 * main_binary::A[0] := main_binary::A[1];
+2**0 * main_binary::A[0] + 2**8 * main_binary::A_byte[0] := main_binary::A[1];
 call_var(9, 0, 1) = main_binary::A_byte[0];
 main_binary::A_byte[-1] = main_binary::A[0];
 call_var(9, -1, 1) = main_binary::A_byte[-1];
-2**24 * main_binary::B_byte[2] + 2**0 * main_binary::B[2] := main_binary::B[3];
+2**0 * main_binary::B[2] + 2**24 * main_binary::B_byte[2] := main_binary::B[3];
 call_var(9, 2, 2) = main_binary::B_byte[2];
-2**16 * main_binary::B_byte[1] + 2**0 * main_binary::B[1] := main_binary::B[2];
+2**0 * main_binary::B[1] + 2**16 * main_binary::B_byte[1] := main_binary::B[2];
 call_var(9, 1, 2) = main_binary::B_byte[1];
-2**8 * main_binary::B_byte[0] + 2**0 * main_binary::B[0] := main_binary::B[1];
+2**0 * main_binary::B[0] + 2**8 * main_binary::B_byte[0] := main_binary::B[1];
 call_var(9, 0, 2) = main_binary::B_byte[0];
 main_binary::B_byte[-1] = main_binary::B[0];
 call_var(9, -1, 2) = main_binary::B_byte[-1];
@@ -607,7 +607,7 @@ params[1] = Sub::b[0];"
         assert_eq!(
             format_code(&code),
             "SubM::a[0] = params[0];
-2**8 * SubM::b[0] + 2**0 * SubM::c[0] := SubM::a[0];
+2**0 * SubM::c[0] + 2**8 * SubM::b[0] := SubM::a[0];
 params[1] = SubM::b[0];
 params[2] = SubM::c[0];
 call_var(1, 0, 0) = SubM::c[0];
@@ -816,4 +816,68 @@ S::Z[0] = params[1];
 params[2] = S::carry[0];"
         );
     }
+
+    #[test]
+    fn bit_decomp_negative_concrete() {
+        let input = "
+        namespace Main(256);
+            col witness a, b, c;
+            [a, b, c] is [S.Y, S.Z,  S.carry];
+        namespace S(256);
+            let BYTE: col = |i| i & 0xff;
+            let X;
+            let Y;
+            let Z;
+            Y = 19;
+            Z = 16;
+            let carry;
+            carry * (1 - carry) = 0;
+            [ X ] in [ BYTE ];
+            [ Y ] in [ BYTE ];
+            [ Z ] in [ BYTE ];
+            X + Y = Z + 256 * carry;
+        ";
+        let code = format_code(&generate_for_block_machine(input, "S", 2, 1).unwrap().code);
+        assert_eq!(
+            code,
+            "\
+S::Y[0] = params[0];
+S::Z[0] = params[1];
+S::X[0] = 253;
+S::carry[0] = 1;
+params[2] = 1;"
+        );
+    }
+
+    #[test]
+    fn bit_decomp_negative_concrete_2() {
+        let input = "
+        namespace Main(256);
+            col witness a, b, c;
+            [a, b, c] is [S.Y, S.Z,  S.carry];
+        namespace S(256);
+            let BYTE: col = |i| i & 0xff;
+            let X;
+            let Y;
+            let Z;
+            Y = 1;
+            Z = 16;
+            let carry;
+            carry * (1 - carry) = 0;
+            [ X ] in [ BYTE ];
+            [ Y ] in [ BYTE ];
+            [ Z ] in [ BYTE ];
+            X + Y = Z + 256 * carry;
+        ";
+        let code = format_code(&generate_for_block_machine(input, "S", 2, 1).unwrap().code);
+        assert_eq!(
+            code,
+            "\
+S::Y[0] = params[0];
+S::Z[0] = params[1];
+S::X[0] = 15;
+S::carry[0] = 0;
+params[2] = 0;"
+        );
+    }
 }

executor/src/witgen/jit/witgen_inference.rs

@@ -835,15 +835,15 @@ namespace Xor(256 * 256);
         assert_eq!(
             code,
             "\
-2**24 * Xor::A_byte[6] + 2**0 * Xor::A[6] := Xor::A[7];
-2**24 * Xor::C_byte[6] + 2**0 * Xor::C[6] := Xor::C[7];
-2**16 * Xor::A_byte[5] + 2**0 * Xor::A[5] := Xor::A[6];
-2**16 * Xor::C_byte[5] + 2**0 * Xor::C[5] := Xor::C[6];
+2**0 * Xor::A[6] + 2**24 * Xor::A_byte[6] := Xor::A[7];
+2**0 * Xor::C[6] + 2**24 * Xor::C_byte[6] := Xor::C[7];
+2**0 * Xor::A[5] + 2**16 * Xor::A_byte[5] := Xor::A[6];
+2**0 * Xor::C[5] + 2**16 * Xor::C_byte[5] := Xor::C[6];
 call_var(0, 6, 0) = Xor::A_byte[6];
 call_var(0, 6, 2) = Xor::C_byte[6];
 machine_call(1, [Known(call_var(0, 6, 0)), Unknown(call_var(0, 6, 1)), Known(call_var(0, 6, 2))]);
-2**8 * Xor::A_byte[4] + 2**0 * Xor::A[4] := Xor::A[5];
-2**8 * Xor::C_byte[4] + 2**0 * Xor::C[4] := Xor::C[5];
+2**0 * Xor::A[4] + 2**8 * Xor::A_byte[4] := Xor::A[5];
+2**0 * Xor::C[4] + 2**8 * Xor::C_byte[4] := Xor::C[5];
 call_var(0, 5, 0) = Xor::A_byte[5];
 call_var(0, 5, 2) = Xor::C_byte[5];
 machine_call(1, [Known(call_var(0, 5, 0)), Unknown(call_var(0, 5, 1)), Known(call_var(0, 5, 2))]);