diff --git a/executor/src/witgen/mod.rs b/executor/src/witgen/mod.rs index 4d62e30bc..e019996c9 100644 --- a/executor/src/witgen/mod.rs +++ b/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, degree: DegreeType, fixed_cols: FixedColumnMap>, witness_cols: WitnessColumnMap>, @@ -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, diff --git a/executor/src/witgen/query_processor.rs b/executor/src/witgen/query_processor.rs index fe4ebee05..e5240ff85 100644 --- a/executor/src/witgen/query_processor.rs +++ b/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> #[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>, 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( diff --git a/test_data/asm/sqrt.asm b/test_data/asm/sqrt.asm index c74ddffe7..0a856cd77 100644 --- a/test_data/asm/sqrt.asm +++ b/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; } } \ No newline at end of file