refactor: use stwo_air_utils for par air generation

[giladchase]

• follows current design in stwo-cairo for this opcode.
• inline write_trace_row
• add LookupData that wraps InteractionClaimGenerator fields, which also holds n_calls.
• fix edge cases if input is smaller than LOG_N_LANES

---

This change is 

[shaharsamocha7]

Reviewed 6 of 7 files at r1, all commit messages.
Reviewable status: 6 of 7 files reviewed, 3 unresolved discussions (waiting on @giladchase)

---

crates/prover/src/components/ret_opcode/prover.rs line 91 at r1 (raw file):

            c.iter()
                .for_each(|v| memory_trace_generator.add_inputs_simd(v))
        });

I think that in the latest version we add the inputs inside write_trace_simd

Code quote:

        lookup_data.memory_inputs.iter().for_each(|c| {
            c.iter()
                .for_each(|v| memory_trace_generator.add_inputs_simd(v))
        });

---

crates/prover/src/components/ret_opcode/prover.rs line 105 at r1 (raw file):

pub struct LookupData {
    pub memory_inputs: Vec<[PackedM31; N_MEMORY_CALLS]>,
    pub memory_outputs: Vec<[PackedQM31; N_MEMORY_CALLS]>,

what are memory outputs?

Code quote:

pub memory_outputs: Vec<[PackedQM31; N_MEMORY_CALLS]>,

---

crates/prover/src/components/ret_opcode/component.rs line 86 at r1 (raw file):

#[derive(Clone, Serialize, Deserialize)]
pub struct Claim {
    pub n_calls: usize,

We renamed it to n_rows in stwo-cairo if you want to be consistent

Code quote:

pub n_calls: usize,

[giladchase]

Reviewable status: 6 of 7 files reviewed, 2 unresolved discussions (waiting on @shaharsamocha7)

---

crates/prover/src/components/ret_opcode/prover.rs line 91 at r1 (raw file):

Previously, shaharsamocha7 wrote…

I think that in the latest version we add the inputs inside write_trace_simd

Might be misunderstanding, but looks like it's also being done in write_trace on main (add_inputs_simd does the same thing)

---

crates/prover/src/components/ret_opcode/prover.rs line 105 at r1 (raw file):

Previously, shaharsamocha7 wrote…

what are memory outputs?

Legacy stwo-cairo for holding results of memory calls during trace generation.
Ohad refactored them out two months ago.

I can fetch the refactor, but I'm not sure how big a job that is

[shaharsamocha7]

Reviewed 3 of 4 files at r2, all commit messages.
Reviewable status: 8 of 9 files reviewed, 1 unresolved discussion (waiting on @giladchase)

---

crates/prover/src/components/ret_opcode/prover.rs line 105 at r1 (raw file):

Previously, giladchase wrote…

Legacy stwo-cairo for holding results of memory calls during trace generation.
Ohad refactored them out two months ago.

I can fetch the refactor, but I'm not sure how big a job that is

I think that eventually we should do that because it doesn't make sense, could be a TODO for now

[giladchase]

Reviewable status: 8 of 9 files reviewed, 1 unresolved discussion (waiting on @shaharsamocha7)

---

crates/prover/src/components/ret_opcode/prover.rs line 105 at r1 (raw file):

Previously, shaharsamocha7 wrote…

I think that eventually we should do that because it doesn't make sense, could be a TODO for now

Added TODO

[shaharsamocha7]

Reviewed 1 of 1 files at r3, all commit messages.
Reviewable status: all files reviewed, 2 unresolved discussions (waiting on @giladchase)

---

crates/prover/src/components/ret_opcode/prover.rs line 87 at r3 (raw file):

                .resize(size, self.inputs.first().unwrap().clone());
            bit_reverse_coset_to_circle_domain_order(&mut self.inputs);
        }

FYI We removed this feature in stwo, we need to think if we also want to apply that change here

Code quote:

        let need_padding = n_rows != size;

        if need_padding {
            self.inputs
                .resize(size, self.inputs.first().unwrap().clone());
            bit_reverse_coset_to_circle_domain_order(&mut self.inputs);
        }

---

crates/prover/src/components/ret_opcode/prover.rs line 143 at r3 (raw file):

                col_gen.write_frac(i, PackedQM31::one(), denom);
            }
            col_gen.finalize_col();

I'm not sure this logic is correct, what exactly memory_inputs, memory_outputs has?

Code quote:

            let mut col_gen = logup_gen.new_col();
            for (i, (&addr, &output)) in zip_eq(
                &self.lookup_data.memory_inputs[col_index],
                &self.lookup_data.memory_outputs[col_index],
            )
            .enumerate()
            {
                let address_and_value = vec![addr, output.into_packed_m31s()[0]];
                let denom = lookup_elements.combine(&address_and_value);
                col_gen.write_frac(i, PackedQM31::one(), denom);
            }
            col_gen.finalize_col();

[giladchase]

Reviewable status: 6 of 10 files reviewed, 2 unresolved discussions (waiting on @Alon-Ti and @shaharsamocha7)

---

crates/prover/src/components/ret_opcode/prover.rs line 87 at r3 (raw file):

Previously, shaharsamocha7 wrote…

FYI We removed this feature in stwo, we need to think if we also want to apply that change here

Removed. If it comes up adding it won't be a problem

---

crates/prover/src/components/ret_opcode/prover.rs line 143 at r3 (raw file):

Previously, shaharsamocha7 wrote…

I'm not sure this logic is correct, what exactly memory_inputs, memory_outputs has?

NVM, i already started supported the non-legacy LookupData up the stack, so i backported it here since noones seems to remember how thouse inputs/outputs worked 😅

[Alon-Ti]

Reviewed 2 of 4 files at r4, 1 of 1 files at r5, all commit messages.
Reviewable status: 9 of 10 files reviewed, 5 unresolved discussions (waiting on @shaharsamocha7)

---

crates/prover/src/components/ret_opcode/component.rs line 14 at r5 (raw file):

pub const RET_N_TRACE_CELLS: usize = 5;
pub const RET_INSTRUCTION: M31 = M31::from_u32_unchecked(171);

Add a TODO to reallocate this number based on all instructions (maybe even set it to zero so that it's clear this is not an actual value).

---

crates/prover/src/components/ret_opcode/component.rs line 26 at r5 (raw file):

impl Eval {
    pub fn log_size(&self) -> u32 {
        std::cmp::max(self.claim.n_rows.next_power_of_two().ilog2(), LOG_N_LANES)

This is going to be everywhere, right? Maybe export it to a function?

---

crates/prover/src/components/ret_opcode/component.rs line 108 at r5 (raw file):

impl InteractionClaim {
    pub fn mix_into(&self, channel: &mut impl Channel) {
        let (total_sum, claimed_sum) = self.logup_sums;

Update stwo to get rid of this.

[giladchase]

Reviewable status: 9 of 10 files reviewed, 4 unresolved discussions (waiting on @Alon-Ti and @shaharsamocha7)

---

crates/prover/src/components/ret_opcode/component.rs line 14 at r5 (raw file):

Previously, Alon-Ti wrote…

Add a TODO to reallocate this number based on all instructions (maybe even set it to zero so that it's clear this is not an actual value).

Done, used the same todo from other PRs for easier grep