Merge pull request #1032 from powdr-labs/evaluate-functions-in-queries

`QueryProcessor`: Give access to functions

executor/src/witgen/mod.rs

@@ -182,6 +182,7 @@ impl<'a, 'b, T: FieldElement> WitnessGenerator<'a, 'b, T> {
 
 /// Data that is fixed for witness generation.
 pub struct FixedData<'a, T> {
+    analyzed: &'a Analyzed<T>,
     degree: DegreeType,
     fixed_cols: FixedColumnMap<FixedColumn<'a, T>>,
     witness_cols: WitnessColumnMap<WitnessColumn<'a, T>>,
@@ -234,6 +235,7 @@ impl<'a, T: FieldElement> FixedData<'a, T> {
         let fixed_cols =
             FixedColumnMap::from(fixed_col_values.iter().map(|(n, v)| FixedColumn::new(n, v)));
         FixedData {
+            analyzed,
             degree: analyzed.degree(),
             fixed_cols,
             witness_cols,

executor/src/witgen/query_processor.rs

@@ -1,7 +1,10 @@
 use std::{fmt::Display, rc::Rc};
 
 use num_traits::ToPrimitive;
-use powdr_ast::analyzed::{AlgebraicReference, Expression, PolyID, PolynomialType};
+use powdr_ast::analyzed::{
+    types::TypedExpression, AlgebraicReference, Expression, FunctionValueDefinition, PolyID,
+    PolynomialType,
+};
 use powdr_number::{DegreeType, FieldElement};
 use powdr_pil_analyzer::evaluator::{self, Custom, EvalError, SymbolLookup, Value};
 
@@ -91,7 +94,6 @@ impl<'a, 'b, T: FieldElement, QueryCallback: super::QueryCallback<T>>
 
 #[derive(Clone)]
 struct Symbols<'a, T: FieldElement> {
-    // TODO we should also provide access to non-column symbols.
     fixed_data: &'a FixedData<'a, T>,
     rows: &'a RowPair<'a, 'a, T>,
 }
@@ -100,9 +102,24 @@ impl<'a, T: FieldElement> SymbolLookup<'a, T, Reference<'a>> for Symbols<'a, T>
     fn lookup(&self, name: &'a str) -> Result<Value<'a, T, Reference<'a>>, EvalError> {
         match self.fixed_data.try_column_by_name(name) {
             Some(poly_id) => Ok(Value::Custom(Reference { name, poly_id })),
-            None => Err(EvalError::SymbolNotFound(format!(
-                "Symbol {name} not found."
-            ))),
+            None => match self.fixed_data.analyzed.definitions.get(&name.to_string()) {
+                Some((_, value)) => {
+                    let value = value
+                        .as_ref()
+                        .expect("Witness columns should have been found by try_column_by_name()");
+                    match value {
+                        FunctionValueDefinition::Expression(TypedExpression { e, ty: _ }) => {
+                            evaluator::evaluate(e, self)
+                        }
+                        _ => panic!(
+                            "Arrays and queries should have been found by try_column_by_name()"
+                        ),
+                    }
+                }
+                None => Err(EvalError::SymbolNotFound(format!(
+                    "Symbol {name} not found."
+                ))),
+            },
         }
     }
     fn eval_function_application(

test_data/asm/sqrt.asm

@@ -10,10 +10,13 @@ machine Sqrt(latch, operation_id) {
 
     // Witness generation is not smart enough to figure out that
     // there is a unique witness, so we provide it as a hint.
-    // This is a dummy example that hard-codes the answer for an input of 4.
-    // Once we have a sqrt function that we can run to compute the query result,
-    // this can be used to compute the hint from x.
-    col witness y(i) query ("hint", 2);
+    // This is a dummy example that hard-codes the answer for inputs 1 and 4.
+    let sqrt_hint: int -> int = |x| match x {
+        1 => 1,
+        4 => 2
+    };
+
+    col witness y(i) query ("hint", sqrt_hint(x(i)));
     
     y * y = x;
     
@@ -49,6 +52,9 @@ machine Main {
         A <== sqrt(4);
         assert_zero A - 2;
 
+        A <== sqrt(1);
+        assert_zero A - 1;
+
         return;
     }
 }