perf: fold constant byte operations in std/math/uints

std/math/uints/bytes.go

@@ -166,16 +166,45 @@ func (bf *Bytes) twoArgFn(tbl *logderivprecomp.Precomputed, a ...U8) U8 {
 	if len(a) == 1 {
 		return a[0]
 	}
-	ret := tbl.Query(a[0].Val, a[1].Val)[0]
+	ret := bf.queryOrFold(tbl, a[0].Val, a[1].Val)
 	for i := 2; i < len(a); i++ {
-		ret = tbl.Query(ret, a[i].Val)[0]
+		ret = bf.queryOrFold(tbl, ret, a[i].Val)
 	}
-	// because the response comes from the lookup table, then (assuming that the
-	// function which built the table is correct) we can assume that the value
-	// is in range. Thus we set the internal flag to true.
+	// because the response comes from the lookup table (or from constant
+	// folding of width-checked operands), then (assuming that the function
+	// which built the table is correct) we can assume that the value is in
+	// range. Thus we set the internal flag to true.
 	return bf.packInternal(ret)
 }
 
+// queryOrFold evaluates the byte operation implemented by tbl natively when
+// both operands are compile-time constants, avoiding the lookup query and its
+// constraints. The operands are already width-checked ([Bytes.enforceWidth]
+// for the inputs, previous lookups or folds for the intermediates), but we
+// still guard on the width so that an out-of-range constant falls back to the
+// lookup path instead of being truncated silently.
+func (bf *Bytes) queryOrFold(tbl *logderivprecomp.Precomputed, x, y frontend.Variable) frontend.Variable {
+	cx, xIsConst := bf.api.ConstantValue(x)
+	if !xIsConst || cx.BitLen() > 8 {
+		return tbl.Query(x, y)[0]
+	}
+	cy, yIsConst := bf.api.ConstantValue(y)
+	if !yIsConst || cy.BitLen() > 8 {
+		return tbl.Query(x, y)[0]
+	}
+	xb, yb := uint8(cx.Uint64()), uint8(cy.Uint64())
+	switch tbl {
+	case bf.xorT:
+		return xb ^ yb
+	case bf.andT:
+		return xb & yb
+	case bf.orT:
+		return xb | yb
+	default:
+		return tbl.Query(x, y)[0]
+	}
+}
+
 func (bf *Bytes) Not(a U8) U8 {
 	ret := bf.xorT.Query(a.Val, bf.allOne.Val)
 	// the response comes from the lookup table, thus we can assume that the

std/math/uints/bytes_test.go

@@ -0,0 +1,101 @@
+package uints
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/consensys/gnark-crypto/ecc"
+	"github.com/consensys/gnark/frontend"
+	"github.com/consensys/gnark/frontend/cs/r1cs"
+	"github.com/consensys/gnark/test"
+)
+
+// constFoldMixedCircuit chains compile-time constants with a witness operand
+// so that both the folded and the lookup paths are exercised in one query
+// chain: the leading constant pair folds, the witness operand forces a lookup
+// and the trailing constant queries against a non-constant intermediate.
+type constFoldMixedCircuit struct {
+	In       U8
+	Expected U8
+	mode     int
+}
+
+func (c *constFoldMixedCircuit) Define(api frontend.API) error {
+	uapi, err := NewBytes(api)
+	if err != nil {
+		return fmt.Errorf("NewBytes: %w", err)
+	}
+	var res U8
+	switch c.mode {
+	case 0:
+		res = uapi.And(NewU8(0x3c), NewU8(0x5a), c.In, NewU8(0xf0))
+	case 1:
+		res = uapi.Or(NewU8(0x3c), NewU8(0x5a), c.In, NewU8(0xf0))
+	case 2:
+		res = uapi.Xor(NewU8(0x3c), NewU8(0x5a), c.In, NewU8(0xf0))
+	}
+	uapi.AssertIsEqual(res, c.Expected)
+	return nil
+}
+
+func TestByteOpConstantFoldMixed(t *testing.T) {
+	assert := test.NewAssert(t)
+	in := uint8(0xa7)
+	assert.Run(func(assert *test.Assert) {
+		expected := 0x3c & 0x5a & in & 0xf0
+		assert.CheckCircuit(&constFoldMixedCircuit{mode: 0}, test.WithValidAssignment(&constFoldMixedCircuit{In: NewU8(in), Expected: NewU8(expected)}))
+		assert.CheckCircuit(&constFoldMixedCircuit{mode: 0}, test.WithInvalidAssignment(&constFoldMixedCircuit{In: NewU8(in), Expected: NewU8(expected ^ 1)}))
+	}, "and")
+	assert.Run(func(assert *test.Assert) {
+		expected := 0x3c | 0x5a | in | 0xf0
+		assert.CheckCircuit(&constFoldMixedCircuit{mode: 1}, test.WithValidAssignment(&constFoldMixedCircuit{In: NewU8(in), Expected: NewU8(expected)}))
+		assert.CheckCircuit(&constFoldMixedCircuit{mode: 1}, test.WithInvalidAssignment(&constFoldMixedCircuit{In: NewU8(in), Expected: NewU8(expected ^ 1)}))
+	}, "or")
+	assert.Run(func(assert *test.Assert) {
+		expected := 0x3c ^ 0x5a ^ in ^ 0xf0
+		assert.CheckCircuit(&constFoldMixedCircuit{mode: 2}, test.WithValidAssignment(&constFoldMixedCircuit{In: NewU8(in), Expected: NewU8(expected)}))
+		assert.CheckCircuit(&constFoldMixedCircuit{mode: 2}, test.WithInvalidAssignment(&constFoldMixedCircuit{In: NewU8(in), Expected: NewU8(expected ^ 1)}))
+	}, "xor")
+}
+
+// constFoldOnlyCircuit performs byte operations exclusively on compile-time
+// constants; with folding it must not emit any lookup query.
+type constFoldOnlyCircuit struct {
+	Dummy frontend.Variable
+}
+
+func (c *constFoldOnlyCircuit) Define(api frontend.API) error {
+	uapi, err := NewBytes(api)
+	if err != nil {
+		return fmt.Errorf("NewBytes: %w", err)
+	}
+	and := uapi.And(NewU8(0x3c), NewU8(0x5a), NewU8(0xf0))
+	or := uapi.Or(NewU8(0x3c), NewU8(0x5a), NewU8(0x0f))
+	xor := uapi.Xor(NewU8(0x3c), NewU8(0x5a), NewU8(0xff))
+	uapi.AssertIsEqual(and, NewU8(0x3c&0x5a&0xf0))
+	uapi.AssertIsEqual(or, NewU8(0x3c|0x5a|0x0f))
+	uapi.AssertIsEqual(xor, NewU8(0x3c^0x5a^0xff))
+	return nil
+}
+
+// byteOpBaselineCircuit instantiates the same tables but performs no byte
+// operations. Constant-only operations must not add constraints on top of it.
+type byteOpBaselineCircuit struct {
+	Dummy frontend.Variable
+}
+
+func (c *byteOpBaselineCircuit) Define(api frontend.API) error {
+	if _, err := NewBytes(api); err != nil {
+		return fmt.Errorf("NewBytes: %w", err)
+	}
+	return nil
+}
+
+func TestByteOpConstantFoldAddsNoConstraints(t *testing.T) {
+	assert := test.NewAssert(t)
+	folded, err := frontend.Compile(ecc.BN254.ScalarField(), r1cs.NewBuilder, &constFoldOnlyCircuit{})
+	assert.NoError(err)
+	baseline, err := frontend.Compile(ecc.BN254.ScalarField(), r1cs.NewBuilder, &byteOpBaselineCircuit{})
+	assert.NoError(err)
+	assert.Equal(baseline.GetNbConstraints(), folded.GetNbConstraints(), "constant-only byte operations added constraints")
+}