From 6da7c6c4763212b297a93d4fe2f0271fc5d198b3 Mon Sep 17 00:00:00 2001
From: malatrax <simon.malatrait@grenoble-inp.org>
Date: Wed, 11 Dec 2024 10:24:30 +0100
Subject: [PATCH] feat: add I/O lookup elements

---
 .../brainfuck_prover/src/brainfuck_air/mod.rs |  9 ++-
 .../src/components/io/table.rs                | 68 +++++++++++++++++--
 2 files changed, 70 insertions(+), 7 deletions(-)

diff --git a/crates/brainfuck_prover/src/brainfuck_air/mod.rs b/crates/brainfuck_prover/src/brainfuck_air/mod.rs
index 0a3c0af..df61c64 100644
--- a/crates/brainfuck_prover/src/brainfuck_air/mod.rs
+++ b/crates/brainfuck_prover/src/brainfuck_air/mod.rs
@@ -1,4 +1,5 @@
 use crate::components::{
+    io::table::IoElements,
     memory::{
         self,
         component::{MemoryComponent, MemoryEval},
@@ -62,6 +63,8 @@ impl BrainfuckClaim {
 /// All the interaction elements (drawn from the channel)
 /// required by the various components during the interaction phase.
 pub struct BrainfuckInteractionElements {
+    pub input_lookup_elements: IoElements,
+    pub output_lookup_elements: IoElements,
     pub memory_lookup_elements: MemoryElements,
 }
 
@@ -69,7 +72,11 @@ impl BrainfuckInteractionElements {
     /// Draw all the interaction elements needed for
     /// all the components of the Brainfuck ZK-VM.
     pub fn draw(channel: &mut impl Channel) -> Self {
-        Self { memory_lookup_elements: MemoryElements::draw(channel) }
+        Self {
+            input_lookup_elements: IoElements::draw(channel),
+            output_lookup_elements: IoElements::draw(channel),
+            memory_lookup_elements: MemoryElements::draw(channel),
+        }
     }
 }
 
diff --git a/crates/brainfuck_prover/src/components/io/table.rs b/crates/brainfuck_prover/src/components/io/table.rs
index 4796d80..e7998a8 100644
--- a/crates/brainfuck_prover/src/components/io/table.rs
+++ b/crates/brainfuck_prover/src/components/io/table.rs
@@ -1,12 +1,16 @@
 use crate::components::{IoClaim, TraceColumn, TraceEval};
 use brainfuck_vm::{instruction::InstructionType, registers::Registers};
-use stwo_prover::core::{
-    backend::{
-        simd::{column::BaseColumn, m31::LOG_N_LANES},
-        Column,
+use stwo_prover::{
+    constraint_framework::{logup::LookupElements, Relation, RelationEFTraitBound},
+    core::{
+        backend::{
+            simd::{column::BaseColumn, m31::LOG_N_LANES},
+            Column,
+        },
+        channel::Channel,
+        fields::m31::BaseField,
+        poly::circle::{CanonicCoset, CircleEvaluation},
     },
-    fields::m31::BaseField,
-    poly::circle::{CanonicCoset, CircleEvaluation},
 };
 
 /// Represents a single row in the I/O Table.
@@ -180,6 +184,58 @@ impl TraceColumn for IoColumn {
     }
 }
 
+/// The number of random elements necessary for the I/O lookup arguments.
+const IO_LOOKUP_ELEMENTS: usize = 1;
+
+/// The interaction elements are drawn for the extension column of the I/O components.
+///
+/// The logUp protocol uses these elements to combine the values of the different
+/// registers of the main trace to create a random linear combination
+/// of them, and use it in the denominator of the fractions in the logUp protocol.
+///
+/// There is a single lookup element in the I/O component: `mv`.
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct IoElements(LookupElements<IO_LOOKUP_ELEMENTS>);
+
+impl IoElements {
+    /// Provides dummy lookup elements.
+    pub fn dummy() -> Self {
+        Self(LookupElements::dummy())
+    }
+
+    /// Draw random elements from the Fiat-Shamir [`Channel`].
+    ///
+    /// These elements are randomly secured, and will be use
+    /// to generate the interaction trace with the logUp protocol.
+    pub fn draw(channel: &mut impl Channel) -> Self {
+        Self(LookupElements::draw(channel))
+    }
+}
+
+impl<F: Clone, EF: RelationEFTraitBound<F>> Relation<F, EF> for IoElements {
+    /// Combine multiple values from a basefield (e.g. [`BaseField`])
+    /// and combine them to a value from an extension field (e.g. [`PackedSecureField`])
+    ///
+    /// This is used when computing the interaction values from the main trace values.
+    fn combine(&self, values: &[F]) -> EF {
+        values
+            .iter()
+            .zip(self.0.alpha_powers)
+            .fold(EF::zero(), |acc, (value, power)| acc + EF::from(power) * value.clone()) -
+            self.0.z.into()
+    }
+
+    /// Returns the name of the struct.
+    fn get_name(&self) -> &str {
+        stringify!(IoElements)
+    }
+
+    /// Returns the number interaction elements.
+    fn get_size(&self) -> usize {
+        IO_LOOKUP_ELEMENTS
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;