Modify agg_kzg to fold polynomials in parallel

Currently the code will do it sequentially. To change it to parallel, we will need to allocate a polynomial for each thread since we do not want to modify the polynomials passed in.

We can offer both APIs for the user, where the first will simply do a Copy beforehand

Code snippets

// func BatchOpenSinglePoint(domain *kzg.Domain, polynomials []kzg.Polynomial, commitKey *kzg.CommitKey) (*BatchOpeningProof, error) {
//  numPolynomials := len(polynomials)

//  _polys := make([]kzg.Polynomial, numPolynomials)
//  for i := 0; i < numPolynomials; i++ {
//      polyLen := len(polynomials[i])
//      _polys[i] = make([]fr.Element, polyLen)
//      for j := 0; j < polyLen; j++ {
//          _polys[i][j] = polynomials[i][j]
//      }
//  }

//  return BatchOpenSinglePointNoCopy(domain, _polys, commitKey)
// }

func foldPolynomials(polynomials []kzg.Polynomial, challenges []fr.Element) (kzg.Polynomial, error) {
    numPolynomials := len(polynomials)
    numChallenges := len(challenges)

    if numPolynomials != numChallenges {
        return nil, errors.New("number of polynomials is different to the number of challenges provided")
    }

    // WORKS
    // result := make(kzg.Polynomial, len(polynomials[0]))
    // copy(result, polynomials[0])

    // for i := 1; i < len(polynomials); i++ {
    //  partialSum := vectorScale(polynomials[i], challenges[i])
    //  vectorAddition(result, result, partialSum)
    // }
    // return result, nil
    // WORKS

    partialSums := make([]kzg.Polynomial, len(polynomials))

    var wg sync.WaitGroup
    wg.Add(len(polynomials) - 1)
    for i := 1; i < len(polynomials); i++ {
        go func(at int) {
            partialSums[at] = vectorScale(polynomials[at], challenges[at])
            wg.Done()
        }(i)
    }
    wg.Wait()

    result := make(kzg.Polynomial, len(polynomials[0]))
    copy(result, polynomials[0])
    for i := 1; i < len(polynomials); i++ {
        vectorAddition(result, result, partialSums[i])
    }
    return result, nil
}

func vectorScale(vector []fr.Element, scale fr.Element) []fr.Element {
    vectorSize := len(vector)

    result := make([]fr.Element, vectorSize)
    for i := 0; i < vectorSize; i++ {
        result[i].Mul(&vector[i], &scale)
    }
    return result
}

func vectorAddition(result []fr.Element, lhs []fr.Element, rhs []fr.Element) {
    lhsSize := len(lhs)
    rhsSize := len(rhs)
    if lhsSize != rhsSize {
        fmt.Printf("%d %d", lhsSize, rhsSize)
        panic("both slices must be the same length")
    }

    for i := 0; i < lhsSize; i++ {
        result[i].Add(&lhs[i], &rhs[i])
    }
}

crate-crypto / go-eth-kzg

Modify agg_kzg to fold polynomials in parallel #1