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