Reduce RefCell::borrow_mut() calls in Collector::enforce_zero

[ConstanceBeguier]

Crate

ragu_circuits

Severity

low

Repo commit

c9649de

Description

Collector::enforce_zero builds linear constraints by repeatedly calling borrow_mut() inside TermEnforcer::add_term. This results in one dynamic RefCell borrow per term, which adds avoidable overhead.

Code

File name: ragu_circuits/src/sy.rs

struct TermEnforcer<'table, 'sy, F: Field, R: Rank>(
    &'table RefCell<VirtualTable<'sy, F, R>>,
    Coeff<F>,
);

impl<'table, 'sy, F: Field, R: Rank> LinearExpression<Wire<'table, 'sy, F, R>, F>
    for TermEnforcer<'table, 'sy, F, R>
{
    fn add_term(self, wire: &Wire<'table, 'sy, F, R>, coeff: Coeff<F>) -> Self {
        self.0.borrow_mut().add(wire.index, coeff * self.1);
        self
    }
    ...
}

impl<'table, 'sy, F: Field, R: Rank> Driver<'table> for Collector<'table, 'sy, F, R> {
    ...
    fn enforce_zero(&mut self, lc: impl Fn(Self::LCenforce) -> Self::LCenforce) -> Result<()> {
        let q = self.linear_constraints;
        if q == R::num_coeffs() {
            return Err(Error::LinearBoundExceeded(R::num_coeffs()));
        }
        self.linear_constraints += 1;

        lc(TermEnforcer(
            self.virtual_table,
            Coeff::Arbitrary(self.current_y),
        ));

        self.current_y *= self.y_inv;

        Ok(())
    }
    ...

}

Recommendations

use core::cell::RefMut;

struct TermEnforcer<'table, 'sy, F: Field, R: Rank>(
    RefMut<'table, VirtualTable<'sy, F, R>>,
    Coeff<F>,
);

impl<'table, 'sy, F: Field, R: Rank> LinearExpression<Wire<'table, 'sy, F, R>, F>
    for TermEnforcer<'table, 'sy, F, R>
{
    fn add_term(mut self, wire: &Wire<'table, 'sy, F, R>, coeff: Coeff<F>) -> Self {
        // No repeated borrow_mut(): we already own the RefMut
        self.0.add(wire.index, coeff * self.1);
        self
    }
    ...
}

impl<'table, 'sy, F: Field, R: Rank> Driver<'table> for Collector<'table, 'sy, F, R> {
    ...
    fn enforce_zero(&mut self, lc: impl Fn(Self::LCenforce) -> Self::LCenforce) -> Result<()> {
        let q = self.linear_constraints;
        if q == R::num_coeffs() {
            return Err(Error::LinearBoundExceeded(R::num_coeffs()));
        }
        self.linear_constraints += 1;

        let table = self.virtual_table.borrow_mut(); // single borrow per constraint
        lc(TermEnforcer(table, Coeff::Arbitrary(self.current_y)));

        self.current_y *= self.y_inv;

        Ok(())
    }
    ...

}