/* Cooley-Tukey Radix2 DIT Algorithm.
   Takes in omegas in normal ordering.
*/

decl a: bit<32>[32];
decl prime: bit<32>[1];
decl omegas: bit<32>[32];

let q: bit<32> = prime[0];
let N: ubit<6> = 32;
let log2N: ubit<6> = 5;

let idx: ubit<6> = 0;
---
while (idx < N) {
  let rev_idx: ubit<6> = 0;
  let j: ubit<6> = 0;
  while (j < log2N) {
    // Reverse `idx`.
    rev_idx := rev_idx | (((idx >> j) & (1 as ubit<6>)) << ((log2N - 1) - j));
    ---
    j := j + 1;
  }
  if (rev_idx > idx) {
    // Swap `a[idx]`, `a[rev_idx]`.
    let temp: bit<32> = a[idx];
    ---
    a[idx] := a[rev_idx];
    ---
    a[rev_idx] := temp;
  }
  idx := idx + 1;
}

let mIndex: ubit<6> = 2;
let iters: ubit<6> = 0;

while (iters < log2N) {
  let i: ubit<6> = 0;
  let end: ubit<6> = N / mIndex;
  let mHalved: ubit<6> = mIndex >> 1;
  while (i < end) {
    let oidx: ubit<6> = 0;
    let iStep: ubit<6> = i * mIndex;
    let j: ubit<6> = 0;
    while (j < mHalved) {
      let k: ubit<6> = iStep + j;
      let U: bit<32> = a[k];
      ---
      let V: bit<32> = a[k + mHalved] * omegas[oidx];
      ---
      a[k] := (U + V) % q;
      ---
      a[k + mHalved] := (U - V) % q;
      oidx := oidx + end;
      ---
      j := j + 1;
    }
    i := i + 1;
  }
  mIndex := mIndex << 1;
  ---
  iters := iters + 1;
}