nihav-codec-support/src/dsp/fft.rs

   1 //! FFT and RDFT implementation.
   2 use std::f32::{self, consts};
   3 use std::ops::{Not, Neg, Add, AddAssign, Sub, SubAssign, Mul, MulAssign, Div};
   4 use std::fmt;
   5
   6 /// Complex number.
   7 #[repr(C)]
   8 #[derive(Debug,Clone,Copy,PartialEq)]
   9 pub struct FFTComplex {
  10     /// Real part of the numner.
  11     pub re: f32,
  12     /// Complex part of the number.
  13     pub im: f32,
  14 }
  15
  16 impl FFTComplex {
  17     /// Calculates `exp(i * val)`.
  18     pub fn exp(val: f32) -> Self {
  19         FFTComplex { re: val.cos(), im: val.sin() }
  20     }
  21     /// Returns `-Im + i * Re`.
  22     pub fn rotate(self) -> Self {
  23         FFTComplex { re: -self.im, im: self.re }
  24     }
  25     /// Multiplies complex number by scalar.
  26     pub fn scale(self, scale: f32) -> Self {
  27         FFTComplex { re: self.re * scale, im: self.im * scale }
  28     }
  29     /// Returns squared modulus value of the complex number.
  30     pub fn sq_modulus(self) -> f32 {
  31         self.re * self.re + self.im * self.im
  32     }
  33     /// Returns reciprocal of the complex number.
  34     pub fn reciprocal(self) -> Self {
  35         !self.scale(self.sq_modulus())
  36     }
  37 }
  38
  39 impl Neg for FFTComplex {
  40     type Output = FFTComplex;
  41     fn neg(self) -> Self::Output {
  42         FFTComplex { re: -self.re, im: -self.im }
  43     }
  44 }
  45
  46 impl Not for FFTComplex {
  47     type Output = FFTComplex;
  48     fn not(self) -> Self::Output {
  49         FFTComplex { re: self.re, im: -self.im }
  50     }
  51 }
  52
  53 impl Add for FFTComplex {
  54     type Output = FFTComplex;
  55     fn add(self, other: Self) -> Self::Output {
  56         FFTComplex { re: self.re + other.re, im: self.im + other.im }
  57     }
  58 }
  59
  60 impl AddAssign for FFTComplex {
  61     fn add_assign(&mut self, other: Self) {
  62         self.re += other.re;
  63         self.im += other.im;
  64     }
  65 }
  66
  67 impl Sub for FFTComplex {
  68     type Output = FFTComplex;
  69     fn sub(self, other: Self) -> Self::Output {
  70         FFTComplex { re: self.re - other.re, im: self.im - other.im }
  71     }
  72 }
  73
  74 impl SubAssign for FFTComplex {
  75     fn sub_assign(&mut self, other: Self) {
  76         self.re -= other.re;
  77         self.im -= other.im;
  78     }
  79 }
  80
  81 impl Mul for FFTComplex {
  82     type Output = FFTComplex;
  83     fn mul(self, other: Self) -> Self::Output {
  84         FFTComplex { re: self.re * other.re - self.im * other.im,
  85                      im: self.im * other.re + self.re * other.im }
  86     }
  87 }
  88
  89 impl MulAssign for FFTComplex {
  90     fn mul_assign(&mut self, other: Self) {
  91         let re = self.re * other.re - self.im * other.im;
  92         let im = self.im * other.re + self.re * other.im;
  93         self.re = re;
  94         self.im = im;
  95     }
  96 }
  97
  98 impl Div for FFTComplex {
  99     type Output = FFTComplex;
 100     fn div(self, other: Self) -> Self::Output {
 101         self * other.reciprocal()
 102     }
 103 }
 104
 105 impl fmt::Display for FFTComplex {
 106     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 107         write!(f, "({}, {})", self.re, self.im)
 108     }
 109 }
 110
 111 /// Complex number with zero value.
 112 pub const FFTC_ZERO: FFTComplex = FFTComplex { re: 0.0, im: 0.0 };
 113
 114 /// Calculates forward or inverse FFT in the straightforward way.
 115 pub fn generic_fft(data: &mut [FFTComplex], forward: bool) {
 116     let mut tmp = Vec::with_capacity(data.len());
 117     tmp.resize(data.len(), FFTC_ZERO);
 118     let base = if forward { -consts::PI * 2.0 / (data.len() as f32) }
 119                else       {  consts::PI * 2.0 / (data.len() as f32) };
 120     for k in 0..data.len() {
 121         let mut sum = FFTC_ZERO;
 122         for n in 0..data.len() {
 123             let w = FFTComplex::exp(base * ((n * k) as f32));
 124             sum += data[n] * w;
 125         }
 126         tmp[k] = sum;
 127     }
 128     for k in 0..data.len() {
 129         data[k] = tmp[k];
 130     }
 131 }
 132
 133 struct FFTData {
 134     table:      Vec<FFTComplex>,
 135     tmp:        Vec<FFTComplex>,
 136     twiddle:    Vec<FFTComplex>,
 137     size:       usize,
 138     step:       usize,
 139     div:        usize,
 140 }
 141
 142 struct FFTGeneric {}
 143
 144 const FFT3_CONST: f32 = 0.86602540378443864677;
 145 const FFT5_CONST1: FFTComplex = FFTComplex { re: 0.80901699437494742410, im: 0.58778525229247312915 };
 146 const FFT5_CONST2: FFTComplex = FFTComplex { re: 0.30901699437494742411, im: 0.95105651629515357211 };
 147
 148 fn twiddle5(a: FFTComplex, b: FFTComplex, c: FFTComplex) -> (FFTComplex, FFTComplex) {
 149     let re = a.re * c.re;
 150     let im = a.im * c.re;
 151     let diffre = b.im * c.im;
 152     let diffim = b.re * c.im;
 153
 154     (FFTComplex { re: re - diffre, im: im + diffim }, FFTComplex { re: re + diffre, im: im - diffim })
 155 }
 156
 157 impl FFTGeneric {
 158     fn new_data(size: usize, forward: bool) -> FFTData {
 159         let mut table: Vec<FFTComplex> = Vec::with_capacity(size * size);
 160         table.resize(size * size, FFTC_ZERO);
 161         let base = consts::PI * 2.0 / (size as f32);
 162         if forward {
 163             for n in 0..size {
 164                 for k in 0..size {
 165                     table[n * size + k] = FFTComplex::exp(-base * ((n * k) as f32));
 166                 }
 167             }
 168         } else {
 169             for n in 0..size {
 170                 for k in 0..size {
 171                     table[n * size + k] = FFTComplex::exp( base * ((n * k) as f32));
 172                 }
 173             }
 174         }
 175         let mut tmp = Vec::with_capacity(size);
 176         tmp.resize(size, FFTC_ZERO);
 177         FFTData { table, tmp, twiddle: Vec::new(), size, step: 0, div: 0 }
 178     }
 179     fn fft(tbl: &mut FFTData, size: usize, data: &mut [FFTComplex], step: usize) {
 180         if size == 3 {
 181             let s0 = data[step * 0];
 182             let s1 = data[step * 1];
 183             let s2 = data[step * 2];
 184             let t0 = s1 + s2;
 185             data[step * 0] += t0;
 186             let t1 = s0 - t0.scale(0.5);
 187             let t2 = (s2 - s1).rotate().scale(FFT3_CONST);
 188             data[step * 1] = t1 + t2;
 189             data[step * 2] = t1 - t2;
 190             return;
 191         }
 192         if size == 5 {
 193             let s0 = data[step * 0];
 194             let s1 = data[step * 1];
 195             let s2 = data[step * 2];
 196             let s3 = data[step * 3];
 197             let s4 = data[step * 4];
 198
 199             let t0 = s1 + s4;
 200             let t1 = s1 - s4;
 201             let t2 = s2 + s3;
 202             let t3 = s2 - s3;
 203             let (t4, t5) = twiddle5(t0, t1, FFT5_CONST2);
 204             let (t6, t7) = twiddle5(t2, t3, FFT5_CONST1);
 205             let (t8, t9) = twiddle5(t0, t1, FFT5_CONST1);
 206             let (ta, tb) = twiddle5(t2, t3, FFT5_CONST2);
 207
 208             data[step * 0] = s0 + t0 + t2;
 209             data[step * 1] = s0 + t5 - t6;
 210             data[step * 2] = s0 - t8 + ta;
 211             data[step * 3] = s0 - t9 + tb;
 212             data[step * 4] = s0 + t4 - t7;
 213             return;
 214         }
 215         for k in 0..tbl.size {
 216             tbl.tmp[k] = FFTC_ZERO;
 217             for n in 0..tbl.size {
 218                 tbl.tmp[k] += data[n * step] * tbl.table[k * tbl.size + n];
 219             }
 220         }
 221         for n in 0..tbl.size {
 222             data[n * step] = tbl.tmp[n];
 223         }
 224     }
 225     fn ifft(tbl: &mut FFTData, size: usize, data: &mut [FFTComplex], step: usize) {
 226         if size == 3 {
 227             let s0 = data[step * 0];
 228             let s1 = data[step * 1];
 229             let s2 = data[step * 2];
 230             let t0 = s1 + s2;
 231             data[step * 0] += t0;
 232             let t1 = s0 - t0.scale(0.5);
 233             let t2 = (s2 - s1).rotate().scale(FFT3_CONST);
 234             data[step * 1] = t1 - t2;
 235             data[step * 2] = t1 + t2;
 236             return;
 237         }
 238         if size == 5 {
 239             let s0 = data[step * 0];
 240             let s1 = data[step * 1];
 241             let s2 = data[step * 2];
 242             let s3 = data[step * 3];
 243             let s4 = data[step * 4];
 244
 245             let t0 = s1 + s4;
 246             let t1 = s1 - s4;
 247             let t2 = s2 + s3;
 248             let t3 = s2 - s3;
 249             let (t4, t5) = twiddle5(t0, t1, FFT5_CONST2);
 250             let (t6, t7) = twiddle5(t2, t3, FFT5_CONST1);
 251             let (t8, t9) = twiddle5(t0, t1, FFT5_CONST1);
 252             let (ta, tb) = twiddle5(t2, t3, FFT5_CONST2);
 253
 254             data[step * 0] = s0 + t0 + t2;
 255             data[step * 1] = s0 + t4 - t7;
 256             data[step * 2] = s0 - t9 + tb;
 257             data[step * 3] = s0 - t8 + ta;
 258             data[step * 4] = s0 + t5 - t6;
 259             return;
 260         }
 261         Self::fft(tbl, size, data, step);
 262     }
 263 }
 264
 265 struct FFTSplitRadix {}
 266
 267 impl FFTSplitRadix {
 268     fn new_data(bits: u8, _forward: bool) -> FFTData {
 269         let size = 1 << bits;
 270         let mut table = Vec::with_capacity(size);
 271         for _ in 0..4 { table.push(FFTC_ZERO); }
 272         for b in 3..=bits {
 273             let qsize = (1 << (b - 2)) as usize;
 274             let base = -consts::PI / ((qsize * 2) as f32);
 275             for k in 0..qsize {
 276                 table.push(FFTComplex::exp(base * ((k * 1) as f32)));
 277                 table.push(FFTComplex::exp(base * ((k * 3) as f32)));
 278             }
 279         }
 280         FFTData { table, tmp: Vec::new(), twiddle: Vec::new(), size, step: 0, div: 0 }
 281     }
 282     fn fft(fftdata: &mut FFTData, bits: u8, data: &mut [FFTComplex]) {
 283         if bits == 0 { return; }
 284         if bits == 1 {
 285             let sum01 = data[0] + data[1];
 286             let dif01 = data[0] - data[1];
 287             data[0] = sum01;
 288             data[1] = dif01;
 289             return;
 290         }
 291         if bits == 2 {
 292             let sum01 = data[0] + data[2];
 293             let dif01 = data[0] - data[2];
 294             let sum23 = data[1] + data[3];
 295             let dif23 = data[1] - data[3];
 296             data[0] = sum01 + sum23;
 297             data[1] = dif01 - dif23.rotate();
 298             data[2] = sum01 - sum23;
 299             data[3] = dif01 + dif23.rotate();
 300             return;
 301         }
 302         let qsize = (1 << (bits - 2)) as usize;
 303         let hsize = (1 << (bits - 1)) as usize;
 304         let q3size = qsize + hsize;
 305
 306         Self::fft(fftdata, bits - 1, &mut data[0     ..hsize]);
 307         Self::fft(fftdata, bits - 2, &mut data[hsize ..q3size]);
 308         Self::fft(fftdata, bits - 2, &mut data[q3size..]);
 309         let off = hsize;
 310         {
 311             let t3 =  data[0 + hsize] + data[0 + q3size];
 312             let t4 = (data[0 + hsize] - data[0 + q3size]).rotate();
 313             let e1 = data[0];
 314             let e2 = data[0 + qsize];
 315             data[0]          = e1 + t3;
 316             data[0 + qsize]  = e2 - t4;
 317             data[0 + hsize]  = e1 - t3;
 318             data[0 + q3size] = e2 + t4;
 319         }
 320         for k in 1..qsize {
 321             let t1 = fftdata.table[off + k * 2 + 0] * data[k + hsize];
 322             let t2 = fftdata.table[off + k * 2 + 1] * data[k + q3size];
 323             let t3 =  t1 + t2;
 324             let t4 = (t1 - t2).rotate();
 325             let e1 = data[k];
 326             let e2 = data[k + qsize];
 327             data[k]             = e1 + t3;
 328             data[k + qsize]     = e2 - t4;
 329             data[k + hsize]     = e1 - t3;
 330             data[k + qsize * 3] = e2 + t4;
 331         }
 332     }
 333     fn ifft(fftdata: &mut FFTData, bits: u8, data: &mut [FFTComplex]) {
 334         if bits == 0 { return; }
 335         if bits == 1 {
 336             let sum01 = data[0] + data[1];
 337             let dif01 = data[0] - data[1];
 338             data[0] = sum01;
 339             data[1] = dif01;
 340             return;
 341         }
 342         if bits == 2 {
 343             let sum01 = data[0] + data[2];
 344             let dif01 = data[0] - data[2];
 345             let sum23 = data[1] + data[3];
 346             let dif23 = data[1] - data[3];
 347             data[0] = sum01 + sum23;
 348             data[1] = dif01 + dif23.rotate();
 349             data[2] = sum01 - sum23;
 350             data[3] = dif01 - dif23.rotate();
 351             return;
 352         }
 353         let qsize = (1 << (bits - 2)) as usize;
 354         let hsize = (1 << (bits - 1)) as usize;
 355         let q3size = qsize + hsize;
 356
 357         Self::ifft(fftdata, bits - 1, &mut data[0     ..hsize]);
 358         Self::ifft(fftdata, bits - 2, &mut data[hsize ..q3size]);
 359         Self::ifft(fftdata, bits - 2, &mut data[q3size..]);
 360         let off = hsize;
 361         {
 362             let t3 =  data[0 + hsize] + data[0 + q3size];
 363             let t4 = (data[0 + hsize] - data[0 + q3size]).rotate();
 364             let e1 = data[0];
 365             let e2 = data[0 + qsize];
 366             data[0]          = e1 + t3;
 367             data[0 + qsize]  = e2 + t4;
 368             data[0 + hsize]  = e1 - t3;
 369             data[0 + q3size] = e2 - t4;
 370         }
 371         for k in 1..qsize {
 372             let t1 = !fftdata.table[off + k * 2 + 0] * data[k + hsize];
 373             let t2 = !fftdata.table[off + k * 2 + 1] * data[k + q3size];
 374             let t3 =  t1 + t2;
 375             let t4 = (t1 - t2).rotate();
 376             let e1 = data[k];
 377             let e2 = data[k + qsize];
 378             data[k]             = e1 + t3;
 379             data[k + qsize]     = e2 + t4;
 380             data[k + hsize]     = e1 - t3;
 381             data[k + qsize * 3] = e2 - t4;
 382         }
 383     }
 384 }
 385
 386 struct FFT15 {}
 387
 388 const FFT15_INSWAP:  [usize; 20] = [ 0, 5, 10, 42, 3, 8, 13, 42, 6, 11, 1, 42, 9, 14, 4, 42, 12, 2, 7, 42 ];
 389 const FFT15_OUTSWAP: [usize; 20] = [ 0, 10, 5, 42, 6, 1, 11, 42, 12, 7, 2, 42, 3, 13, 8, 42, 9, 4, 14, 42 ];
 390
 391 impl FFT15 {
 392     fn new_data(size: usize, _forward: bool) -> FFTData {
 393         FFTData { table: Vec::new(), tmp: Vec::new(), twiddle: Vec::new(), size, step: 0, div: 0 }
 394     }
 395     fn fft3(dst: &mut [FFTComplex], src: &[FFTComplex], step: usize, n: usize) {
 396         let s0 = src[0];
 397         let s1 = src[1];
 398         let s2 = src[2];
 399
 400         let t0 = s1 + s2;
 401         let t1 = s0 - t0.scale(0.5);
 402         let t2 = (s2 - s1).rotate().scale(FFT3_CONST);
 403
 404         dst[FFT15_OUTSWAP[n * 4 + 0] * step] = s0 + t0;
 405         dst[FFT15_OUTSWAP[n * 4 + 1] * step] = t1 + t2;
 406         dst[FFT15_OUTSWAP[n * 4 + 2] * step] = t1 - t2;
 407     }
 408     fn ifft3(dst: &mut [FFTComplex], src: &[FFTComplex], step: usize, n: usize) {
 409         let s0 = src[0];
 410         let s1 = src[1];
 411         let s2 = src[2];
 412
 413         let t0 = s1 + s2;
 414         let t1 = s0 - t0.scale(0.5);
 415         let t2 = (s2 - s1).rotate().scale(FFT3_CONST);
 416
 417         dst[FFT15_OUTSWAP[n * 4 + 0] * step] = s0 + t0;
 418         dst[FFT15_OUTSWAP[n * 4 + 1] * step] = t1 - t2;
 419         dst[FFT15_OUTSWAP[n * 4 + 2] * step] = t1 + t2;
 420     }
 421     fn fft5(dst: &mut [FFTComplex], src: &[FFTComplex], step: usize, n: usize) {
 422         let s0 = src[FFT15_INSWAP[n + 0 * 4] * step];
 423         let s1 = src[FFT15_INSWAP[n + 1 * 4] * step];
 424         let s2 = src[FFT15_INSWAP[n + 2 * 4] * step];
 425         let s3 = src[FFT15_INSWAP[n + 3 * 4] * step];
 426         let s4 = src[FFT15_INSWAP[n + 4 * 4] * step];
 427
 428         let t0 = s1 + s4;
 429         let t1 = s1 - s4;
 430         let t2 = s2 + s3;
 431         let t3 = s2 - s3;
 432         let (t4, t5) = twiddle5(t0, t1, FFT5_CONST2);
 433         let (t6, t7) = twiddle5(t2, t3, FFT5_CONST1);
 434         let (t8, t9) = twiddle5(t0, t1, FFT5_CONST1);
 435         let (ta, tb) = twiddle5(t2, t3, FFT5_CONST2);
 436
 437         dst[0 * 3] = s0 + t0 + t2;
 438         dst[1 * 3] = s0 + t5 - t6;
 439         dst[2 * 3] = s0 - t8 + ta;
 440         dst[3 * 3] = s0 - t9 + tb;
 441         dst[4 * 3] = s0 + t4 - t7;
 442     }
 443     fn ifft5(dst: &mut [FFTComplex], src: &[FFTComplex], step: usize, n: usize) {
 444         let s0 = src[FFT15_INSWAP[n + 0 * 4] * step];
 445         let s1 = src[FFT15_INSWAP[n + 1 * 4] * step];
 446         let s2 = src[FFT15_INSWAP[n + 2 * 4] * step];
 447         let s3 = src[FFT15_INSWAP[n + 3 * 4] * step];
 448         let s4 = src[FFT15_INSWAP[n + 4 * 4] * step];
 449
 450         let t0 = s1 + s4;
 451         let t1 = s1 - s4;
 452         let t2 = s2 + s3;
 453         let t3 = s2 - s3;
 454         let (t4, t5) = twiddle5(t0, t1, FFT5_CONST2);
 455         let (t6, t7) = twiddle5(t2, t3, FFT5_CONST1);
 456         let (t8, t9) = twiddle5(t0, t1, FFT5_CONST1);
 457         let (ta, tb) = twiddle5(t2, t3, FFT5_CONST2);
 458
 459         dst[0 * 3] = s0 + t0 + t2;
 460         dst[1 * 3] = s0 + t4 - t7;
 461         dst[2 * 3] = s0 - t9 + tb;
 462         dst[3 * 3] = s0 - t8 + ta;
 463         dst[4 * 3] = s0 + t5 - t6;
 464     }
 465     fn fft(_fftdata: &mut FFTData, data: &mut [FFTComplex], step: usize) {
 466         let mut tmp = [FFTC_ZERO; 15];
 467         for n in 0..3 {
 468             Self::fft5(&mut tmp[n..], data, step, n);
 469         }
 470         for n in 0..5 {
 471             Self::fft3(data, &tmp[n * 3..][..3], step, n);
 472         }
 473     }
 474     fn ifft(_fftdata: &mut FFTData, data: &mut [FFTComplex], step: usize) {
 475         let mut tmp = [FFTC_ZERO; 15];
 476         for n in 0..3 {
 477             Self::ifft5(&mut tmp[n..], data, step, n);
 478         }
 479         for n in 0..5 {
 480             Self::ifft3(data, &tmp[n * 3..][..3], step, n);
 481         }
 482     }
 483 }
 484
 485
 486 enum FFTMode {
 487     Generic(usize),
 488     SplitRadix(u8),
 489     Prime15,
 490 }
 491
 492 impl FFTMode {
 493     fn permute(&self, perms: &mut [usize]) {
 494         match *self {
 495             FFTMode::Generic(_)         => {},
 496             FFTMode::SplitRadix(bits)   => {
 497                 let div = perms.len() >> bits;
 498                 gen_sr_perms(perms, 1 << bits);
 499                 if div > 1 {
 500                     for i in 0..(1 << bits) {
 501                         perms[i] *= div;
 502                     }
 503                     for i in 1..div {
 504                         for j in 0..(1 << bits) {
 505                             perms[(i << bits) + j] = perms[j] + i;
 506                         }
 507                     }
 508                 }
 509             },
 510             FFTMode::Prime15            => {},
 511         };
 512     }
 513     fn do_fft(&self, fftdata: &mut FFTData, data: &mut [FFTComplex]) {
 514         match *self {
 515             FFTMode::Generic(size)      => FFTGeneric::fft(fftdata, size, data, 1),
 516             FFTMode::SplitRadix(bits)   => FFTSplitRadix::fft(fftdata, bits, data),
 517             FFTMode::Prime15            => FFT15::fft(fftdata, data, 1),
 518         };
 519     }
 520     fn do_fft2(&self, fftdata: &mut FFTData, data: &mut [FFTComplex], step: usize) {
 521         match *self {
 522             FFTMode::Generic(size)      => FFTGeneric::fft(fftdata, size, data, step),
 523             FFTMode::SplitRadix(_)      => unreachable!(),
 524             FFTMode::Prime15            => FFT15::fft(fftdata, data, step),
 525         };
 526     }
 527     fn do_ifft(&self, fftdata: &mut FFTData, data: &mut [FFTComplex]) {
 528         match *self {
 529             FFTMode::Generic(size)      => FFTGeneric::ifft(fftdata, size, data, 1),
 530             FFTMode::SplitRadix(bits)   => FFTSplitRadix::ifft(fftdata, bits, data),
 531             FFTMode::Prime15            => FFT15::ifft(fftdata, data, 1),
 532         };
 533     }
 534     fn do_ifft2(&self, fftdata: &mut FFTData, data: &mut [FFTComplex], step: usize) {
 535         match *self {
 536             FFTMode::Generic(size)      => FFTGeneric::ifft(fftdata, size, data, step),
 537             FFTMode::SplitRadix(_)      => unreachable!(),
 538             FFTMode::Prime15            => FFT15::ifft(fftdata, data, step),
 539         };
 540     }
 541     fn get_size(&self) -> usize {
 542         match *self {
 543             FFTMode::Generic(size)      => size,
 544             FFTMode::SplitRadix(bits)   => 1 << bits,
 545             FFTMode::Prime15            => 15,
 546         }
 547     }
 548 }
 549
 550 /// FFT working context.
 551 pub struct FFT {
 552     perms:  Vec<usize>,
 553     swaps:  Vec<usize>,
 554     ffts:   Vec<(FFTMode, FFTData)>,
 555 }
 556
 557 impl FFT {
 558     /// Calculates Fourier transform.
 559     pub fn do_fft(&mut self, src: &[FFTComplex], dst: &mut [FFTComplex]) {
 560         for k in 0..src.len() { dst[k] = src[self.perms[k]]; }
 561         self.do_fft_core(dst);
 562     }
 563     /// Calculates inverse Fourier transform.
 564     pub fn do_ifft(&mut self, src: &[FFTComplex], dst: &mut [FFTComplex]) {
 565         for k in 0..src.len() { dst[k] = src[self.perms[k]]; }
 566         self.do_ifft_core(dst);
 567     }
 568     /// Performs inplace FFT.
 569     pub fn do_fft_inplace(&mut self, data: &mut [FFTComplex]) {
 570         for idx in 0..self.swaps.len() {
 571             let nidx = self.swaps[idx];
 572             if idx != nidx {
 573                 data.swap(nidx, idx);
 574             }
 575         }
 576         self.do_fft_core(data);
 577     }
 578     /// Performs inplace inverse FFT.
 579     pub fn do_ifft_inplace(&mut self, data: &mut [FFTComplex]) {
 580         for idx in 0..self.swaps.len() {
 581             let nidx = self.swaps[idx];
 582             if idx != nidx {
 583                 data.swap(nidx, idx);
 584             }
 585         }
 586         self.do_ifft_core(data);
 587     }
 588     fn do_fft_core(&mut self, data: &mut [FFTComplex]) {
 589         for el in self.ffts.iter_mut() {
 590             let (mode, ref mut fftdata) = el;
 591             let bsize = mode.get_size();
 592             let div = fftdata.div;
 593             let step = fftdata.step;
 594             if step == 1 {
 595                 mode.do_fft(fftdata, data);
 596                 for i in 1..div {
 597                     mode.do_fft(fftdata, &mut data[i * bsize..]);
 598                 }
 599             } else {
 600                 mode.do_fft2(fftdata, data, div);
 601                 let mut toff = bsize;
 602                 for i in 1..div {
 603                     for j in 1..bsize {
 604                         data[i + j * div] *= fftdata.twiddle[toff + j];
 605                     }
 606                     mode.do_fft2(fftdata, &mut data[i..], div);
 607                     toff += bsize;
 608                 }
 609             }
 610         }
 611     }
 612     fn do_ifft_core(&mut self, data: &mut [FFTComplex]) {
 613         for el in self.ffts.iter_mut() {
 614             let (mode, ref mut fftdata) = el;
 615             let bsize = mode.get_size();
 616             let div = fftdata.div;
 617             let step = fftdata.step;
 618             if step == 1 {
 619                 mode.do_ifft(fftdata, data);
 620                 for i in 1..div {
 621                     mode.do_ifft(fftdata, &mut data[i * bsize..]);
 622                 }
 623             } else {
 624                 mode.do_ifft2(fftdata, data, div);
 625                 let mut toff = bsize;
 626                 for i in 1..div {
 627                     for j in 1..bsize {
 628                         data[i + j * div] *= fftdata.twiddle[toff + j];
 629                     }
 630                     mode.do_ifft2(fftdata, &mut data[i..], div);
 631                     toff += bsize;
 632                 }
 633             }
 634         }
 635     }
 636 }
 637
 638 /// [`FFT`] context creator.
 639 ///
 640 /// [`FFT`]: ./struct.FFT.html
 641 pub struct FFTBuilder {
 642 }
 643
 644 /*fn reverse_bits(inval: u32) -> u32 {
 645     const REV_TAB: [u8; 16] = [
 646         0b0000, 0b1000, 0b0100, 0b1100, 0b0010, 0b1010, 0b0110, 0b1110,
 647         0b0001, 0b1001, 0b0101, 0b1101, 0b0011, 0b1011, 0b0111, 0b1111,
 648     ];
 649
 650     let mut ret = 0;
 651     let mut val = inval;
 652     for _ in 0..8 {
 653         ret = (ret << 4) | (REV_TAB[(val & 0xF) as usize] as u32);
 654         val = val >> 4;
 655     }
 656     ret
 657 }
 658
 659 fn swp_idx(idx: usize, bits: u32) -> usize {
 660     let s = reverse_bits(idx as u32) as usize;
 661     s >> (32 - bits)
 662 }*/
 663
 664 fn gen_sr_perms(swaps: &mut [usize], size: usize) {
 665     if size <= 4 { return; }
 666     let mut evec:  Vec<usize> = Vec::with_capacity(size / 2);
 667     let mut ovec1: Vec<usize> = Vec::with_capacity(size / 4);
 668     let mut ovec2: Vec<usize> = Vec::with_capacity(size / 4);
 669     for k in 0..size/4 {
 670         evec.push (swaps[k * 4 + 0]);
 671         ovec1.push(swaps[k * 4 + 1]);
 672         evec.push (swaps[k * 4 + 2]);
 673         ovec2.push(swaps[k * 4 + 3]);
 674     }
 675     for k in 0..size/2 { swaps[k]            = evec[k]; }
 676     for k in 0..size/4 { swaps[k +   size/2] = ovec1[k]; }
 677     for k in 0..size/4 { swaps[k + 3*size/4] = ovec2[k]; }
 678     gen_sr_perms(&mut swaps[0..size/2],        size/2);
 679     gen_sr_perms(&mut swaps[size/2..3*size/4], size/4);
 680     gen_sr_perms(&mut swaps[3*size/4..],       size/4);
 681 }
 682
 683 fn gen_swaps_for_perm(swaps: &mut Vec<usize>, perms: &[usize]) {
 684     let mut idx_arr: Vec<usize> = Vec::with_capacity(perms.len());
 685     for i in 0..perms.len() { idx_arr.push(i); }
 686     let mut run_size = 0;
 687     let mut run_pos  = 0;
 688     for idx in 0..perms.len() {
 689         if perms[idx] == idx_arr[idx] {
 690             if run_size == 0 { run_pos = idx; }
 691             run_size += 1;
 692         } else {
 693             for i in 0..run_size {
 694                 swaps.push(run_pos + i);
 695             }
 696             run_size = 0;
 697             let mut spos = idx + 1;
 698             while idx_arr[spos] != perms[idx] { spos += 1; }
 699             idx_arr[spos] = idx_arr[idx];
 700             idx_arr[idx]  = perms[idx];
 701             swaps.push(spos);
 702         }
 703     }
 704 }
 705
 706 impl FFTBuilder {
 707     fn generate_twiddle(data: &mut FFTData, size: usize, cur_size: usize, forward: bool) {
 708         if size == cur_size { return; }
 709         data.twiddle = Vec::with_capacity(size);
 710         let div = size / cur_size;
 711         let base = if forward { -2.0 * consts::PI / (size as f32) } else { 2.0 * consts::PI / (size as f32) };
 712         for n in 0..div {
 713             for k in 0..cur_size {
 714                 data.twiddle.push(FFTComplex::exp(base * ((k * n) as f32)));
 715             }
 716         }
 717     }
 718     /// Constructs a new `FFT` context.
 719     pub fn new_fft(size: usize, forward: bool) -> FFT {
 720         let mut ffts: Vec<(FFTMode, FFTData)> = Vec::with_capacity(1);
 721         let mut perms: Vec<usize> = Vec::with_capacity(size);
 722         let mut swaps: Vec<usize> = Vec::with_capacity(size);
 723         let mut rem_size = size;
 724         if rem_size.trailing_zeros() > 0 {
 725             let bits = rem_size.trailing_zeros() as u8;
 726             let mut data = FFTSplitRadix::new_data(bits, forward);
 727             Self::generate_twiddle(&mut data, size, 1 << bits, forward);
 728             data.step = 1;
 729             data.div = rem_size >> bits;
 730             ffts.push((FFTMode::SplitRadix(bits), data));
 731             rem_size >>= bits;
 732         }
 733         if (rem_size % 15) == 0 {
 734             let mut data = FFT15::new_data(size, forward);
 735             Self::generate_twiddle(&mut data, size, 15, forward);
 736             data.step = size / rem_size;
 737             data.div = size / rem_size;
 738             ffts.push((FFTMode::Prime15, data));
 739             rem_size /= 15;
 740         }
 741         if rem_size > 1 {
 742             let mut data = FFTGeneric::new_data(rem_size, forward);
 743             Self::generate_twiddle(&mut data, size, rem_size, forward);
 744             data.step = size / rem_size;
 745             data.div  = size / rem_size;
 746             ffts.push((FFTMode::Generic(rem_size), data));
 747         }
 748
 749         for i in 0..size {
 750             perms.push(i);
 751         }
 752         for (mode, _) in ffts.iter().rev() {
 753             mode.permute(&mut perms);
 754         }
 755         gen_swaps_for_perm(&mut swaps, perms.as_slice());
 756
 757         FFT { perms, swaps, ffts }
 758     }
 759 }
 760
 761 /// RDFT working context.
 762 pub struct RDFT {
 763     table:  Vec<FFTComplex>,
 764     fft:    FFT,
 765     fwd:    bool,
 766     size:   usize,
 767     fwd_fft:    bool,
 768 }
 769
 770 fn crossadd(a: FFTComplex, b: FFTComplex) -> FFTComplex {
 771     FFTComplex { re: a.re + b.re, im: a.im - b.im }
 772 }
 773
 774 impl RDFT {
 775     /// Calculates RDFT.
 776     pub fn do_rdft(&mut self, src: &[FFTComplex], dst: &mut [FFTComplex]) {
 777         dst.copy_from_slice(src);
 778         self.do_rdft_inplace(dst);
 779     }
 780     /// Calculates inplace RDFT.
 781     pub fn do_rdft_inplace(&mut self, buf: &mut [FFTComplex]) {
 782         if !self.fwd {
 783             for n in 0..self.size/2 {
 784                 let in0 = buf[n + 1];
 785                 let in1 = buf[self.size - n - 1];
 786
 787                 let t0 = crossadd(in0, in1);
 788                 let t1 = FFTComplex { re: in1.im + in0.im, im: in1.re - in0.re };
 789                 let tab = self.table[n];
 790                 let t2 = FFTComplex { re: t1.im * tab.im + t1.re * tab.re, im: t1.im * tab.re - t1.re * tab.im };
 791
 792                 buf[n + 1] = FFTComplex { re: t0.im - t2.im, im: t0.re - t2.re }; // (t0 - t2).conj().rotate()
 793                 buf[self.size - n - 1] = (t0 + t2).rotate();
 794             }
 795             let a = buf[0].re;
 796             let b = buf[0].im;
 797             buf[0].re = a - b;
 798             buf[0].im = a + b;
 799         }
 800         if self.fwd_fft {
 801             self.fft.do_fft_inplace(buf);
 802         } else {
 803             self.fft.do_ifft_inplace(buf);
 804         }
 805         if self.fwd {
 806             for n in 0..self.size/2 {
 807                 let in0 = buf[n + 1];
 808                 let in1 = buf[self.size - n - 1];
 809
 810                 let t0 = crossadd(in0, in1).scale(0.5);
 811                 let t1 = FFTComplex { re: in0.im + in1.im, im: in0.re - in1.re };
 812                 let t2 = t1 * self.table[n];
 813
 814                 buf[n + 1] = crossadd(t0, t2);
 815                 buf[self.size - n - 1] = FFTComplex { re: t0.re - t2.re, im: -(t0.im + t2.im) };
 816             }
 817             let a = buf[0].re;
 818             let b = buf[0].im;
 819             buf[0].re = a + b;
 820             buf[0].im = a - b;
 821         } else {
 822             for n in 0..self.size {
 823                 buf[n] = FFTComplex{ re: buf[n].im, im: buf[n].re };
 824             }
 825         }
 826     }
 827 }
 828
 829 /// [`RDFT`] context creator.
 830 ///
 831 /// [`RDFT`]: ./struct.FFT.html
 832 pub struct RDFTBuilder {
 833 }
 834
 835 impl RDFTBuilder {
 836     /// Constructs a new `RDFT` context.
 837     pub fn new_rdft(size: usize, forward: bool, forward_fft: bool) -> RDFT {
 838         let mut table: Vec<FFTComplex> = Vec::with_capacity(size / 4);
 839         let (base, scale) = if forward { (consts::PI / (size as f32), 0.5) } else { (-consts::PI / (size as f32), 1.0) };
 840         for i in 0..size/2 {
 841             table.push(FFTComplex::exp(base * ((i + 1) as f32)).scale(scale));
 842         }
 843         let fft = FFTBuilder::new_fft(size, forward_fft);
 844         RDFT { table, fft, size, fwd: forward, fwd_fft: forward_fft }
 845     }
 846 }
 847
 848
 849 #[cfg(test)]
 850 mod test {
 851     use super::*;
 852
 853     fn test_fft(size: usize) {
 854         println!("testing FFT {}", size);
 855         let mut fin:   Vec<FFTComplex> = Vec::with_capacity(size);
 856         let mut fout1: Vec<FFTComplex> = Vec::with_capacity(size);
 857         let mut fout2: Vec<FFTComplex> = Vec::with_capacity(size);
 858         fin.resize(size, FFTC_ZERO);
 859         fout1.resize(size, FFTC_ZERO);
 860         fout2.resize(size, FFTC_ZERO);
 861         let mut fft = FFTBuilder::new_fft(size, true);
 862         let mut seed: u32 = 42;
 863         for i in 0..fin.len() {
 864             seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
 865             let val = (seed >> 16) as i16;
 866             fin[i].re = (val as f32) / 256.0;
 867             seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
 868             let val = (seed >> 16) as i16;
 869             fin[i].im = (val as f32) / 256.0;
 870         }
 871         fft.do_fft(&fin, &mut fout1);
 872         fout2.copy_from_slice(&fin);
 873         generic_fft(&mut fout2, true);
 874
 875         for i in 0..fin.len() {
 876             assert!((fout1[i].re - fout2[i].re).abs() < 1.0);
 877             assert!((fout1[i].im - fout2[i].im).abs() < 1.0);
 878         }
 879         let mut ifft = FFTBuilder::new_fft(size, false);
 880         ifft.do_ifft_inplace(&mut fout1);
 881         generic_fft(&mut fout2, false);
 882
 883         let sc = 1.0 / (size as f32);
 884         for i in 0..fin.len() {
 885             assert!((fin[i].re - fout1[i].re * sc).abs() < 1.0);
 886             assert!((fin[i].im - fout1[i].im * sc).abs() < 1.0);
 887             assert!((fout1[i].re - fout2[i].re).abs() * sc < 1.0);
 888             assert!((fout1[i].im - fout2[i].im).abs() * sc < 1.0);
 889         }
 890     }
 891
 892     #[test]
 893     fn test_ffts() {
 894         test_fft(3);
 895         test_fft(5);
 896         test_fft(16);
 897         test_fft(15);
 898         test_fft(60);
 899         test_fft(256);
 900         test_fft(240);
 901     }
 902
 903     #[test]
 904     fn test_rdft() {
 905         let mut fin:   [FFTComplex; 128] = [FFTC_ZERO; 128];
 906         let mut fout1: [FFTComplex; 128] = [FFTC_ZERO; 128];
 907         let mut rdft = RDFTBuilder::new_rdft(fin.len(), true, true);
 908         let mut seed: u32 = 42;
 909         for i in 0..fin.len() {
 910             seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
 911             let val = (seed >> 16) as i16;
 912             fin[i].re = (val as f32) / 256.0;
 913             seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
 914             let val = (seed >> 16) as i16;
 915             fin[i].im = (val as f32) / 256.0;
 916         }
 917         rdft.do_rdft(&fin, &mut fout1);
 918         let mut irdft = RDFTBuilder::new_rdft(fin.len(), false, true);
 919         irdft.do_rdft_inplace(&mut fout1);
 920
 921         for i in 0..fin.len() {
 922             let tst = fout1[i].scale(0.5/(fout1.len() as f32));
 923             assert!((tst.re - fin[i].re).abs() < 1.0);
 924             assert!((tst.im - fin[i].im).abs() < 1.0);
 925         }
 926     }
 927 }