feat: adds api.MAC(..)

frontend/api.go

@@ -30,6 +30,11 @@ type API interface {
 	// Add returns res = i1+i2+...in
 	Add(i1, i2 Variable, in ...Variable) Variable
 
+	// MulAcc sets and return a = a + (b*c)
+	// ! may mutate a without allocating a new result
+	// ! always use MulAcc(...) result for correctness
+	MulAcc(a, b, c Variable) Variable
+
 	// Neg returns -i
 	Neg(i1 Variable) Variable
 

frontend/builder.go

@@ -53,7 +53,7 @@ type Compiler interface {
 	// Commit returns a commitment to the given variables, to be used as initial randomness in
 	// Fiat-Shamir when the statement to prove is particularly large.
 	// TODO cite paper
-	// This API is experimental
+	// ! Experimental
 	// TENTATIVE: Functions regarding fiat-shamir-ed proofs over enormous statements  TODO finalize
 	Commit(...Variable) (Variable, error)
 }

frontend/cs/r1cs/api.go

@@ -40,19 +40,69 @@ import (
 func (builder *builder) Add(i1, i2 frontend.Variable, in ...frontend.Variable) frontend.Variable {
 	// extract frontend.Variables from input
 	vars, s := builder.toVariables(append([]frontend.Variable{i1, i2}, in...)...)
-	return builder.add(vars, false, s)
+	return builder.add(vars, false, s, nil)
+}
+
+func (builder *builder) MulAcc(a, b, c frontend.Variable) frontend.Variable {
+	// do the multiplication into builder.mbuf1
+	mulBC := func() {
+		// reset the buffer
+		builder.mbuf1 = builder.mbuf1[:0]
+
+		n1, v1Constant := builder.constantValue(b)
+		n2, v2Constant := builder.constantValue(c)
+
+		// v1 and v2 are both unknown, this is the only case we add a constraint
+		if !v1Constant && !v2Constant {
+			res := builder.newInternalVariable()
+			builder.cs.AddConstraint(builder.newR1C(b, c, res))
+			builder.mbuf1 = append(builder.mbuf1, res...)
+			return
+		}
+
+		// v1 and v2 are constants, we multiply big.Int values and return resulting constant
+		if v1Constant && v2Constant {
+			builder.cs.Mul(&n1, &n2)
+			builder.mbuf1 = append(builder.mbuf1, expr.NewTerm(0, n1))
+			return
+		}
 
+		if v1Constant {
+			builder.mbuf1 = append(builder.mbuf1, builder.toVariable(c)...)
+			builder.mulConstant(builder.mbuf1, n1, true)
+			return
+		}
+		builder.mbuf1 = append(builder.mbuf1, builder.toVariable(b)...)
+		builder.mulConstant(builder.mbuf1, n2, true)
+	}
+	mulBC()
+
+	_a := builder.toVariable(a)
+	// copy _a in buffer, use _a as result; so if _a was already a linear expression and
+	// results fits, _a is mutated without performing a new memalloc
+	builder.mbuf2 = builder.mbuf2[:0]
+	builder.add([]expr.LinearExpression{_a, builder.mbuf1}, false, 0, &builder.mbuf2)
+	_a = _a[:0]
+	if len(builder.mbuf2) <= cap(_a) {
+		// it fits, no mem alloc
+		_a = append(_a, builder.mbuf2...)
+	} else {
+		// allocate a expression linear with extended capacity
+		_a = make(expr.LinearExpression, len(builder.mbuf2), len(builder.mbuf2)*3)
+		copy(_a, builder.mbuf2)
+	}
+	return _a
 }
 
 // Sub returns res = i1 - i2
 func (builder *builder) Sub(i1, i2 frontend.Variable, in ...frontend.Variable) frontend.Variable {
 	// extract frontend.Variables from input
 	vars, s := builder.toVariables(append([]frontend.Variable{i1, i2}, in...)...)
-	return builder.add(vars, true, s)
+	return builder.add(vars, true, s, nil)
 }
 
 // returns res = Σ(vars) or res = vars[0] - Σ(vars[1:]) if sub == true.
-func (builder *builder) add(vars []expr.LinearExpression, sub bool, capacity int) frontend.Variable {
+func (builder *builder) add(vars []expr.LinearExpression, sub bool, capacity int, res *expr.LinearExpression) frontend.Variable {
 	// we want to merge all terms from input linear expressions
 	// if they are duplicate, we reduce; that is, if multiple terms in different vars have the
 	// same variable id.
@@ -68,7 +118,10 @@ func (builder *builder) add(vars []expr.LinearExpression, sub bool, capacity int
 	}
 	builder.heap.heapify()
 
-	res := make(expr.LinearExpression, 0, capacity)
+	if res == nil {
+		t := make(expr.LinearExpression, 0, capacity)
+		res = &t
+	}
 	curr := -1
 
 	// process all the terms from all the inputs, in sorted order
@@ -87,37 +140,42 @@ func (builder *builder) add(vars []expr.LinearExpression, sub bool, capacity int
 		if t.Coeff.IsZero() {
 			continue // is this really needed?
 		}
-		if curr != -1 && t.VID == res[curr].VID {
+		if curr != -1 && t.VID == (*res)[curr].VID {
 			// accumulate, it's the same variable ID
 			if sub && lID != 0 {
-				builder.cs.Sub(&res[curr].Coeff, &t.Coeff)
+				builder.cs.Sub(&(*res)[curr].Coeff, &t.Coeff)
 			} else {
-				builder.cs.Add(&res[curr].Coeff, &t.Coeff)
+				builder.cs.Add(&(*res)[curr].Coeff, &t.Coeff)
 			}
-			if res[curr].Coeff.IsZero() {
+			if (*res)[curr].Coeff.IsZero() {
 				// remove self.
-				res = res[:curr]
+				(*res) = (*res)[:curr]
 				curr--
 			}
 		} else {
 			// append, it's a new variable ID
-			res = append(res, *t)
+			(*res) = append((*res), *t)
 			curr++
 			if sub && lID != 0 {
-				builder.cs.Neg(&res[curr].Coeff)
+				builder.cs.Neg(&(*res)[curr].Coeff)
 			}
 		}
 	}
 
-	if len(res) == 0 {
+	if len((*res)) == 0 {
 		// keep the linear expression valid (assertIsSet)
-		res = expr.NewLinearExpression(0, constraint.Coeff{})
+		(*res) = append((*res), expr.NewTerm(0, constraint.Coeff{}))
 	}
 	// if the linear expression LE is too long then record an equality
 	// constraint LE * 1 = t and return short linear expression instead.
-	res = builder.compress(res)
+	compressed := builder.compress((*res))
+	if len(compressed) != len(*res) {
+		// we compressed, but don't want to override buffer
+		*res = (*res)[:0]
+		*res = append(*res, compressed...)
+	}
 
-	return res
+	return *res
 }
 
 // Neg returns -i
@@ -151,7 +209,6 @@ func (builder *builder) Mul(i1, i2 frontend.Variable, in ...frontend.Variable) f
 		// v1 and v2 are constants, we multiply big.Int values and return resulting constant
 		if v1Constant && v2Constant {
 			builder.cs.Mul(&n1, &n2)
-			// n1.Mul(n1, n2).Mod(n1, builder.q)
 			return expr.NewLinearExpression(0, n1)
 		}
 

frontend/cs/r1cs/api_assertions.go

@@ -83,14 +83,12 @@ func (builder *builder) AssertIsBoolean(i1 frontend.Variable) {
 func (builder *builder) AssertIsLessOrEqual(_v frontend.Variable, bound frontend.Variable) {
 	v := builder.toVariable(_v)
 
-	switch b := bound.(type) {
-	case expr.LinearExpression:
+	if b, ok := bound.(expr.LinearExpression); ok {
 		assertIsSet(b)
 		builder.mustBeLessOrEqVar(v, b)
-	default:
-		builder.mustBeLessOrEqCst(v, utils.FromInterface(b))
+	} else {
+		builder.mustBeLessOrEqCst(v, utils.FromInterface(bound))
 	}
-
 }
 
 func (builder *builder) mustBeLessOrEqVar(a, bound expr.LinearExpression) {

frontend/cs/r1cs/builder.go

@@ -59,15 +59,25 @@ type builder struct {
 	q    *big.Int
 	tOne constraint.Coeff
 	heap minHeap // helps merge k sorted linear expressions
+
+	// buffers used to do in place api.MAC
+	mbuf1 expr.LinearExpression
+	mbuf2 expr.LinearExpression
 }
 
 // initialCapacity has quite some impact on frontend performance, especially on large circuits size
 // we may want to add build tags to tune that
 func newBuilder(field *big.Int, config frontend.CompileConfig) *builder {
+	macCapacity := 100
+	if config.CompressThreshold != 0 {
+		macCapacity = config.CompressThreshold
+	}
 	builder := builder{
 		mtBooleans: make(map[uint64][]expr.LinearExpression, config.Capacity/10),
 		config:     config,
 		heap:       make(minHeap, 0, 100),
+		mbuf1:      make(expr.LinearExpression, 0, macCapacity),
+		mbuf2:      make(expr.LinearExpression, 0, macCapacity),
 	}
 
 	// by default the circuit is given a public wire equal to 1
@@ -306,6 +316,9 @@ func (builder *builder) toVariable(input interface{}) expr.LinearExpression {
 		// this is already a "kwown" variable
 		assertIsSet(t)
 		return t
+	case *expr.LinearExpression:
+		assertIsSet(*t)
+		return *t
 	case constraint.Coeff:
 		return expr.NewLinearExpression(0, t)
 	case *constraint.Coeff:
@@ -352,14 +365,11 @@ func (builder *builder) NewHint(f hint.Function, nbOutputs int, inputs ...fronte
 	// TODO @gbotrel hint input pass
 	// ensure inputs are set and pack them in a []uint64
 	for i, in := range inputs {
-		switch t := in.(type) {
-		case expr.LinearExpression:
+		if t, ok := in.(expr.LinearExpression); ok {
 			assertIsSet(t)
 			hintInputs[i] = builder.getLinearExpression(t)
-		default:
-			// make a term
-			// c := utils.FromInterface(t)
-			c := builder.cs.FromInterface(t)
+		} else {
+			c := builder.cs.FromInterface(in)
 			term := builder.cs.MakeTerm(&c, 0)
 			term.MarkConstant()
 			hintInputs[i] = constraint.LinearExpression{term}

frontend/cs/scs/api.go

@@ -52,6 +52,12 @@ func (builder *scs) Add(i1, i2 frontend.Variable, in ...frontend.Variable) front
 
 }
 
+func (builder *scs) MulAcc(a, b, c frontend.Variable) frontend.Variable {
+	// TODO can we do better here to limit allocations?
+	// technically we could do that in one PlonK constraint (against 2 for separate Add & Mul)
+	return builder.Add(a, builder.Mul(b, c))
+}
+
 // neg returns -in
 func (builder *scs) neg(in []frontend.Variable) []frontend.Variable {
 

std/math/emulated/field_assert.go

@@ -67,8 +67,9 @@ func rsh(api frontend.API, v frontend.Variable, startDigit, endDigit int) fronte
 	c.Lsh(c, uint(startDigit))
 
 	for i := 0; i < len(bits); i++ {
-		Σbi = api.Add(Σbi, api.Mul(bits[i], c))
-		ΣbiRShift = api.Add(ΣbiRShift, api.Mul(bits[i], cRShift))
+		Σbi = api.MulAcc(Σbi, bits[i], c)
+		ΣbiRShift = api.MulAcc(ΣbiRShift, bits[i], cRShift)
+
 		c.Lsh(c, 1)
 		cRShift.Lsh(cRShift, 1)
 		api.AssertIsBoolean(bits[i])

std/math/emulated/wrapped_api.go

@@ -76,6 +76,11 @@ func (w *FieldAPI[T]) Add(i1 frontend.Variable, i2 frontend.Variable, in ...fron
 	return res
 }
 
+func (w *FieldAPI[T]) MulAcc(a, b, c frontend.Variable) frontend.Variable {
+	// TODO can we do better here to limit allocations?
+	return w.Add(a, w.Mul(b, c))
+}
+
 func (w *FieldAPI[T]) Neg(i1 frontend.Variable) frontend.Variable {
 	el := w.varToElement(i1)
 	return w.f.Neg(el)

test/engine.go

@@ -30,6 +30,7 @@ import (
 	"github.com/consensys/gnark/logger"
 
 	"github.com/consensys/gnark-crypto/ecc"
+	"github.com/consensys/gnark-crypto/field"
 	"github.com/consensys/gnark/backend"
 	"github.com/consensys/gnark/backend/hint"
 	"github.com/consensys/gnark/frontend"
@@ -156,6 +157,17 @@ func (e *engine) Add(i1, i2 frontend.Variable, in ...frontend.Variable) frontend
 	return res
 }
 
+func (e *engine) MulAcc(a, b, c frontend.Variable) frontend.Variable {
+	bc := field.BigIntPool.Get()
+	bc.Mul(e.toBigInt(b), e.toBigInt(c))
+
+	_a := e.toBigInt(a)
+	_a.Add(_a, bc).Mod(_a, e.modulus())
+
+	field.BigIntPool.Put(bc)
+	return _a
+}
+
 func (e *engine) Sub(i1, i2 frontend.Variable, in ...frontend.Variable) frontend.Variable {
 	cptSub++
 	res := new(big.Int)