Converting "c >= r" or "r <= c" into "r < c + 1" if possible.

Author: lvella

integration/src/const_collapse.rs

@@ -47,17 +47,19 @@ pub fn collapse_const_if_possible(op: &Operator, inputs: &[MaybeConstant]) {
             {
                 // Right operand is constant and can be immediate
                 must_collapse.replace(true);
-            } else if is_commutative(op)
-                && let [
+            } else if let [
                     MaybeConstant::ReferenceConstant {
                         value,
                         must_collapse,
                     },
                     MaybeConstant::NonConstant,
                 ] = inputs
-                && can_be_i16(value)
+                && ((is_commutative(op)
+                    && can_be_i16(value)) || can_turn_to_lt(op, value).is_some())
             {
-                // Left operand is constant and can be immediate (commutative operations only)
+                // Left operand is constant and can be immediate
+                // (either the operation is commutative, or it's
+                // a GE that can be transformed to LT)
                 must_collapse.replace(true);
             }
         }
@@ -98,6 +100,18 @@ pub fn collapse_const_if_possible(op: &Operator, inputs: &[MaybeConstant]) {
             {
                 // Left operand is constant and can be immediate
                 must_collapse.replace(true);
+            } else if let [
+                    _,
+                    MaybeConstant::ReferenceConstant {
+                        value,
+                        must_collapse,
+                    },
+                ] = inputs
+                && can_turn_to_lt(op, value).is_some()
+            {
+                // The constant is on the right, but we can turn the LE into an LT,
+                // so we can use the left operand as immediate.
+                must_collapse.replace(true);
             }
         }
 
@@ -136,6 +150,35 @@ fn is_commutative(op: &Operator) -> bool {
     )
 }
 
+/// If op is a GE or LE, can we turn "c >= x" or "x <= c" into "x < c + 1", where c + 1 fits in i16?
+pub fn can_turn_to_lt(op: &Operator, value: &WasmValue) -> Option<i16> {
+    match op {
+        // Signed case:
+        Operator::I32GeS | Operator::I64GeS | Operator::I32LeS | Operator::I64LeS => match value {
+            WasmValue::I32(v) => v.checked_add(1).and_then(|v| i16::try_from(v).ok()),
+            WasmValue::I64(v) => v.checked_add(1).and_then(|v| i16::try_from(v).ok()),
+            _ => None,
+        },
+        // Unsigned case:
+        // This is weird, because we need to represent an unsigned value as a sign-extended i16.
+        // So, we can represent (c+1) if it falls in the ranges:
+        // [0..0x7FFF] or [0xFFFF8000..0xFFFFFFFF], for u32
+        // [0..0x7FFF] or [0xFFFFFFFFFFFF8000..0xFFFFFFFFFFFFFFFF], for u64
+        Operator::I32GeU | Operator::I64GeU | Operator::I32LeU | Operator::I64LeU => match value {
+            WasmValue::I32(v) => {
+                let uv = *v as u32;
+                uv.checked_add(1).and_then(|v| i16::try_from(v as i32).ok())
+            }
+            WasmValue::I64(v) => {
+                let uv = *v as u64;
+                uv.checked_add(1).and_then(|v| i16::try_from(v as i64).ok())
+            }
+            _ => None,
+        },
+        _ => None,
+    }
+}
+
 fn can_be_i16(value: &WasmValue) -> bool {
     match value {
         WasmValue::I32(v) => i16::try_from(*v).is_ok(),

integration/src/womir_translation.rs

@@ -1,6 +1,6 @@
 use std::{collections::HashMap, ops::Range, vec};
 
-use crate::{instruction_builder as ib, to_field::ToField};
+use crate::{const_collapse::can_turn_to_lt, instruction_builder as ib, to_field::ToField};
 use arrayvec::ArrayVec;
 use itertools::Itertools;
 use openvm_circuit::{
@@ -888,23 +888,32 @@ impl<'a, F: PrimeField32> Settings<'a> for OpenVMSettings<F> {
                     }
                     [WasmOpInput::Register(reg), WasmOpInput::Constant(c)]
                     | [WasmOpInput::Constant(c), WasmOpInput::Register(reg)] => {
-                        // This is a little confusing, because GE and LE are using the same opcode,
-                        // but this is correct, as the operands are reversed between GE and LE.
-                        // Const folding step guarantees that only GE will have the constant on
-                        // the right side, and only LE will have the constant on the left side.
-                        //
-                        // In another words: the order of the operands will guarantee that the inverse
-                        // operation is always a less-than comparison: reg_value < constant_value.
-                        let inverse_op = match op {
+                        let opcode = match op {
                             Op::I32GeS | Op::I32LeS => LessThanOpcode::SLT.global_opcode(),
                             Op::I32GeU | Op::I32LeU => LessThanOpcode::SLTU.global_opcode(),
                             Op::I64GeS | Op::I64LeS => LessThan64Opcode::SLT.global_opcode(),
                             Op::I64GeU | Op::I64LeU => LessThan64Opcode::SLTU.global_opcode(),
                             _ => unreachable!(),
                         };
 
+                        // Check whether this is the case where "c >= r" or "r <= c" can be turned into "r < c + 1".
+                        if ((matches!(op, Op::I32GeS | Op::I32GeU | Op::I64GeS | Op::I64GeU)
+                            && matches!(&inputs[0], WasmOpInput::Constant(_)))
+                            || (matches!(op, Op::I32LeS | Op::I32LeU | Op::I64LeS | Op::I64LeU)
+                                && matches!(&inputs[0], WasmOpInput::Register(_))))
+                            && let Some(inc_c) = can_turn_to_lt(&op, c)
+                        {
+                            return Directive::Instruction(ib::instr_i(
+                                opcode.as_usize(),
+                                output,
+                                reg.start as usize,
+                                i16_to_imm_field(inc_c),
+                            ))
+                            .into();
+                        }
+
                         let c = const_i16_as_field(c);
-                        ib::instr_i(inverse_op.as_usize(), inverse_result, reg.start as usize, c)
+                        ib::instr_i(opcode.as_usize(), inverse_result, reg.start as usize, c)
                     }
                     _ => unreachable!("combination of inputs not possible for GE/LE operations"),
                 };
@@ -2155,7 +2164,11 @@ fn const_i16_as_field<F: PrimeField32>(value: &WasmValue) -> F {
         WasmValue::I64(v) => v as i16,
         _ => panic!("not valid i16"),
     };
-    F::from_canonical_u32((c as i32 as u32) & 0xff_ff_ff)
+    i16_to_imm_field(c)
+}
+
+fn i16_to_imm_field<F: PrimeField32>(value: i16) -> F {
+    F::from_canonical_u32((value as i32 as u32) & 0xff_ff_ff)
 }
 
 fn mem_offset<F: PrimeField32>(memarg: MemArg, c: &Ctx<F>) -> i32 {