sp1_hypercube/ir/
output.rs

1use slop_air::AirBuilder;
2
3use crate::Word;
4
5use super::{ConstraintCompiler, ExprRef, FuncCtx, GLOBAL_AST};
6
7/// Trait for handling operation outputs in SP1.
8///
9/// This trait provides methods for allocating and assigning operation outputs.
10pub trait SP1OperationOutput<T> {
11    /// Allocates space for the output and returns the allocated value.
12    fn alloc() -> T;
13
14    /// Replace all contents with IrVar(Output...)
15    fn to_output(&self, ctx: &mut FuncCtx) -> T;
16
17    /// Assigns the output.
18    fn assign(&self, other: T);
19}
20
21impl SP1OperationOutput<()> for () {
22    fn alloc() {}
23
24    fn to_output(&self, _: &mut FuncCtx) {}
25
26    fn assign(&self, (): Self) {
27        // Does nothing
28    }
29}
30
31type F = <ConstraintCompiler as AirBuilder>::F;
32
33impl SP1OperationOutput<ExprRef<F>> for ExprRef<F> {
34    fn alloc() -> Self {
35        GLOBAL_AST.lock().unwrap().alloc()
36    }
37
38    fn to_output(&self, ctx: &mut FuncCtx) -> Self {
39        ExprRef::<F>::output_arg(ctx)
40    }
41
42    fn assign(&self, other: Self) {
43        GLOBAL_AST.lock().unwrap().assign(other, *self);
44    }
45}
46
47impl<const N: usize> SP1OperationOutput<[ExprRef<F>; N]> for [ExprRef<F>; N] {
48    fn alloc() -> Self {
49        GLOBAL_AST.lock().unwrap().alloc_array()
50    }
51
52    fn to_output(&self, ctx: &mut FuncCtx) -> Self {
53        core::array::from_fn(|_| ExprRef::<F>::output_arg(ctx))
54    }
55
56    fn assign(&self, other: Self) {
57        for (i, o) in self.iter().zip(other.iter()) {
58            GLOBAL_AST.lock().unwrap().assign(*i, *o);
59        }
60    }
61}
62
63impl SP1OperationOutput<Word<ExprRef<F>>> for Word<ExprRef<F>> {
64    fn alloc() -> Self {
65        let a0 = GLOBAL_AST.lock().unwrap().alloc();
66        let a1 = GLOBAL_AST.lock().unwrap().alloc();
67        let a2 = GLOBAL_AST.lock().unwrap().alloc();
68        let a3 = GLOBAL_AST.lock().unwrap().alloc();
69        Word([a0, a1, a2, a3])
70    }
71
72    fn to_output(&self, ctx: &mut FuncCtx) -> Self {
73        ExprRef::<F>::output_from_struct(ctx)
74    }
75
76    fn assign(&self, other: Self) {
77        for (i, o) in self.0.iter().zip(other.0.iter()) {
78            GLOBAL_AST.lock().unwrap().assign(*i, *o);
79        }
80    }
81}
sp1_hypercube/ir/output.rs

sp1_hypercube/ir/
output.rs
