Set Opcode Flags

autoprecompiles/Cargo.toml

@@ -14,6 +14,7 @@ powdr-parser.workspace = true
 powdr-parser-util.workspace = true
 powdr-pil-analyzer.workspace = true
 powdr-pilopt.workspace = true
+powdr-executor.workspace = true
 
 itertools = "0.13"
 lazy_static = "1.4.0"

autoprecompiles/src/lib.rs

@@ -5,21 +5,21 @@ use powdr_ast::parsed::asm::Part;
 use powdr_ast::{
     analyzed::{
         AlgebraicBinaryOperation, AlgebraicBinaryOperator, AlgebraicExpression, AlgebraicReference,
-        AlgebraicUnaryOperator, Analyzed, BusInteractionIdentity, Identity, PolyID,
-        PolynomialIdentity, PolynomialType,
+        AlgebraicUnaryOperator, Analyzed, BusInteractionIdentity, Identity, PolynomialIdentity,
     },
     parsed::{
         asm::SymbolPath, visitor::AllChildren, ArrayLiteral, BinaryOperation, BinaryOperator,
         Expression, FunctionCall, NamespacedPolynomialReference, Number, PILFile, PilStatement,
         UnaryOperation, UnaryOperator,
     },
 };
+use powdr_executor::witgen::evaluators::symbolic_evaluator::SymbolicEvaluator;
+use powdr_executor::witgen::{AlgebraicVariable, PartialExpressionEvaluator};
 use powdr_parser_util::SourceRef;
 use powdr_pil_analyzer::analyze_ast;
 use powdr_pilopt::optimize;
 use serde::{Deserialize, Serialize};
 use std::collections::{BTreeMap, BTreeSet, HashSet};
-use std::convert::TryFrom;
 
 use powdr_number::{BigUint, FieldElement, LargeInt};
 use powdr_pilopt::simplify_expression;
@@ -960,8 +960,7 @@ pub fn generate_precompile<T: FieldElement>(
                     sub_map_loadstore.extend(loadstore_chip_info(&machine, instr.opcode));
                 }
 
-                let opcode_a = AlgebraicExpression::Number((instr.opcode as u64).into());
-                local_constraints.push((pc_lookup.op - opcode_a).into());
+                add_opcode_constraints(&mut local_constraints, instr.opcode, &pc_lookup.op);
 
                 assert_eq!(instr.args.len(), pc_lookup.args.len());
                 if optimize {
@@ -1149,6 +1148,93 @@ pub fn generate_precompile<T: FieldElement>(
     )
 }
 
+fn add_opcode_constraints<T: FieldElement>(
+    constraints: &mut Vec<SymbolicConstraint<T>>,
+    opcode: usize,
+    expected_opcode: &AlgebraicExpression<T>,
+) {
+    let opcode_a = AlgebraicExpression::Number((opcode as u64).into());
+    match try_compute_opcode_map(expected_opcode) {
+        Ok(opcode_to_flag) => {
+            let active_flag = opcode_to_flag.get(&(opcode as u64)).unwrap();
+            for flag_ref in opcode_to_flag.values() {
+                let expected_value = if flag_ref == active_flag {
+                    AlgebraicExpression::Number(1u32.into())
+                } else {
+                    AlgebraicExpression::Number(0u32.into())
+                };
+                let flag_expr = AlgebraicExpression::Reference(flag_ref.clone());
+                let constraint = flag_expr - expected_value;
+                constraints.push(constraint.into());
+            }
+        }
+        Err(_) => {
+            // We were not able to extract the flags, so we keep them as witness columns
+            // and add a constraint that the expected opcode needs to equal the compile-time opcode.
+            constraints.push((expected_opcode.clone() - opcode_a).into());
+        }
+    }
+}
+
+fn try_compute_opcode_map<T: FieldElement>(
+    expected_opcode: &AlgebraicExpression<T>,
+) -> Result<BTreeMap<u64, AlgebraicReference>, ()> {
+    // Parse the expected opcode as an algebraic expression:
+    // flag1 * c1 + flag2 * c2 + ... + offset
+    let imm = BTreeMap::new();
+    let affine_expression = PartialExpressionEvaluator::new(SymbolicEvaluator, &imm)
+        .evaluate(expected_opcode)
+        .map_err(|_| ())?;
+    println!("GGG affine_expression = {affine_expression}");
+
+    // The above would not include any entry for `flag * 0` or `0 * flag`.
+    // If it exists, we collect it here.
+    let zero = AlgebraicExpression::Number(0u32.into());
+    let zero_flag = expected_opcode
+        .all_children()
+        .find_map(|expr| match expr {
+            AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                op: AlgebraicBinaryOperator::Mul,
+                left,
+                right,
+            }) => {
+                if **left == zero {
+                    Some((**right).clone())
+                } else if **right == zero {
+                    Some((**left).clone())
+                } else {
+                    None
+                }
+            }
+            _ => None,
+        })
+        .and_then(|expr| match expr {
+            AlgebraicExpression::Reference(reference) => Some(reference),
+            _ => None,
+        });
+
+    let offset = affine_expression.offset();
+
+    if !offset.is_zero() && zero_flag.is_none() {
+        // We didn't find any flag with factor zero, and the offset is not zero.
+        // Probably something went wrong.
+        return Err(());
+    }
+
+    Ok(affine_expression
+        .nonzero_coefficients()
+        .map(|(k, factor)| {
+            let opcode = (offset + *factor).to_degree();
+            let flag = match k {
+                AlgebraicVariable::Column(column) => (*column).clone(),
+                AlgebraicVariable::Public(_) => unreachable!(),
+            };
+            (opcode, flag)
+        })
+        .chain(zero_flag.map(|flag| (offset.to_degree(), flag.clone())))
+        .collect())
+}
+
 /// Use a SymbolicMachine to create a PILFile and run powdr's optimizations on it.
 /// Then, translate the optimized constraints and interactions back to a SymbolicMachine
 fn powdr_optimize<P: FieldElement>(symbolic_machine: SymbolicMachine<P>) -> SymbolicMachine<P> {