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