github.com/remyoudompheng/bigfft@v0.0.0-20230129092748-24d4a6f8daec/fft.go (about)

     1  // Package bigfft implements multiplication of big.Int using FFT.
     2  //
     3  // The implementation is based on the Schönhage-Strassen method
     4  // using integer FFT modulo 2^n+1.
     5  package bigfft
     6  
     7  import (
     8  	"math/big"
     9  	"unsafe"
    10  )
    11  
    12  const _W = int(unsafe.Sizeof(big.Word(0)) * 8)
    13  
    14  type nat []big.Word
    15  
    16  func (n nat) String() string {
    17  	v := new(big.Int)
    18  	v.SetBits(n)
    19  	return v.String()
    20  }
    21  
    22  // fftThreshold is the size (in words) above which FFT is used over
    23  // Karatsuba from math/big.
    24  //
    25  // TestCalibrate seems to indicate a threshold of 60kbits on 32-bit
    26  // arches and 110kbits on 64-bit arches.
    27  var fftThreshold = 1800
    28  
    29  // Mul computes the product x*y and returns z.
    30  // It can be used instead of the Mul method of
    31  // *big.Int from math/big package.
    32  func Mul(x, y *big.Int) *big.Int {
    33  	xwords := len(x.Bits())
    34  	ywords := len(y.Bits())
    35  	if xwords > fftThreshold && ywords > fftThreshold {
    36  		return mulFFT(x, y)
    37  	}
    38  	return new(big.Int).Mul(x, y)
    39  }
    40  
    41  func mulFFT(x, y *big.Int) *big.Int {
    42  	var xb, yb nat = x.Bits(), y.Bits()
    43  	zb := fftmul(xb, yb)
    44  	z := new(big.Int)
    45  	z.SetBits(zb)
    46  	if x.Sign()*y.Sign() < 0 {
    47  		z.Neg(z)
    48  	}
    49  	return z
    50  }
    51  
    52  // A FFT size of K=1<<k is adequate when K is about 2*sqrt(N) where
    53  // N = x.Bitlen() + y.Bitlen().
    54  
    55  func fftmul(x, y nat) nat {
    56  	k, m := fftSize(x, y)
    57  	xp := polyFromNat(x, k, m)
    58  	yp := polyFromNat(y, k, m)
    59  	rp := xp.Mul(&yp)
    60  	return rp.Int()
    61  }
    62  
    63  // fftSizeThreshold[i] is the maximal size (in bits) where we should use
    64  // fft size i.
    65  var fftSizeThreshold = [...]int64{0, 0, 0,
    66  	4 << 10, 8 << 10, 16 << 10, // 5 
    67  	32 << 10, 64 << 10, 1 << 18, 1 << 20, 3 << 20, // 10
    68  	8 << 20, 30 << 20, 100 << 20, 300 << 20, 600 << 20,
    69  }
    70  
    71  // returns the FFT length k, m the number of words per chunk
    72  // such that m << k is larger than the number of words
    73  // in x*y.
    74  func fftSize(x, y nat) (k uint, m int) {
    75  	words := len(x) + len(y)
    76  	bits := int64(words) * int64(_W)
    77  	k = uint(len(fftSizeThreshold))
    78  	for i := range fftSizeThreshold {
    79  		if fftSizeThreshold[i] > bits {
    80  			k = uint(i)
    81  			break
    82  		}
    83  	}
    84  	// The 1<<k chunks of m words must have N bits so that
    85  	// 2^N-1 is larger than x*y. That is, m<<k > words
    86  	m = words>>k + 1
    87  	return
    88  }
    89  
    90  // valueSize returns the length (in words) to use for polynomial
    91  // coefficients, to compute a correct product of polynomials P*Q
    92  // where deg(P*Q) < K (== 1<<k) and where coefficients of P and Q are
    93  // less than b^m (== 1 << (m*_W)).
    94  // The chosen length (in bits) must be a multiple of 1 << (k-extra).
    95  func valueSize(k uint, m int, extra uint) int {
    96  	// The coefficients of P*Q are less than b^(2m)*K
    97  	// so we need W * valueSize >= 2*m*W+K
    98  	n := 2*m*_W + int(k) // necessary bits
    99  	K := 1 << (k - extra)
   100  	if K < _W {
   101  		K = _W
   102  	}
   103  	n = ((n / K) + 1) * K // round to a multiple of K
   104  	return n / _W
   105  }
   106  
   107  // poly represents an integer via a polynomial in Z[x]/(x^K+1)
   108  // where K is the FFT length and b^m is the computation basis 1<<(m*_W).
   109  // If P = a[0] + a[1] x + ... a[n] x^(K-1), the associated natural number
   110  // is P(b^m).
   111  type poly struct {
   112  	k uint  // k is such that K = 1<<k.
   113  	m int   // the m such that P(b^m) is the original number.
   114  	a []nat // a slice of at most K m-word coefficients.
   115  }
   116  
   117  // polyFromNat slices the number x into a polynomial
   118  // with 1<<k coefficients made of m words.
   119  func polyFromNat(x nat, k uint, m int) poly {
   120  	p := poly{k: k, m: m}
   121  	length := len(x)/m + 1
   122  	p.a = make([]nat, length)
   123  	for i := range p.a {
   124  		if len(x) < m {
   125  			p.a[i] = make(nat, m)
   126  			copy(p.a[i], x)
   127  			break
   128  		}
   129  		p.a[i] = x[:m]
   130  		x = x[m:]
   131  	}
   132  	return p
   133  }
   134  
   135  // Int evaluates back a poly to its integer value.
   136  func (p *poly) Int() nat {
   137  	length := len(p.a)*p.m + 1
   138  	if na := len(p.a); na > 0 {
   139  		length += len(p.a[na-1])
   140  	}
   141  	n := make(nat, length)
   142  	m := p.m
   143  	np := n
   144  	for i := range p.a {
   145  		l := len(p.a[i])
   146  		c := addVV(np[:l], np[:l], p.a[i])
   147  		if np[l] < ^big.Word(0) {
   148  			np[l] += c
   149  		} else {
   150  			addVW(np[l:], np[l:], c)
   151  		}
   152  		np = np[m:]
   153  	}
   154  	n = trim(n)
   155  	return n
   156  }
   157  
   158  func trim(n nat) nat {
   159  	for i := range n {
   160  		if n[len(n)-1-i] != 0 {
   161  			return n[:len(n)-i]
   162  		}
   163  	}
   164  	return nil
   165  }
   166  
   167  // Mul multiplies p and q modulo X^K-1, where K = 1<<p.k.
   168  // The product is done via a Fourier transform.
   169  func (p *poly) Mul(q *poly) poly {
   170  	// extra=2 because:
   171  	// * some power of 2 is a K-th root of unity when n is a multiple of K/2.
   172  	// * 2 itself is a square (see fermat.ShiftHalf)
   173  	n := valueSize(p.k, p.m, 2)
   174  
   175  	pv, qv := p.Transform(n), q.Transform(n)
   176  	rv := pv.Mul(&qv)
   177  	r := rv.InvTransform()
   178  	r.m = p.m
   179  	return r
   180  }
   181  
   182  // A polValues represents the value of a poly at the powers of a
   183  // K-th root of unity θ=2^(l/2) in Z/(b^n+1)Z, where b^n = 2^(K/4*l).
   184  type polValues struct {
   185  	k      uint     // k is such that K = 1<<k.
   186  	n      int      // the length of coefficients, n*_W a multiple of K/4.
   187  	values []fermat // a slice of K (n+1)-word values
   188  }
   189  
   190  // Transform evaluates p at θ^i for i = 0...K-1, where
   191  // θ is a K-th primitive root of unity in Z/(b^n+1)Z.
   192  func (p *poly) Transform(n int) polValues {
   193  	k := p.k
   194  	inputbits := make([]big.Word, (n+1)<<k)
   195  	input := make([]fermat, 1<<k)
   196  	// Now computed q(ω^i) for i = 0 ... K-1
   197  	valbits := make([]big.Word, (n+1)<<k)
   198  	values := make([]fermat, 1<<k)
   199  	for i := range values {
   200  		input[i] = inputbits[i*(n+1) : (i+1)*(n+1)]
   201  		if i < len(p.a) {
   202  			copy(input[i], p.a[i])
   203  		}
   204  		values[i] = fermat(valbits[i*(n+1) : (i+1)*(n+1)])
   205  	}
   206  	fourier(values, input, false, n, k)
   207  	return polValues{k, n, values}
   208  }
   209  
   210  // InvTransform reconstructs p (modulo X^K - 1) from its
   211  // values at θ^i for i = 0..K-1.
   212  func (v *polValues) InvTransform() poly {
   213  	k, n := v.k, v.n
   214  
   215  	// Perform an inverse Fourier transform to recover p.
   216  	pbits := make([]big.Word, (n+1)<<k)
   217  	p := make([]fermat, 1<<k)
   218  	for i := range p {
   219  		p[i] = fermat(pbits[i*(n+1) : (i+1)*(n+1)])
   220  	}
   221  	fourier(p, v.values, true, n, k)
   222  	// Divide by K, and untwist q to recover p.
   223  	u := make(fermat, n+1)
   224  	a := make([]nat, 1<<k)
   225  	for i := range p {
   226  		u.Shift(p[i], -int(k))
   227  		copy(p[i], u)
   228  		a[i] = nat(p[i])
   229  	}
   230  	return poly{k: k, m: 0, a: a}
   231  }
   232  
   233  // NTransform evaluates p at θω^i for i = 0...K-1, where
   234  // θ is a (2K)-th primitive root of unity in Z/(b^n+1)Z
   235  // and ω = θ².
   236  func (p *poly) NTransform(n int) polValues {
   237  	k := p.k
   238  	if len(p.a) >= 1<<k {
   239  		panic("Transform: len(p.a) >= 1<<k")
   240  	}
   241  	// θ is represented as a shift.
   242  	θshift := (n * _W) >> k
   243  	// p(x) = a_0 + a_1 x + ... + a_{K-1} x^(K-1)
   244  	// p(θx) = q(x) where
   245  	// q(x) = a_0 + θa_1 x + ... + θ^(K-1) a_{K-1} x^(K-1)
   246  	//
   247  	// Twist p by θ to obtain q.
   248  	tbits := make([]big.Word, (n+1)<<k)
   249  	twisted := make([]fermat, 1<<k)
   250  	src := make(fermat, n+1)
   251  	for i := range twisted {
   252  		twisted[i] = fermat(tbits[i*(n+1) : (i+1)*(n+1)])
   253  		if i < len(p.a) {
   254  			for i := range src {
   255  				src[i] = 0
   256  			}
   257  			copy(src, p.a[i])
   258  			twisted[i].Shift(src, θshift*i)
   259  		}
   260  	}
   261  
   262  	// Now computed q(ω^i) for i = 0 ... K-1
   263  	valbits := make([]big.Word, (n+1)<<k)
   264  	values := make([]fermat, 1<<k)
   265  	for i := range values {
   266  		values[i] = fermat(valbits[i*(n+1) : (i+1)*(n+1)])
   267  	}
   268  	fourier(values, twisted, false, n, k)
   269  	return polValues{k, n, values}
   270  }
   271  
   272  // InvTransform reconstructs a polynomial from its values at
   273  // roots of x^K+1. The m field of the returned polynomial
   274  // is unspecified.
   275  func (v *polValues) InvNTransform() poly {
   276  	k := v.k
   277  	n := v.n
   278  	θshift := (n * _W) >> k
   279  
   280  	// Perform an inverse Fourier transform to recover q.
   281  	qbits := make([]big.Word, (n+1)<<k)
   282  	q := make([]fermat, 1<<k)
   283  	for i := range q {
   284  		q[i] = fermat(qbits[i*(n+1) : (i+1)*(n+1)])
   285  	}
   286  	fourier(q, v.values, true, n, k)
   287  
   288  	// Divide by K, and untwist q to recover p.
   289  	u := make(fermat, n+1)
   290  	a := make([]nat, 1<<k)
   291  	for i := range q {
   292  		u.Shift(q[i], -int(k)-i*θshift)
   293  		copy(q[i], u)
   294  		a[i] = nat(q[i])
   295  	}
   296  	return poly{k: k, m: 0, a: a}
   297  }
   298  
   299  // fourier performs an unnormalized Fourier transform
   300  // of src, a length 1<<k vector of numbers modulo b^n+1
   301  // where b = 1<<_W.
   302  func fourier(dst []fermat, src []fermat, backward bool, n int, k uint) {
   303  	var rec func(dst, src []fermat, size uint)
   304  	tmp := make(fermat, n+1)  // pre-allocate temporary variables.
   305  	tmp2 := make(fermat, n+1) // pre-allocate temporary variables.
   306  
   307  	// The recursion function of the FFT.
   308  	// The root of unity used in the transform is ω=1<<(ω2shift/2).
   309  	// The source array may use shifted indices (i.e. the i-th
   310  	// element is src[i << idxShift]).
   311  	rec = func(dst, src []fermat, size uint) {
   312  		idxShift := k - size
   313  		ω2shift := (4 * n * _W) >> size
   314  		if backward {
   315  			ω2shift = -ω2shift
   316  		}
   317  
   318  		// Easy cases.
   319  		if len(src[0]) != n+1 || len(dst[0]) != n+1 {
   320  			panic("len(src[0]) != n+1 || len(dst[0]) != n+1")
   321  		}
   322  		switch size {
   323  		case 0:
   324  			copy(dst[0], src[0])
   325  			return
   326  		case 1:
   327  			dst[0].Add(src[0], src[1<<idxShift]) // dst[0] = src[0] + src[1]
   328  			dst[1].Sub(src[0], src[1<<idxShift]) // dst[1] = src[0] - src[1]
   329  			return
   330  		}
   331  
   332  		// Let P(x) = src[0] + src[1<<idxShift] * x + ... + src[K-1 << idxShift] * x^(K-1)
   333  		// The P(x) = Q1(x²) + x*Q2(x²)
   334  		// where Q1's coefficients are src with indices shifted by 1
   335  		// where Q2's coefficients are src[1<<idxShift:] with indices shifted by 1
   336  
   337  		// Split destination vectors in halves.
   338  		dst1 := dst[:1<<(size-1)]
   339  		dst2 := dst[1<<(size-1):]
   340  		// Transform Q1 and Q2 in the halves.
   341  		rec(dst1, src, size-1)
   342  		rec(dst2, src[1<<idxShift:], size-1)
   343  
   344  		// Reconstruct P's transform from transforms of Q1 and Q2.
   345  		// dst[i]            is dst1[i] + ω^i * dst2[i]
   346  		// dst[i + 1<<(k-1)] is dst1[i] + ω^(i+K/2) * dst2[i]
   347  		//
   348  		for i := range dst1 {
   349  			tmp.ShiftHalf(dst2[i], i*ω2shift, tmp2) // ω^i * dst2[i]
   350  			dst2[i].Sub(dst1[i], tmp)
   351  			dst1[i].Add(dst1[i], tmp)
   352  		}
   353  	}
   354  	rec(dst, src, k)
   355  }
   356  
   357  // Mul returns the pointwise product of p and q.
   358  func (p *polValues) Mul(q *polValues) (r polValues) {
   359  	n := p.n
   360  	r.k, r.n = p.k, p.n
   361  	r.values = make([]fermat, len(p.values))
   362  	bits := make([]big.Word, len(p.values)*(n+1))
   363  	buf := make(fermat, 8*n)
   364  	for i := range r.values {
   365  		r.values[i] = bits[i*(n+1) : (i+1)*(n+1)]
   366  		z := buf.Mul(p.values[i], q.values[i])
   367  		copy(r.values[i], z)
   368  	}
   369  	return
   370  }