feat(avm)!: support expressions in lookups (part 1)

barretenberg/cpp/src/barretenberg/vm2/common/expression.hpp

@@ -0,0 +1,113 @@
+#pragma once
+
+#include "barretenberg/vm2/common/field.hpp"
+#include "barretenberg/vm2/generated/columns.hpp"
+
+namespace bb::avm2 {
+
+/**
+ * @brief Relation expression representing a constant value.
+ */
+struct ConstantExpression {
+    FF value;
+    static constexpr bool __IS_EXPRESSION__ = true;
+
+    constexpr explicit ConstantExpression(FF value)
+        : value(value)
+    {}
+};
+
+/**
+ * @brief Relation expression representing a column.
+ */
+struct ColumnExpression {
+    ColumnAndShifts column;
+    static constexpr bool __IS_EXPRESSION__ = true;
+
+    constexpr explicit ColumnExpression(ColumnAndShifts column)
+        : column(column)
+    {}
+};
+
+/**
+ * @brief Relation expression representing a unary negation.
+ */
+template <typename InnerType> struct NegExpression {
+    InnerType inner;
+    static constexpr bool __IS_EXPRESSION__ = true;
+
+    constexpr explicit NegExpression(InnerType inner)
+        : inner(std::forward<InnerType>(inner))
+    {}
+};
+
+/**
+ * @brief Relation expression representing an addition.
+ */
+template <typename LhsType, typename RhsType> struct AddExpression {
+    LhsType lhs;
+    RhsType rhs;
+
+    constexpr explicit AddExpression(LhsType&& lhs, RhsType&& rhs)
+        : lhs(std::forward<LhsType>(lhs))
+        , rhs(std::forward<RhsType>(rhs))
+    {}
+};
+
+/**
+ * @brief Relation expression representing a subtraction.
+ */
+template <typename LhsType, typename RhsType> struct SubExpression {
+    LhsType lhs;
+    RhsType rhs;
+    static constexpr bool __IS_EXPRESSION__ = true;
+
+    constexpr explicit SubExpression(LhsType&& lhs, RhsType&& rhs)
+        : lhs(std::forward<LhsType>(lhs))
+        , rhs(std::forward<RhsType>(rhs))
+    {}
+};
+
+/**
+ * @brief Relation expression representing a multiplication.
+ */
+template <typename LhsType, typename RhsType> struct MulExpression {
+    LhsType lhs;
+    RhsType rhs;
+    static constexpr bool __IS_EXPRESSION__ = true;
+
+    constexpr explicit MulExpression(LhsType&& lhs, RhsType&& rhs)
+        : lhs(std::forward<LhsType>(lhs))
+        , rhs(std::forward<RhsType>(rhs))
+    {}
+};
+
+template <typename InnerType>
+constexpr NegExpression<InnerType> operator-(InnerType&& inner)
+    requires(InnerType::__IS_EXPRESSION__)
+{
+    return NegExpression<InnerType>(std::forward<InnerType>(inner));
+}
+
+template <typename LhsType, typename RhsType>
+constexpr AddExpression<LhsType, RhsType> operator+(LhsType&& lhs, RhsType&& rhs)
+    requires(LhsType::__IS_EXPRESSION__ && RhsType::__IS_EXPRESSION__)
+{
+    return AddExpression<LhsType, RhsType>(std::forward<LhsType>(lhs), std::forward<RhsType>(rhs));
+}
+
+template <typename LhsType, typename RhsType>
+constexpr SubExpression<LhsType, RhsType> operator-(LhsType&& lhs, RhsType&& rhs)
+    requires(LhsType::__IS_EXPRESSION__ && RhsType::__IS_EXPRESSION__)
+{
+    return SubExpression<LhsType, RhsType>(std::forward<LhsType>(lhs), std::forward<RhsType>(rhs));
+}
+
+template <typename LhsType, typename RhsType>
+constexpr MulExpression<LhsType, RhsType> operator*(LhsType&& lhs, RhsType&& rhs)
+    requires(LhsType::__IS_EXPRESSION__ && RhsType::__IS_EXPRESSION__)
+{
+    return MulExpression<LhsType, RhsType>(std::forward<LhsType>(lhs), std::forward<RhsType>(rhs));
+}
+
+} // namespace bb::avm2

barretenberg/cpp/src/barretenberg/vm2/common/expression_evaluation.hpp

@@ -0,0 +1,80 @@
+#pragma once
+
+#include "barretenberg/vm2/common/expression.hpp"
+
+namespace bb::avm2 {
+
+/**
+ * @brief Returns the degree of an expression.
+ *
+ * It can be used at copile time.
+ */
+class DegreeExpressionEvaluator {
+  public:
+    static constexpr size_t evaluate(const ConstantExpression&) { return 0; }
+    static constexpr size_t evaluate(const ColumnExpression&) { return 1; }
+
+    template <typename InnerType> static constexpr size_t evaluate(const NegExpression<InnerType>& expression)
+    {
+        return evaluate(expression.inner);
+    }
+
+    template <typename LhsType, typename RhsType>
+    static constexpr size_t evaluate(const AddExpression<LhsType, RhsType>& expression)
+    {
+        return std::max(evaluate(expression.lhs), evaluate(expression.rhs));
+    }
+
+    template <typename LhsType, typename RhsType>
+    static constexpr size_t evaluate(const SubExpression<LhsType, RhsType>& expression)
+    {
+        return std::max(evaluate(expression.lhs), evaluate(expression.rhs));
+    }
+
+    template <typename LhsType, typename RhsType>
+    static constexpr size_t evaluate(const MulExpression<LhsType, RhsType>& expression)
+    {
+        return evaluate(expression.lhs) + evaluate(expression.rhs);
+    }
+};
+
+/**
+ * @brief Evaluates an expression to a value.
+ */
+template <typename ValueType> class ValueExpressionEvaluator {
+  public:
+    using ColumnEvaluator = std::function<ValueType(const ColumnExpression&)>;
+    ValueExpressionEvaluator(ColumnEvaluator column_evaluator)
+        : column_evaluator_(std::move(column_evaluator))
+    {}
+
+    ValueType evaluate(const ConstantExpression& expression) const { return expression.value; }
+    ValueType evaluate(const ColumnExpression& expression) const { return column_evaluator_(expression); }
+
+    template <typename InnerType> ValueType evaluate(const NegExpression<InnerType>& expression) const
+    {
+        return -evaluate(expression.inner);
+    }
+
+    template <typename LhsType, typename RhsType>
+    ValueType evaluate(const AddExpression<LhsType, RhsType>& expression) const
+    {
+        return evaluate(expression.lhs) + evaluate(expression.rhs);
+    }
+
+    template <typename LhsType, typename RhsType>
+    ValueType evaluate(const SubExpression<LhsType, RhsType>& expression) const
+    {
+        return evaluate(expression.lhs) - evaluate(expression.rhs);
+    }
+
+    template <typename LhsType, typename RhsType>
+    ValueType evaluate(const MulExpression<LhsType, RhsType>& expression) const
+    {
+        return evaluate(expression.lhs) * evaluate(expression.rhs);
+    }
+
+  private:
+    ColumnEvaluator column_evaluator_;
+};
+} // namespace bb::avm2

barretenberg/cpp/src/barretenberg/vm2/common/expression_evaluation.test.cpp

@@ -0,0 +1,22 @@
+#include "barretenberg/vm2/common/expression_evaluation.hpp"
+
+namespace bb::avm2 {
+
+static_assert(DegreeExpressionEvaluator::evaluate(ConstantExpression(33) +
+                                                  ColumnExpression(ColumnAndShifts::memory_clk)) == 1);
+static_assert(DegreeExpressionEvaluator::evaluate(ColumnExpression(ColumnAndShifts::context_stack_is_static) *
+                                                  ColumnExpression(ColumnAndShifts::memory_clk)) == 2);
+static_assert(DegreeExpressionEvaluator::evaluate(-ConstantExpression(33)) == 0);
+static_assert(DegreeExpressionEvaluator::evaluate(-ColumnExpression(ColumnAndShifts::memory_clk)) == 1);
+static_assert(DegreeExpressionEvaluator::evaluate(-(-ConstantExpression(33))) == 0);
+static_assert(DegreeExpressionEvaluator::evaluate(-(-ColumnExpression(ColumnAndShifts::memory_clk))) == 1);
+static_assert(DegreeExpressionEvaluator::evaluate(ConstantExpression(33) -
+                                                  ColumnExpression(ColumnAndShifts::memory_clk)) == 1);
+static_assert(DegreeExpressionEvaluator::evaluate(ColumnExpression(ColumnAndShifts::context_stack_is_static) -
+                                                  ColumnExpression(ColumnAndShifts::memory_clk)) == 1);
+static_assert(DegreeExpressionEvaluator::evaluate(ConstantExpression(33) *
+                                                  ColumnExpression(ColumnAndShifts::memory_clk)) == 1);
+static_assert(DegreeExpressionEvaluator::evaluate(ColumnExpression(ColumnAndShifts::context_stack_is_static) *
+                                                  ColumnExpression(ColumnAndShifts::memory_clk)) == 2);
+
+} // namespace bb::avm2

barretenberg/cpp/src/barretenberg/vm2/constraining/relations/interactions_base.hpp

@@ -1,10 +1,14 @@
 #pragma once
 
 #include <cstddef>
+#include <type_traits>
+#include <utility>
 
 #include "barretenberg/common/tuple.hpp"
 #include "barretenberg/relations/generic_lookup/generic_lookup_relation.hpp"
 #include "barretenberg/relations/generic_permutation/generic_permutation_relation.hpp"
+#include "barretenberg/vm2/common/expression.hpp"
+#include "barretenberg/vm2/common/expression_evaluation.hpp"
 #include "barretenberg/vm2/generated/columns.hpp"
 
 namespace bb::avm2 {
@@ -20,34 +24,55 @@ template <typename Settings_> struct lookup_settings : public Settings_ {
     static constexpr size_t READ_TERM_DEGREE = 0;
     static constexpr size_t WRITE_TERM_DEGREE = 0;
 
+    // TODO(fcarreiro): support any degree.
+    // static constexpr size_t MAX_DEGREE = std::max(
+    //     { DegreeExpressionEvaluator::evaluate(Settings_::SRC_SELECTOR_EXPR),
+    //       DegreeExpressionEvaluator::evaluate(Settings_::DST_SELECTOR_EXPR),
+    //       []<size_t... ISource, size_t... IDest>(std::index_sequence<ISource...>, std::index_sequence<IDest...>) {
+    //           return std::max(DegreeExpressionEvaluator::evaluate(std::get<ISource>(Settings_::SRC_EXPRS))...,
+    //                           DegreeExpressionEvaluator::evaluate(std::get<IDest>(Settings_::DST_EXPRS))...);
+    //       }(std::make_index_sequence<std::tuple_size_v<decltype(Settings_::SRC_EXPRS)>>{},
+    //         std::make_index_sequence<std::tuple_size_v<decltype(Settings_::DST_EXPRS)>>{}) });
+    // // The GenericLookupRelation only safely supports degree <= 1 for each expression.
+    // static_assert(MAX_DEGREE <= 1, "MAX_DEGREE of lookup must be less than 1");
+
     template <typename AllEntities> static inline auto inverse_polynomial_is_computed_at_row(const AllEntities& in)
     {
-        return (in.get(static_cast<ColumnAndShifts>(Settings_::SRC_SELECTOR)) == 1 ||
-                in.get(static_cast<ColumnAndShifts>(Settings_::DST_SELECTOR)) == 1);
+        using RetType =
+            decltype(std::declval<const std::remove_reference_t<AllEntities>&>().get(std::declval<ColumnAndShifts>()));
+        auto ev =
+            ValueExpressionEvaluator<RetType>([&in](const ColumnExpression& expr) { return in.get(expr.column); });
+        return (ev.evaluate(Settings_::SRC_SELECTOR_EXPR) == 1 || ev.evaluate(Settings_::DST_SELECTOR_EXPR) == 1);
     }
 
     template <typename Accumulator, typename AllEntities>
     static inline auto compute_inverse_exists(const AllEntities& in)
     {
         using View = typename Accumulator::View;
-        const auto is_operation = View(in.get(static_cast<ColumnAndShifts>(Settings_::SRC_SELECTOR)));
-        const auto is_table_entry = View(in.get(static_cast<ColumnAndShifts>(Settings_::DST_SELECTOR)));
+        auto ev =
+            ValueExpressionEvaluator<View>([&in](const ColumnExpression& expr) { return View(in.get(expr.column)); });
+        const auto is_operation = ev.evaluate(Settings_::SRC_SELECTOR_EXPR);
+        const auto is_table_entry = ev.evaluate(Settings_::DST_SELECTOR_EXPR);
         return (is_operation + is_table_entry - is_operation * is_table_entry);
     }
 
     template <typename AllEntities> static inline auto get_entities(AllEntities&& in)
     {
         return []<size_t... ISource, size_t... IDest>(
                    AllEntities&& in, std::index_sequence<ISource...>, std::index_sequence<IDest...>) {
+            using RetType = decltype(std::declval<const std::remove_reference_t<AllEntities>&>().get(
+                std::declval<ColumnAndShifts>()));
+            auto ev =
+                ValueExpressionEvaluator<RetType>([&in](const ColumnExpression& expr) { return in.get(expr.column); });
             return flat_tuple::forward_as_tuple(in.get(static_cast<ColumnAndShifts>(Settings_::INVERSES)),
                                                 in.get(static_cast<ColumnAndShifts>(Settings_::COUNTS)),
-                                                in.get(static_cast<ColumnAndShifts>(Settings_::SRC_SELECTOR)),
-                                                in.get(static_cast<ColumnAndShifts>(Settings_::DST_SELECTOR)),
-                                                in.get(Settings_::SRC_COLUMNS[ISource])...,
-                                                in.get(Settings_::DST_COLUMNS[IDest])...);
+                                                ev.evaluate(Settings_::SRC_SELECTOR_EXPR),
+                                                ev.evaluate(Settings_::DST_SELECTOR_EXPR),
+                                                ev.evaluate(std::get<ISource>(Settings_::SRC_EXPRS))...,
+                                                ev.evaluate(std::get<IDest>(Settings_::DST_EXPRS))...);
         }(std::forward<AllEntities>(in),
-               std::make_index_sequence<Settings_::SRC_COLUMNS.size()>{},
-               std::make_index_sequence<Settings_::DST_COLUMNS.size()>{});
+               std::make_index_sequence<std::tuple_size_v<decltype(Settings_::SRC_EXPRS)>>{},
+               std::make_index_sequence<std::tuple_size_v<decltype(Settings_::DST_EXPRS)>>{});
     }
 
     template <typename AllEntities> static inline auto get_const_entities(const AllEntities& in)