1package prover
  2
  3import (
  4	"math"
  5	"math/big"
  6
  7	"github.com/iden3/go-circom-prover-verifier/types"
  8	"github.com/iden3/go-iden3-crypto/ff"
  9)
 10
 11type rootsT struct {
 12	roots [][]*ff.Element
 13	w     []*ff.Element
 14}
 15
 16func newRootsT() rootsT {
 17	var roots rootsT
 18
 19	rem := new(big.Int).Sub(types.R, big.NewInt(1))
 20	s := 0
 21	for rem.Bit(0) == 0 { // rem.Bit==0 when even
 22		s++
 23		rem = new(big.Int).Rsh(rem, 1)
 24	}
 25	roots.w = make([]*ff.Element, s+1)
 26	roots.w[s] = ff.NewElement().SetBigInt(fExp(big.NewInt(5), rem))
 27
 28	n := s - 1
 29	for n >= 0 {
 30		roots.w[n] = ff.NewElement().Mul(roots.w[n+1], roots.w[n+1])
 31		n--
 32	}
 33	roots.roots = make([][]*ff.Element, 50) // TODO WIP
 34
 35	roots.setRoots(15)
 36	return roots
 37}
 38
 39func (roots rootsT) setRoots(n int) {
 40	for i := n; i >= 0 && nil == roots.roots[i]; i-- { // TODO tmp i<=len(r)
 41		r := ff.NewElement().SetBigInt(big.NewInt(1))
 42		nroots := 1 << i
 43		var rootsi []*ff.Element
 44		for j := 0; j < nroots; j++ {
 45			rootsi = append(rootsi, r)
 46			r = ff.NewElement().Mul(r, roots.w[i])
 47		}
 48		roots.roots[i] = rootsi
 49	}
 50}
 51
 52func fftroots(roots rootsT, pall []*ff.Element, bits, offset, step int) []*ff.Element {
 53	n := 1 << bits
 54	if n == 1 {
 55		return []*ff.Element{pall[offset]}
 56	} else if n == 2 {
 57		return []*ff.Element{
 58			ff.NewElement().Add(pall[offset], pall[offset+step]), // TODO tmp
 59			ff.NewElement().Sub(pall[offset], pall[offset+step]),
 60		}
 61	}
 62
 63	ndiv2 := n >> 1
 64	p1 := fftroots(roots, pall, bits-1, offset, step*2)
 65	p2 := fftroots(roots, pall, bits-1, offset+step, step*2)
 66
 67	out := make([]*ff.Element, n)
 68	for i := 0; i < ndiv2; i++ {
 69		out[i] = ff.NewElement().Add(p1[i], ff.NewElement().Mul(roots.roots[bits][i], p2[i]))
 70		out[i+ndiv2] = ff.NewElement().Sub(p1[i], ff.NewElement().Mul(roots.roots[bits][i], p2[i]))
 71	}
 72	return out
 73}
 74
 75func fft(p []*ff.Element) []*ff.Element {
 76	if len(p) <= 1 {
 77		return p
 78	}
 79	bits := math.Log2(float64(len(p)-1)) + 1
 80	roots := newRootsT()
 81	roots.setRoots(int(bits))
 82	m := 1 << int(bits)
 83	ep := extend(p, m)
 84	res := fftroots(roots, ep, int(bits), 0, 1)
 85	return res
 86}
 87
 88func ifft(p []*ff.Element) []*ff.Element {
 89	res := fft(p)
 90	bits := math.Log2(float64(len(p)-1)) + 1
 91	m := 1 << int(bits)
 92
 93	twoinvm := ff.NewElement().SetBigInt(fInv(fMul(big.NewInt(1), big.NewInt(int64(m)))))
 94
 95	var resn []*ff.Element
 96	for i := 0; i < m; i++ {
 97		resn = append(resn, ff.NewElement().Mul(res[(m-i)%m], twoinvm))
 98	}
 99
100	return resn
101}
102
103func extend(p []*ff.Element, e int) []*ff.Element {
104	if e == len(p) {
105		return p
106	}
107	z := arrayOfZeroesE(e - len(p))
108	return append(p, z...)
109}