feat: better emulation for small fields in large fields (#1682)

Co-authored-by: Ivo Kubjas <ivo.kubjas@consensys.net>

Author: gbotrel

std/math/emulated/field.go

@@ -51,6 +51,13 @@ type Field[T FieldParams] struct {
 	checker          frontend.Rangechecker
 
 	deferredChecks []deferredChecker
+
+	// smallFieldMode indicates that the emulated field is small enough that
+	// products fit in the native field and we can use scalar batched verification
+	// instead of polynomial identity testing. This provides significant constraint
+	// reduction for small field emulation (e.g., KoalaBear on BLS12-377).
+	smallFieldMode     bool
+	smallFieldModeOnce sync.Once
 }
 
 type ctxKey[T FieldParams] struct{}
@@ -310,3 +317,46 @@ func sum[T constraints.Ordered](a ...T) T {
 	}
 	return m
 }
+
+// useSmallFieldOptimization returns true if we can use the small field
+// optimization for multiplication. The optimization is possible when:
+//   - NbLimbs == 1 (emulated field fits in a single native limb)
+//   - 2 * modBits + margin < nativeBits - 2 (products fit with margin for batching)
+//
+// When these conditions are met, we can use scalar batched verification instead
+// of polynomial identity testing, which significantly reduces constraint counts.
+func (f *Field[T]) useSmallFieldOptimization() bool {
+	f.smallFieldModeOnce.Do(func() {
+		// Small field optimization only works when NbLimbs == 1
+		if f.fParams.NbLimbs() != 1 {
+			f.smallFieldMode = false
+			return
+		}
+
+		// Small field optimization doesn't work when we're already using extension field
+		// for multiplication checks (native field is small)
+		if f.extensionApi != nil {
+			f.smallFieldMode = false
+			return
+		}
+
+		// Check that products fit in the native field with margin for batching.
+		// We need: 2 * modBits + batchingMargin < nativeBits - 2
+		// The margin accounts for:
+		// - γ^i scaling factors in the batched sum
+		// - Multiple terms being summed together
+		// We use 32 bits margin which allows for batching millions of operations.
+		modBits := uint(f.fParams.Modulus().BitLen())
+		nativeBits := uint(f.api.Compiler().FieldBitLen())
+		const batchingMargin = 32
+
+		f.smallFieldMode = 2*modBits+batchingMargin < nativeBits-2
+		if f.smallFieldMode {
+			f.log.Debug().
+				Uint("modBits", modBits).
+				Uint("nativeBits", nativeBits).
+				Msg("using small field optimization for emulated multiplication")
+		}
+	})
+	return f.smallFieldMode
+}

std/math/emulated/field_mul.go

@@ -233,6 +233,16 @@ func (f *Field[T]) mulMod(a, b *Element[T], _ uint, p *Element[T]) *Element[T] {
 	}
 	f.enforceWidthConditional(a)
 	f.enforceWidthConditional(b)
+
+	// Use small field optimization if available and no custom modulus
+	if p == nil && f.useSmallFieldOptimization() {
+		// For small field mode, ensure elements are single-limb
+		// If they have more limbs due to witness initialization, convert them
+		aLimb := f.toSingleLimbElement(a)
+		bLimb := f.toSingleLimbElement(b)
+		return f.smallMulMod(aLimb, bLimb)
+	}
+
 	f.enforceWidthConditional(p)
 	k, r, c, err := f.callMulHint(a, b, true, p)
 	if err != nil {
@@ -257,9 +267,18 @@ func (f *Field[T]) checkZero(a *Element[T], p *Element[T]) {
 	if a.isStrictZero() {
 		return
 	}
+
+	f.enforceWidthConditional(a)
+
+	// Use small field optimization if available and no custom modulus
+	if p == nil && f.useSmallFieldOptimization() {
+		aLimb := f.toSingleLimbElement(a)
+		f.smallCheckZero(aLimb)
+		return
+	}
+
 	// the method works similarly to mulMod, but we know that we are multiplying
 	// by one and expected result should be zero.
-	f.enforceWidthConditional(a)
 	f.enforceWidthConditional(p)
 	b := f.One()
 	k, r, c, err := f.callMulHint(a, b, false, p)