feat: make constraint.Element generic interface (#1463)

Author: ivokub

backend/witness/vector.go

@@ -13,7 +13,9 @@ import (
 	fr_bn254 "github.com/consensys/gnark-crypto/ecc/bn254/fr"
 	fr_bw6633 "github.com/consensys/gnark-crypto/ecc/bw6-633/fr"
 	fr_bw6761 "github.com/consensys/gnark-crypto/ecc/bw6-761/fr"
-	"github.com/consensys/gnark/internal/tinyfield"
+	"github.com/consensys/gnark-crypto/field/babybear"
+	"github.com/consensys/gnark-crypto/field/koalabear"
+	"github.com/consensys/gnark/internal/smallfields/tinyfield"
 	"github.com/consensys/gnark/internal/utils"
 )
 
@@ -37,9 +39,14 @@ func newVector(field *big.Int, size int) (any, error) {
 	default:
 		if field.Cmp(tinyfield.Modulus()) == 0 {
 			return make(tinyfield.Vector, size), nil
-		} else {
-			return nil, errors.New("unsupported modulus")
 		}
+		if field.Cmp(babybear.Modulus()) == 0 {
+			return make(babybear.Vector, size), nil
+		}
+		if field.Cmp(koalabear.Modulus()) == 0 {
+			return make(koalabear.Vector, size), nil
+		}
+		return nil, errors.New("unsupported modulus")
 	}
 }
 
@@ -77,6 +84,14 @@ func newFrom(from any, n int) (any, error) {
 		a := make(tinyfield.Vector, n)
 		copy(a, wt)
 		return a, nil
+	case babybear.Vector:
+		a := make(babybear.Vector, n)
+		copy(a, wt)
+		return a, nil
+	case koalabear.Vector:
+		a := make(koalabear.Vector, n)
+		copy(a, wt)
+		return a, nil
 	default:
 		return nil, errors.New("unsupported modulus")
 	}
@@ -100,6 +115,10 @@ func leafType(v any) reflect.Type {
 		return reflect.TypeOf(fr_bw6633.Element{})
 	case tinyfield.Vector:
 		return reflect.TypeOf(tinyfield.Element{})
+	case babybear.Vector:
+		return reflect.TypeOf(babybear.Element{})
+	case koalabear.Vector:
+		return reflect.TypeOf(koalabear.Element{})
 	default:
 		panic("invalid input")
 	}
@@ -155,6 +174,18 @@ func set(v any, index int, value any) error {
 		}
 		_, err := pv[index].SetInterface(value)
 		return err
+	case babybear.Vector:
+		if index >= len(pv) {
+			return errors.New("out of bounds")
+		}
+		_, err := pv[index].SetInterface(value)
+		return err
+	case koalabear.Vector:
+		if index >= len(pv) {
+			return errors.New("out of bounds")
+		}
+		_, err := pv[index].SetInterface(value)
+		return err
 	default:
 		panic("invalid input")
 	}
@@ -219,6 +250,20 @@ func iterate(v any) chan any {
 			}
 			close(chValues)
 		}()
+	case babybear.Vector:
+		go func() {
+			for i := 0; i < len(pv); i++ {
+				chValues <- &(pv)[i]
+			}
+			close(chValues)
+		}()
+	case koalabear.Vector:
+		go func() {
+			for i := 0; i < len(pv); i++ {
+				chValues <- &(pv)[i]
+			}
+			close(chValues)
+		}()
 	default:
 		panic("invalid input")
 	}
@@ -243,6 +288,10 @@ func resize(v any, n int) any {
 		return make(fr_bw6633.Vector, n)
 	case tinyfield.Vector:
 		return make(tinyfield.Vector, n)
+	case babybear.Vector:
+		return make(babybear.Vector, n)
+	case koalabear.Vector:
+		return make(koalabear.Vector, n)
 	default:
 		panic("invalid input")
 	}

backend/witness/witness.go

@@ -54,9 +54,11 @@ import (
 	fr_bn254 "github.com/consensys/gnark-crypto/ecc/bn254/fr"
 	fr_bw6633 "github.com/consensys/gnark-crypto/ecc/bw6-633/fr"
 	fr_bw6761 "github.com/consensys/gnark-crypto/ecc/bw6-761/fr"
+	"github.com/consensys/gnark-crypto/field/babybear"
+	"github.com/consensys/gnark-crypto/field/koalabear"
 	"github.com/consensys/gnark/debug"
 	"github.com/consensys/gnark/frontend/schema"
-	"github.com/consensys/gnark/internal/tinyfield"
+	"github.com/consensys/gnark/internal/smallfields/tinyfield"
 )
 
 var ErrInvalidWitness = errors.New("invalid witness")
@@ -189,6 +191,10 @@ func (w *witness) WriteTo(wr io.Writer) (n int64, err error) {
 		m, err = t.WriteTo(wr)
 	case tinyfield.Vector:
 		m, err = t.WriteTo(wr)
+	case babybear.Vector:
+		m, err = t.WriteTo(wr)
+	case koalabear.Vector:
+		m, err = t.WriteTo(wr)
 	default:
 		panic("invalid input")
 	}
@@ -235,6 +241,12 @@ func (w *witness) ReadFrom(r io.Reader) (n int64, err error) {
 	case tinyfield.Vector:
 		m, err = t.ReadFrom(r)
 		w.vector = t
+	case babybear.Vector:
+		m, err = t.ReadFrom(r)
+		w.vector = t
+	case koalabear.Vector:
+		m, err = t.ReadFrom(r)
+		w.vector = t
 	default:
 		panic("invalid input")
 	}