Skip to content

Commit

Permalink
added interaction scope
Browse files Browse the repository at this point in the history
  • Loading branch information
@kevjue
kevjue committed Aug 9, 2024
1 parent c1451e1 commit 4c46cbc
Show file tree
Hide file tree
Showing 5 changed files with 160 additions and 128 deletions.
191 changes: 101 additions & 90 deletions core/src/air/builder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,18 +17,23 @@ use crate::lookup::InteractionKind;
use crate::memory::MemoryAccessCols;
use crate::{bytes::ByteOpcode, memory::MemoryCols};

pub enum MessageScope {
Shard,
Global,
}

/// A Builder with the ability to encode the existance of interactions with other AIRs by sending
/// and receiving messages.
pub trait MessageBuilder<M> {
fn send(&mut self, message: M);
fn send(&mut self, message: M, scope: MessageScope);

fn receive(&mut self, message: M);
fn receive(&mut self, message: M, scope: MessageScope);
}

impl<AB: EmptyMessageBuilder, M> MessageBuilder<M> for AB {
fn send(&mut self, _message: M) {}
fn send(&mut self, _message: M, _scope: MessageScope) {}

fn receive(&mut self, _message: M) {}
fn receive(&mut self, _message: M, _scope: MessageScope) {}
}

/// A message builder for which sending and receiving messages is a no-op.
Expand Down Expand Up @@ -125,19 +130,22 @@ pub trait ByteAirBuilder: BaseAirBuilder {
channel: impl Into<Self::Expr>,
multiplicity: impl Into<Self::Expr>,
) {
self.send(AirInteraction::new(
vec![
opcode.into(),
a1.into(),
a2.into(),
b.into(),
c.into(),
shard.into(),
channel.into(),
],
multiplicity.into(),
InteractionKind::Byte,
));
self.send(
AirInteraction::new(
vec![
opcode.into(),
a1.into(),
a2.into(),
b.into(),
c.into(),
shard.into(),
channel.into(),
],
multiplicity.into(),
InteractionKind::Byte,
),
MessageScope::Global,
);
}

/// Receives a byte operation to be processed.
Expand Down Expand Up @@ -177,19 +185,22 @@ pub trait ByteAirBuilder: BaseAirBuilder {
channel: impl Into<Self::Expr>,
multiplicity: impl Into<Self::Expr>,
) {
self.receive(AirInteraction::new(
vec![
opcode.into(),
a1.into(),
a2.into(),
b.into(),
c.into(),
shard.into(),
channel.into(),
],
multiplicity.into(),
InteractionKind::Byte,
));
self.receive(
AirInteraction::new(
vec![
opcode.into(),
a1.into(),
a2.into(),
b.into(),
c.into(),
shard.into(),
channel.into(),
],
multiplicity.into(),
InteractionKind::Byte,
),
MessageScope::Global,
);
}
}

Expand Down Expand Up @@ -325,11 +336,10 @@ pub trait AluAirBuilder: BaseAirBuilder {
.chain(once(nonce.into()))
.collect();

self.send(AirInteraction::new(
values,
multiplicity.into(),
InteractionKind::Alu,
));
self.send(
AirInteraction::new(values, multiplicity.into(), InteractionKind::Alu),
MessageScope::Global,
);
}

/// Receives an ALU operation to be processed.
Expand All @@ -354,11 +364,10 @@ pub trait AluAirBuilder: BaseAirBuilder {
.chain(once(nonce.into()))
.collect();

self.receive(AirInteraction::new(
values,
multiplicity.into(),
InteractionKind::Alu,
));
self.receive(
AirInteraction::new(values, multiplicity.into(), InteractionKind::Alu),
MessageScope::Global,
);
}

/// Sends an syscall operation to be processed (with "ECALL" opcode).
Expand All @@ -374,19 +383,22 @@ pub trait AluAirBuilder: BaseAirBuilder {
arg2: impl Into<Self::Expr> + Clone,
multiplicity: impl Into<Self::Expr>,
) {
self.send(AirInteraction::new(
vec![
shard.clone().into(),
channel.clone().into(),
clk.clone().into(),
nonce.clone().into(),
syscall_id.clone().into(),
arg1.clone().into(),
arg2.clone().into(),
],
multiplicity.into(),
InteractionKind::Syscall,
));
self.send(
AirInteraction::new(
vec![
shard.clone().into(),
channel.clone().into(),
clk.clone().into(),
nonce.clone().into(),
syscall_id.clone().into(),
arg1.clone().into(),
arg2.clone().into(),
],
multiplicity.into(),
InteractionKind::Syscall,
),
MessageScope::Global,
);
}

/// Receives a syscall operation to be processed.
Expand All @@ -402,19 +414,22 @@ pub trait AluAirBuilder: BaseAirBuilder {
arg2: impl Into<Self::Expr> + Clone,
multiplicity: impl Into<Self::Expr>,
) {
self.receive(AirInteraction::new(
vec![
shard.clone().into(),
channel.clone().into(),
clk.clone().into(),
nonce.clone().into(),
syscall_id.clone().into(),
arg1.clone().into(),
arg2.clone().into(),
],
multiplicity.into(),
InteractionKind::Syscall,
));
self.receive(
AirInteraction::new(
vec![
shard.clone().into(),
channel.clone().into(),
clk.clone().into(),
nonce.clone().into(),
syscall_id.clone().into(),
arg1.clone().into(),
arg2.clone().into(),
],
multiplicity.into(),
InteractionKind::Syscall,
),
MessageScope::Global,
);
}
}

Expand Down Expand Up @@ -466,18 +481,16 @@ pub trait MemoryAirBuilder: BaseAirBuilder {
.collect();

// The previous values get sent with multiplicity = 1, for "read".
self.send(AirInteraction::new(
prev_values,
do_check.clone(),
InteractionKind::Memory,
));
self.send(
AirInteraction::new(prev_values, do_check.clone(), InteractionKind::Memory),
MessageScope::Global,
);

// The current values get "received", i.e. multiplicity = -1
self.receive(AirInteraction::new(
current_values,
do_check.clone(),
InteractionKind::Memory,
));
self.receive(
AirInteraction::new(current_values, do_check.clone(), InteractionKind::Memory),
MessageScope::Global,
);
}

/// Constraints a memory read or write to a slice of `MemoryAccessCols`.
Expand Down Expand Up @@ -622,11 +635,10 @@ pub trait ProgramAirBuilder: BaseAirBuilder {
.chain(once(shard.into()))
.collect();

self.send(AirInteraction::new(
values,
multiplicity.into(),
InteractionKind::Program,
));
self.send(
AirInteraction::new(values, multiplicity.into(), InteractionKind::Program),
MessageScope::Global,
);
}

/// Receives an instruction.
Expand All @@ -645,11 +657,10 @@ pub trait ProgramAirBuilder: BaseAirBuilder {
.chain(once(shard.into()))
.collect();

self.receive(AirInteraction::new(
values,
multiplicity.into(),
InteractionKind::Program,
));
self.receive(
AirInteraction::new(values, multiplicity.into(), InteractionKind::Program),
MessageScope::Global,
);
}
}

Expand Down Expand Up @@ -714,12 +725,12 @@ pub trait SP1AirBuilder:
}

impl<'a, AB: AirBuilder + MessageBuilder<M>, M> MessageBuilder<M> for FilteredAirBuilder<'a, AB> {
fn send(&mut self, message: M) {
self.inner.send(message);
fn send(&mut self, message: M, scope: MessageScope) {
self.inner.send(message, scope);
}

fn receive(&mut self, message: M) {
self.inner.receive(message);
fn receive(&mut self, message: M, scope: MessageScope) {
self.inner.receive(message, scope);
}
}

Expand Down
47 changes: 26 additions & 21 deletions core/src/lookup/builder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ use p3_matrix::dense::RowMajorMatrix;
use p3_uni_stark::{Entry, SymbolicExpression, SymbolicVariable};

use crate::{
air::{AirInteraction, MessageBuilder},
air::{AirInteraction, MessageBuilder, MessageScope},
stark::PROOF_MAX_NUM_PVS,
};

Expand Down Expand Up @@ -91,7 +91,7 @@ impl<F: Field> PairBuilder for InteractionBuilder<F> {
}

impl<F: Field> MessageBuilder<AirInteraction<SymbolicExpression<F>>> for InteractionBuilder<F> {
fn send(&mut self, message: AirInteraction<SymbolicExpression<F>>) {
fn send(&mut self, message: AirInteraction<SymbolicExpression<F>>, scope: MessageScope) {
let values = message
.values
.into_iter()
Expand All @@ -101,10 +101,10 @@ impl<F: Field> MessageBuilder<AirInteraction<SymbolicExpression<F>>> for Interac
let multiplicity = symbolic_to_virtual_pair(&message.multiplicity);

self.sends
.push(Interaction::new(values, multiplicity, message.kind));
.push(Interaction::new(values, multiplicity, message.kind, scope));
}

fn receive(&mut self, message: AirInteraction<SymbolicExpression<F>>) {
fn receive(&mut self, message: AirInteraction<SymbolicExpression<F>>, scope: MessageScope) {
let values = message
.values
.into_iter()
Expand All @@ -114,7 +114,7 @@ impl<F: Field> MessageBuilder<AirInteraction<SymbolicExpression<F>>> for Interac
let multiplicity = symbolic_to_virtual_pair(&message.multiplicity);

self.receives
.push(Interaction::new(values, multiplicity, message.kind));
.push(Interaction::new(values, multiplicity, message.kind, scope));
}
}

Expand Down Expand Up @@ -250,22 +250,27 @@ mod tests {
let y = local[1];
let z = local[2];

builder.send(AirInteraction::new(
vec![x.into(), y.into()],
AB::F::from_canonical_u32(3).into(),
InteractionKind::Alu,
));
builder.send(AirInteraction::new(
vec![x + y, z.into()],
AB::F::from_canonical_u32(5).into(),
InteractionKind::Alu,
));

builder.receive(AirInteraction::new(
vec![x.into()],
y.into(),
InteractionKind::Byte,
));
builder.send(
AirInteraction::new(
vec![x.into(), y.into()],
AB::F::from_canonical_u32(3).into(),
InteractionKind::Alu,
),
MessageScope::Global,
);
builder.send(
AirInteraction::new(
vec![x + y, z.into()],
AB::F::from_canonical_u32(5).into(),
InteractionKind::Alu,
),
MessageScope::Global,
);

builder.receive(
AirInteraction::new(vec![x.into()], y.into(), InteractionKind::Byte),
MessageScope::Global,
);
}
}

Expand Down
5 changes: 5 additions & 0 deletions core/src/lookup/interaction.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,14 @@ use core::fmt::Display;
use p3_air::VirtualPairCol;
use p3_field::Field;

use crate::air::MessageScope;

/// An interaction for a lookup or a permutation argument.
pub struct Interaction<F: Field> {
pub values: Vec<VirtualPairCol<F>>,
pub multiplicity: VirtualPairCol<F>,
pub kind: InteractionKind,
pub scope: MessageScope,
}

/// The type of interaction for a lookup argument.
Expand Down Expand Up @@ -60,11 +63,13 @@ impl<F: Field> Interaction<F> {
values: Vec<VirtualPairCol<F>>,
multiplicity: VirtualPairCol<F>,
kind: InteractionKind,
scope: MessageScope,
) -> Self {
Self {
values,
multiplicity,
kind,
scope,
}
}

Expand Down
Loading

0 comments on commit 4c46cbc

Please sign in to comment.