p3_dft/

radix_2_dit_parallel.rs

use alloc::slice;
use alloc::sync::Arc;
use alloc::vec::Vec;
use core::mem::{MaybeUninit, transmute};

use itertools::{Itertools, izip};
use p3_field::integers::QuotientMap;
use p3_field::{Field, Powers, PrimeCharacteristicRing, TwoAdicField};
use p3_matrix::Matrix;
use p3_matrix::bitrev::{BitReversalPerm, BitReversedMatrixView, BitReversibleMatrix};
use p3_matrix::dense::{RowMajorMatrix, RowMajorMatrixView, RowMajorMatrixViewMut};
use p3_matrix::util::reverse_matrix_index_bits;
use p3_maybe_rayon::prelude::*;
use p3_util::{log2_strict_usize, reverse_bits_len, reverse_slice_index_bits};
use tracing::{debug_span, instrument};

use crate::butterflies::{Butterfly, DitButterfly, ScaledDitButterfly, TwiddleFreeButterfly};
use crate::twiddle_cache::TwiddleCache;
use crate::{Layout, TwoAdicSubgroupDft};

/// A parallel FFT algorithm which divides a butterfly network's layers into two halves.
///
/// For the first half, we apply a butterfly network with smaller blocks in earlier layers,
/// i.e. either DIT or Bowers G. Then we bit-reverse, and for the second half, we continue executing
/// the same network but in bit-reversed order. This way we're always working with small blocks,
/// so within each half, we can have a certain amount of parallelism with no cross-thread
/// communication.
#[derive(Default, Clone, Debug)]
pub struct Radix2DitParallel<F> {
    /// Twiddles based on roots of unity, used in the forward DFT.
    twiddles: Arc<TwiddleCache<usize, VectorPair<F>>>,

    /// A map from `(log_h, shift)` to forward DFT twiddles with that coset shift baked in.
    #[allow(clippy::type_complexity)]
    coset_twiddles: Arc<TwiddleCache<(usize, F), [Vec<F>]>>,

    /// Twiddles based on inverse roots of unity, used in the inverse DFT.
    inverse_twiddles: Arc<TwiddleCache<usize, VectorPair<F>>>,
}

/// A pair of vectors, one with twiddle factors in their natural order, the other bit-reversed.
#[derive(Default, Clone, Debug)]
struct VectorPair<F> {
    twiddles: Vec<F>,
    bitrev_twiddles: Vec<F>,
}

impl<F> Radix2DitParallel<F>
where
    F: TwoAdicField + Ord,
{
    fn get_or_compute_twiddles(&self, log_h: usize) -> Arc<VectorPair<F>> {
        self.twiddles.get_or_compute(log_h, || {
            let half_h = (1 << log_h) >> 1;
            let root = F::two_adic_generator(log_h);
            let twiddles = root.powers().collect_n(half_h);
            let mut bitrev_twiddles = twiddles.clone();
            reverse_slice_index_bits(&mut bitrev_twiddles);

            Arc::new(VectorPair {
                twiddles,
                bitrev_twiddles,
            })
        })
    }

    fn get_or_compute_coset_twiddles(&self, (log_h, shift): (usize, F)) -> Arc<[Vec<F>]> {
        self.coset_twiddles.get_or_compute((log_h, shift), || {
            let mid = log_h.div_ceil(2);
            let h = 1 << log_h;
            let root = F::two_adic_generator(log_h);
            (0..log_h)
                .map(|layer| {
                    let shift_power = shift.exp_power_of_2(layer);
                    let powers = Powers {
                        base: root.exp_power_of_2(layer),
                        current: shift_power,
                    };
                    let mut twiddles = powers.collect_n(h >> (layer + 1));
                    let layer_rev = log_h - 1 - layer;
                    if layer_rev >= mid {
                        reverse_slice_index_bits(&mut twiddles);
                    }
                    twiddles
                })
                .collect::<Vec<_>>()
                .into()
        })
    }

    fn get_or_compute_inverse_twiddles(&self, log_h: usize) -> Arc<VectorPair<F>> {
        self.inverse_twiddles.get_or_compute(log_h, || {
            let half_h = (1 << log_h) >> 1;
            let root_inv = F::two_adic_generator(log_h).inverse();
            let twiddles = root_inv.powers().collect_n(half_h);
            let mut bitrev_twiddles = twiddles.clone();
            reverse_slice_index_bits(&mut bitrev_twiddles);

            Arc::new(VectorPair {
                twiddles,
                bitrev_twiddles,
            })
        })
    }
}

impl<F: TwoAdicField + Ord> TwoAdicSubgroupDft<F> for Radix2DitParallel<F> {
    type Evaluations = BitReversedMatrixView<RowMajorMatrix<F>>;

    fn dft_batch(&self, mut mat: RowMajorMatrix<F>) -> Self::Evaluations {
        let h = mat.height();
        let log_h = log2_strict_usize(h);

        // Compute twiddle factors, or take memoized ones if already available.
        let twiddles = self.get_or_compute_twiddles(log_h);

        let mid = log_h.div_ceil(2);

        // The first half looks like a normal DIT.
        reverse_matrix_index_bits(&mut mat);
        first_half(&mut mat, mid, &twiddles.twiddles);

        // For the second half, we flip the DIT, working in bit-reversed order.
        reverse_matrix_index_bits(&mut mat);
        second_half(&mut mat, mid, &twiddles.bitrev_twiddles, None);

        mat.bit_reverse_rows()
    }

    fn coset_dft_batch(&self, mut mat: RowMajorMatrix<F>, shift: F) -> Self::Evaluations {
        reverse_matrix_index_bits(&mut mat);
        coset_dft(self, &mut mat.as_view_mut(), shift, 0, &|_, _| {});
        BitReversalPerm::new_view(mat)
    }

    fn idft_batch(&self, mut mat: RowMajorMatrix<F>) -> RowMajorMatrix<F> {
        let h = mat.height();
        if h == 1 {
            return mat;
        }
        let log_h = log2_strict_usize(h);
        let mid = log_h.div_ceil(2);
        let inverse_twiddles = self.get_or_compute_inverse_twiddles(log_h);

        reverse_matrix_index_bits(&mut mat);
        first_half(&mut mat, mid, &inverse_twiddles.twiddles);
        reverse_matrix_index_bits(&mut mat);
        let h_inv_subfield = F::PrimeSubfield::ONE.div_2exp_u64(log_h as u64);
        let scale = Some(F::from_prime_subfield(h_inv_subfield));
        second_half(&mut mat, mid, &inverse_twiddles.bitrev_twiddles, scale);
        reverse_matrix_index_bits(&mut mat);
        mat
    }

    fn coset_idft_batch(&self, mat: RowMajorMatrix<F>, shift: F) -> RowMajorMatrix<F> {
        let mut coeffs = self.idft_batch(mat);
        crate::util::coset_shift_cols(&mut coeffs, shift.inverse());
        coeffs
    }

    #[instrument(skip_all, level = "debug", fields(dims = %mat.dimensions(), added_bits = added_bits))]
    fn coset_lde_batch_with_transform<T>(
        &self,
        mat: RowMajorMatrix<F>,
        added_bits: usize,
        shift: F,
        transform: T,
    ) -> Self::Evaluations
    where
        T: FnOnce(&mut RowMajorMatrixViewMut<'_, F>, Layout),
    {
        self.lde_with_consumer(mat, added_bits, shift, transform, &|_, _| {})
    }

    #[instrument(skip_all, level = "debug", fields(dims = %mat.dimensions(), added_bits = added_bits))]
    fn coset_lde_batch_with_blocks<T, K, C>(
        &self,
        mat: RowMajorMatrix<F>,
        added_bits: usize,
        shift: F,
        transform: T,
        make_consumer: K,
    ) -> Self::Evaluations
    where
        T: FnOnce(&mut RowMajorMatrixViewMut<'_, F>, Layout),
        K: FnOnce(usize) -> C,
        C: Fn(usize, RowMajorMatrixView<'_, F>) + Sync,
    {
        // Blocks are the second butterfly half's sub-FFTs; see `second_half_general`.
        let log_h = log2_strict_usize(mat.height());
        let mid = log_h.div_ceil(2);
        let consume = make_consumer(1 << (log_h - mid));
        self.lde_with_consumer(mat, added_bits, shift, transform, &consume)
    }
}

impl<F: TwoAdicField + Ord> Radix2DitParallel<F> {
    fn lde_with_consumer<T>(
        &self,
        mut mat: RowMajorMatrix<F>,
        added_bits: usize,
        shift: F,
        transform: T,
        consume: &(impl Fn(usize, RowMajorMatrixView<'_, F>) + Sync),
    ) -> BitReversedMatrixView<RowMajorMatrix<F>>
    where
        T: FnOnce(&mut RowMajorMatrixViewMut<'_, F>, Layout),
    {
        let w = mat.width;
        let h = mat.height();
        let log_h = log2_strict_usize(h);
        let mid = log_h.div_ceil(2);

        let inverse_twiddles = self.get_or_compute_inverse_twiddles(log_h);

        // The first half looks like a normal DIT.
        reverse_matrix_index_bits(&mut mat);
        first_half(&mut mat, mid, &inverse_twiddles.twiddles);

        // For the second half, we flip the DIT, working in bit-reversed order.
        reverse_matrix_index_bits(&mut mat);
        // We'll also scale by 1/h, as per the usual inverse DFT algorithm.
        // If F isn't a PrimeField, (and is thus an extension field) it's much cheaper to
        // invert in F::PrimeSubfield.
        let h_inv_subfield = F::PrimeSubfield::from_int(h).try_inverse();
        let scale = h_inv_subfield.map(F::from_prime_subfield);
        second_half(&mut mat, mid, &inverse_twiddles.bitrev_twiddles, scale);
        // We skip the final bit-reversal, since the next FFT expects bit-reversed input.

        transform(&mut mat.as_view_mut(), Layout::BitReversed);

        let lde_elems = w * (h << added_bits);
        let elems_to_add = lde_elems - w * h;
        debug_span!("reserve_exact").in_scope(|| mat.values.reserve_exact(elems_to_add));

        let g_big = F::two_adic_generator(log_h + added_bits);

        // Resolve tables before the forward transforms, indexed by the exponent of g_big.
        let coset_twiddles: Vec<_> = (0..(1usize << added_bits))
            .into_par_iter()
            .map(|coset_idx| {
                let total_shift = g_big.exp_u64(coset_idx as u64) * shift;
                self.get_or_compute_coset_twiddles((log_h, total_shift))
            })
            .collect();

        let mat_ptr = mat.values.as_mut_ptr();
        let rest_ptr = unsafe { (mat_ptr as *mut MaybeUninit<F>).add(w * h) };
        let first_slice: &mut [F] = unsafe { slice::from_raw_parts_mut(mat_ptr, w * h) };
        let rest_slice: &mut [MaybeUninit<F>] =
            unsafe { slice::from_raw_parts_mut(rest_ptr, lde_elems - w * h) };
        let mut first_coset_mat = RowMajorMatrixViewMut::new(first_slice, w);
        let mut rest_cosets_mat = rest_slice
            .chunks_exact_mut(w * h)
            .map(|slice| RowMajorMatrixViewMut::new(slice, w))
            .collect_vec();

        // Each task writes a disjoint destination while sharing the coefficient matrix.
        // Physical slot k + 1 holds coset reverse_bits_len(k + 1, added_bits).
        let src = first_coset_mat.as_view();
        rest_cosets_mat
            .par_iter_mut()
            .enumerate()
            .for_each(|(k, dest)| {
                let coset_idx = reverse_bits_len(k + 1, added_bits);
                coset_dft_oop(&src, dest, &coset_twiddles[coset_idx], (k + 1) * h, consume);
            });

        // Join all coefficient readers and consumers before coset zero overwrites the coefficients.
        coset_dft(self, &mut first_coset_mat.as_view_mut(), shift, 0, consume);

        // SAFETY: We wrote all values above.
        unsafe {
            mat.values.set_len(lde_elems);
        }
        BitReversalPerm::new_view(mat)
    }
}

#[instrument(level = "debug", skip_all)]
fn coset_dft<F: TwoAdicField + Ord>(
    dft: &Radix2DitParallel<F>,
    mat: &mut RowMajorMatrixViewMut<'_, F>,
    shift: F,
    row_base: usize,
    consume: &(impl Fn(usize, RowMajorMatrixView<'_, F>) + Sync),
) {
    let log_h = log2_strict_usize(mat.height());
    let mid = log_h.div_ceil(2);

    let twiddles = dft.get_or_compute_coset_twiddles((log_h, shift));

    // The first half looks like a normal DIT.
    first_half_general(mat, mid, &twiddles);

    // For the second half, we flip the DIT, working in bit-reversed order.
    reverse_matrix_index_bits(mat);

    second_half_general(mat, mid, &twiddles, row_base, consume);
}

/// Like `coset_dft`, except out-of-place and using precomputed twiddles.
///
/// Tables must use the layer layout of `get_or_compute_coset_twiddles` for the
/// source height and intended shift. Each layer contains `height >> (layer + 1)` entries.
///
/// # Panics
/// Panics if the matrix dimensions or twiddle table dimensions do not match,
/// or if the height is not a positive power of two.
#[instrument(level = "debug", skip_all)]
fn coset_dft_oop<F: Field>(
    src: &RowMajorMatrixView<'_, F>,
    dst_maybe: &mut RowMajorMatrixViewMut<'_, MaybeUninit<F>>,
    twiddles: &[Vec<F>],
    row_base: usize,
    consume: &(impl Fn(usize, RowMajorMatrixView<'_, F>) + Sync),
) {
    assert_eq!(src.dimensions(), dst_maybe.dimensions());

    let log_h = log2_strict_usize(dst_maybe.height());
    // Short tables can leave destination rows uninitialized before the casts below.
    assert_eq!(twiddles.len(), log_h, "incorrect number of twiddle layers");
    for (layer, table) in twiddles.iter().enumerate() {
        assert_eq!(
            table.len(),
            src.height() >> (layer + 1),
            "incorrect twiddle count for layer {layer}"
        );
    }

    let mid = log_h.div_ceil(2);
    if log_h == 0 {
        // This is an edge case where first_half_general_oop doesn't work, as it expects there to be
        // at least one layer in the network, so we just copy instead.
        let src_maybe = unsafe {
            transmute::<&RowMajorMatrixView<'_, F>, &RowMajorMatrixView<'_, MaybeUninit<F>>>(src)
        };
        dst_maybe.copy_from(src_maybe);
    } else {
        // The first half looks like a normal DIT.
        first_half_general_oop(src, dst_maybe, mid, twiddles);
    }

    // SAFETY: The copy or first FFT half initialized every destination element.
    let dst = unsafe {
        transmute::<&mut RowMajorMatrixViewMut<'_, MaybeUninit<F>>, &mut RowMajorMatrixViewMut<'_, F>>(
            dst_maybe,
        )
    };

    if log_h == 0 {
        consume(row_base, dst.as_view());
        return;
    }

    // For the second half, we flip the DIT, working in bit-reversed order.
    reverse_matrix_index_bits(dst);

    second_half_general(dst, mid, twiddles, row_base, consume);
}

/// This can be used as the first half of a DIT butterfly network.
///
/// For layer 0, all twiddle factors are 1 (root^0 = 1), so we use `TwiddleFreeButterfly`
/// to avoid a Montgomery multiply by 1 across the entire matrix.
///
/// For layers 1 to mid-1 included, the first twiddle in each block is also always 1 (`twiddles[0] = 1`),
/// so we special-case the first row-pair of each block to use `TwiddleFreeButterfly` as well.
#[instrument(level = "debug", skip_all)]
fn first_half<F: Field>(mat: &mut RowMajorMatrix<F>, mid: usize, twiddles: &[F]) {
    let log_h = log2_strict_usize(mat.height());

    // max block size: 2^mid
    mat.par_row_chunks_exact_mut(1 << mid)
        .for_each(|mut submat| {
            let mut backwards = false;
            for layer in 0..mid {
                if layer == 0 {
                    // For layer 0, half_block_size=1 and each block clones the twiddle
                    // iterator from the start, consuming only twiddles[0] = root^0 = 1.
                    // Use TwiddleFreeButterfly to skip the multiply entirely.
                    dit_layer_twiddle_free(&mut submat, backwards);
                } else {
                    let layer_rev = log_h - 1 - layer;
                    let layer_pow = 1 << layer_rev;
                    // For layers 1..mid-1, twiddles[0] = root^0 = 1 is always the first
                    // twiddle consumed per block. Use the optimized version that applies
                    // TwiddleFreeButterfly for the first row-pair of each block.
                    dit_layer_first_one(
                        &mut submat,
                        layer,
                        twiddles.iter().step_by(layer_pow),
                        backwards,
                    );
                }
                backwards = !backwards;
            }
        });
}

/// Like `first_half`, except supporting different twiddle factors per layer, enabling coset shifts
/// to be baked into them.
#[instrument(level = "debug", skip_all)]
fn first_half_general<F: Field>(
    mat: &mut RowMajorMatrixViewMut<'_, F>,
    mid: usize,
    twiddles: &[Vec<F>],
) {
    let log_h = log2_strict_usize(mat.height());
    mat.par_row_chunks_exact_mut(1 << mid)
        .for_each(|mut submat| {
            let mut backwards = false;
            for layer in 0..mid {
                let layer_rev = log_h - 1 - layer;
                dit_layer(&mut submat, layer, twiddles[layer_rev].iter(), backwards);
                backwards = !backwards;
            }
        });
}

/// Like `first_half_general`, except out-of-place.
///
/// Assumes there's at least one layer in the network, i.e. `src.height() > 1`.
///
/// # Panics
/// Panics (via `log2_strict_usize` and arithmetic underflow) if `src.height() < 2`.
#[instrument(level = "debug", skip_all)]
fn first_half_general_oop<F: Field>(
    src: &RowMajorMatrixView<'_, F>,
    dst_maybe: &mut RowMajorMatrixViewMut<'_, MaybeUninit<F>>,
    mid: usize,
    twiddles: &[Vec<F>],
) {
    let log_h = log2_strict_usize(src.height());
    src.par_row_chunks_exact(1 << mid)
        .zip(dst_maybe.par_row_chunks_exact_mut(1 << mid))
        .for_each(|(src_submat, mut dst_submat_maybe)| {
            debug_assert_eq!(src_submat.dimensions(), dst_submat_maybe.dimensions());

            // The first layer is special, done out-of-place.
            // (Recall from the mid definition that there must be at least one layer here.)
            let layer_rev = log_h - 1;
            dit_layer_oop(
                &src_submat,
                &mut dst_submat_maybe,
                0,
                twiddles[layer_rev].iter(),
            );

            // submat is now initialized.
            let mut dst_submat = unsafe {
                transmute::<RowMajorMatrixViewMut<'_, MaybeUninit<F>>, RowMajorMatrixViewMut<'_, F>>(
                    dst_submat_maybe,
                )
            };

            // Subsequent layers.
            let mut backwards = true;
            for layer in 1..mid {
                let layer_rev = log_h - 1 - layer;
                dit_layer(
                    &mut dst_submat,
                    layer,
                    twiddles[layer_rev].iter(),
                    backwards,
                );
                backwards = !backwards;
            }
        });
}

/// This can be used as the second half of a DIT butterfly network. It works in bit-reversed order.
///
/// The optional `scale` parameter is used to scale the matrix by a constant factor. Rather than
/// doing a separate pass over memory, we fold the scaling into the first butterfly layer to
/// eliminate an extra memory pass.
#[instrument(level = "debug", skip_all)]
#[inline(always)] // To avoid branch on scale
fn second_half<F: Field>(
    mat: &mut RowMajorMatrix<F>,
    mid: usize,
    twiddles_rev: &[F],
    scale: Option<F>,
) {
    let log_h = log2_strict_usize(mat.height());

    // max block size: 2^(log_h - mid)
    mat.par_row_chunks_exact_mut(1 << (log_h - mid))
        .enumerate()
        .for_each(|(thread, mut submat)| {
            let mut backwards = false;
            if let Some(scale) = scale {
                // Fold the scale into the first butterfly layer to avoid a separate
                // memory pass. This merges the O(N) scaling step into the first O(N)
                // butterfly pass.
                let mut scale_applied = false;
                for layer in mid..log_h {
                    let first_block = thread << (layer - mid);
                    if !scale_applied {
                        scale_applied = true;
                        dit_layer_rev_scaled(
                            &mut submat,
                            log_h,
                            layer,
                            twiddles_rev[first_block..].iter().copied(),
                            backwards,
                            Some(scale),
                        );
                    } else {
                        dit_layer_rev(
                            &mut submat,
                            log_h,
                            layer,
                            twiddles_rev[first_block..].iter().copied(),
                            backwards,
                        );
                    }
                    backwards = !backwards;
                }
                // Handle case where there are no layers in the second half (mid == log_h).
                if !scale_applied {
                    submat.scale(scale);
                }
            } else {
                for layer in mid..log_h {
                    let first_block = thread << (layer - mid);
                    dit_layer_rev(
                        &mut submat,
                        log_h,
                        layer,
                        twiddles_rev[first_block..].iter().copied(),
                        backwards,
                    );
                    backwards = !backwards;
                }
            }
        });
}

/// Like `second_half`, except supporting different twiddle factors per layer, enabling coset shifts
/// to be baked into them.
#[instrument(level = "debug", skip_all)]
fn second_half_general<F: Field>(
    mat: &mut RowMajorMatrixViewMut<'_, F>,
    mid: usize,
    twiddles_rev: &[Vec<F>],
    row_base: usize,
    consume: &(impl Fn(usize, RowMajorMatrixView<'_, F>) + Sync),
) {
    let log_h = log2_strict_usize(mat.height());
    mat.par_row_chunks_exact_mut(1 << (log_h - mid))
        .enumerate()
        .for_each(|(thread, mut submat)| {
            let mut backwards = false;
            for layer in mid..log_h {
                let layer_rev = log_h - 1 - layer;
                let first_block = thread << (layer - mid);
                dit_layer_rev(
                    &mut submat,
                    log_h,
                    layer,
                    twiddles_rev[layer_rev][first_block..].iter().copied(),
                    backwards,
                );
                backwards = !backwards;
            }
            consume(row_base + thread * submat.height(), submat.as_view());
        });
}

/// One layer of a DIT butterfly network where all twiddle factors are 1 (i.e., layer 0).
///
/// This is equivalent to `dit_layer` with `layer=0` and `twiddles[0]=1`, but uses
/// `TwiddleFreeButterfly` to avoid a Montgomery multiplication by 1 in the hot loop.
///
/// Correctness: For layer=0, `half_block_size=1` and each block clones the twiddle
/// iterator from position 0, consuming only `twiddles[0] = generator^0 = 1`.
/// Since multiplying by 1 is a no-op, `TwiddleFreeButterfly` gives identical results.
fn dit_layer_twiddle_free<F: Field>(submat: &mut RowMajorMatrixViewMut<'_, F>, backwards: bool) {
    // layer=0 means half_block_size=1, block_size=2.
    let width = submat.width();
    debug_assert!(submat.height() >= 2);

    let process_block = move |block: &mut [F]| {
        // Each block is exactly 2 rows: lo = block[0..width], hi = block[width..2*width]
        let (lo, hi) = block.split_at_mut(width);
        TwiddleFreeButterfly.apply_to_rows(lo, hi);
    };

    let blocks = submat.values.chunks_mut(2 * width);
    if backwards {
        for block in blocks.rev() {
            process_block(block);
        }
    } else {
        for block in blocks {
            process_block(block);
        }
    }
}

/// One layer of a DIT butterfly network where the first twiddle factor per block is always 1.
///
/// This is used in `first_half` for layers 1..mid-1 of the standard (non-coset) DFT/inverse DFT,
/// where `twiddles[0] = root^0 = 1`. The first row-pair of each block uses `TwiddleFreeButterfly`
/// to avoid one Montgomery multiplication per block, while subsequent row-pairs use `DitButterfly`.
///
/// Correctness: The twiddle iterator yields `twiddles[0], twiddles[step], twiddles[2*step], ...`
/// where `twiddles[0] = root^0 = 1`. Only used when this property holds.
fn dit_layer_first_one<'a, F: Field>(
    submat: &mut RowMajorMatrixViewMut<'_, F>,
    layer: usize,
    twiddles: impl Iterator<Item = &'a F> + Clone,
    backwards: bool,
) {
    let half_block_size = 1 << layer;
    let block_size = half_block_size * 2;
    let width = submat.width();
    debug_assert!(submat.height() >= block_size);
    debug_assert!(
        half_block_size >= 2,
        "layer must be >= 1 for dit_layer_first_one"
    );

    let process_block = move |block: &mut [F]| {
        let (lows, highs) = block.split_at_mut(half_block_size * width);
        let mut tw_iter = twiddles.clone();
        // First row-pair: twiddle is always 1, use TwiddleFreeButterfly to skip the multiply.
        let _ = tw_iter.next(); // consume twiddles[0] = 1
        let (lo0, lo_rest) = lows.split_at_mut(width);
        let (hi0, hi_rest) = highs.split_at_mut(width);
        TwiddleFreeButterfly.apply_to_rows(lo0, hi0);
        // Remaining row-pairs use DitButterfly with their respective twiddle factors.
        for (lo, hi, twiddle) in izip!(
            lo_rest.chunks_mut(width),
            hi_rest.chunks_mut(width),
            tw_iter
        ) {
            DitButterfly(*twiddle).apply_to_rows(lo, hi);
        }
    };

    let blocks = submat.values.chunks_mut(block_size * width);
    if backwards {
        for block in blocks.rev() {
            process_block(block);
        }
    } else {
        for block in blocks {
            process_block(block);
        }
    }
}

/// One layer of a DIT butterfly network.
fn dit_layer<'a, F: Field>(
    submat: &mut RowMajorMatrixViewMut<'_, F>,
    layer: usize,
    twiddles: impl Iterator<Item = &'a F> + Clone,
    backwards: bool,
) {
    let half_block_size = 1 << layer;
    let block_size = half_block_size * 2;
    let width = submat.width();
    debug_assert!(submat.height() >= block_size);

    let process_block = move |block: &mut [F]| {
        let (lows, highs) = block.split_at_mut(half_block_size * width);
        for (lo, hi, twiddle) in izip!(
            lows.chunks_mut(width),
            highs.chunks_mut(width),
            twiddles.clone()
        ) {
            DitButterfly(*twiddle).apply_to_rows(lo, hi);
        }
    };

    let blocks = submat.values.chunks_mut(block_size * width);
    if backwards {
        for block in blocks.rev() {
            process_block(block);
        }
    } else {
        for block in blocks {
            process_block(block);
        }
    }
}

/// One layer of a DIT butterfly network, out-of-place.
fn dit_layer_oop<'a, F: Field>(
    src: &RowMajorMatrixView<'_, F>,
    dst: &mut RowMajorMatrixViewMut<'_, MaybeUninit<F>>,
    layer: usize,
    twiddles: impl Iterator<Item = &'a F> + Clone,
) {
    debug_assert_eq!(src.dimensions(), dst.dimensions());
    let half_block_size = 1 << layer;
    let block_size = half_block_size * 2;
    let width = dst.width();
    debug_assert!(dst.height() >= block_size);

    let process_blocks = move |src_block: &[F], dst_block: &mut [MaybeUninit<F>]| {
        let (src_lows, src_highs) = src_block.split_at(half_block_size * width);
        let (dst_lows, dst_highs) = dst_block.split_at_mut(half_block_size * width);

        for (src_lo, dst_lo, src_hi, dst_hi, twiddle) in izip!(
            src_lows.chunks(width),
            dst_lows.chunks_mut(width),
            src_highs.chunks(width),
            dst_highs.chunks_mut(width),
            twiddles.clone()
        ) {
            DitButterfly(*twiddle).apply_to_rows_oop(src_lo, dst_lo, src_hi, dst_hi);
        }
    };

    let src_chunks = src.values.chunks(block_size * width);
    let dst_chunks = dst.values.chunks_mut(block_size * width);

    for (src_block, dst_block) in src_chunks.zip(dst_chunks) {
        process_blocks(src_block, dst_block);
    }
}

/// Like `dit_layer_rev`, except with an optional scale factor folded into the butterfly.
///
/// This avoids an extra memory pass when scaling is required (e.g., 1/N in inverse DFT).
/// When `scale` is `None`, this is identical to `dit_layer_rev`.
///
/// When `scale` is `Some(s)`, uses `ScaledDitButterfly::new(twiddle, s)` which precomputes
/// `twiddle * scale` once per block, reducing multiplications in the hot loop from 3 to 2.
fn dit_layer_rev_scaled<F: Field>(
    submat: &mut RowMajorMatrixViewMut<'_, F>,
    log_h: usize,
    layer: usize,
    twiddles_rev: impl DoubleEndedIterator<Item = F> + ExactSizeIterator,
    backwards: bool,
    scale: Option<F>,
) {
    let layer_rev = log_h - 1 - layer;

    let half_block_size = 1 << layer_rev;
    let block_size = half_block_size * 2;
    let width = submat.width();
    debug_assert!(submat.height() >= block_size);

    match scale {
        None => {
            // No scaling: same as regular dit_layer_rev
            let blocks_and_twiddles = submat
                .values
                .chunks_mut(block_size * width)
                .zip(twiddles_rev);
            if backwards {
                for (block, twiddle) in blocks_and_twiddles.rev() {
                    let (lo, hi) = block.split_at_mut(half_block_size * width);
                    DitButterfly(twiddle).apply_to_rows(lo, hi);
                }
            } else {
                for (block, twiddle) in blocks_and_twiddles {
                    let (lo, hi) = block.split_at_mut(half_block_size * width);
                    DitButterfly(twiddle).apply_to_rows(lo, hi);
                }
            }
        }
        Some(s) => {
            // Fold scaling into the butterfly to avoid a separate memory pass.
            // ScaledDitButterfly::new precomputes twiddle * scale once per block,
            // so the hot loop only needs 2 multiplications instead of 3.
            let blocks_and_twiddles = submat
                .values
                .chunks_mut(block_size * width)
                .zip(twiddles_rev);
            if backwards {
                for (block, twiddle) in blocks_and_twiddles.rev() {
                    let (lo, hi) = block.split_at_mut(half_block_size * width);
                    ScaledDitButterfly::new(twiddle, s).apply_to_rows(lo, hi);
                }
            } else {
                for (block, twiddle) in blocks_and_twiddles {
                    let (lo, hi) = block.split_at_mut(half_block_size * width);
                    ScaledDitButterfly::new(twiddle, s).apply_to_rows(lo, hi);
                }
            }
        }
    }
}

/// Like `dit_layer`, except the matrix and twiddles are encoded in bit-reversed order.
/// This can also be viewed as a layer of the Bowers G^T network.
fn dit_layer_rev<F: Field>(
    submat: &mut RowMajorMatrixViewMut<'_, F>,
    log_h: usize,
    layer: usize,
    twiddles_rev: impl DoubleEndedIterator<Item = F> + ExactSizeIterator,
    backwards: bool,
) {
    let layer_rev = log_h - 1 - layer;

    let half_block_size = 1 << layer_rev;
    let block_size = half_block_size * 2;
    let width = submat.width();
    debug_assert!(submat.height() >= block_size);

    let blocks_and_twiddles = submat
        .values
        .chunks_mut(block_size * width)
        .zip(twiddles_rev);
    if backwards {
        for (block, twiddle) in blocks_and_twiddles.rev() {
            let (lo, hi) = block.split_at_mut(half_block_size * width);
            DitButterfly(twiddle).apply_to_rows(lo, hi);
        }
    } else {
        for (block, twiddle) in blocks_and_twiddles {
            let (lo, hi) = block.split_at_mut(half_block_size * width);
            DitButterfly(twiddle).apply_to_rows(lo, hi);
        }
    }
}

#[cfg(test)]
mod tests {
    use p3_baby_bear::BabyBear;
    use p3_field::TwoAdicField;
    use p3_matrix::Matrix;
    use p3_matrix::dense::RowMajorMatrix;
    use rand::SeedableRng;
    use rand::rngs::SmallRng;

    use super::*;

    type F = BabyBear;

    #[test]
    #[should_panic(expected = "incorrect number of twiddle layers")]
    fn coset_dft_oop_rejects_missing_twiddle_layer() {
        let dft = Radix2DitParallel::<F>::default();
        let mut twiddles = dft
            .get_or_compute_coset_twiddles((3, F::GENERATOR))
            .to_vec();
        twiddles.pop();
        let src = RowMajorMatrix::new(alloc::vec![F::ONE; 8], 1);
        let mut dst = RowMajorMatrix::new(alloc::vec![MaybeUninit::uninit(); 8], 1);

        coset_dft_oop(
            &src.as_view(),
            &mut dst.as_view_mut(),
            &twiddles,
            0,
            &|_, _| {},
        );
    }

    #[test]
    #[should_panic(expected = "incorrect twiddle count for layer 2")]
    fn coset_dft_oop_rejects_short_twiddle_layer() {
        let dft = Radix2DitParallel::<F>::default();
        let mut twiddles = dft
            .get_or_compute_coset_twiddles((3, F::GENERATOR))
            .to_vec();
        twiddles[2].clear();
        let src = RowMajorMatrix::new(alloc::vec![F::ONE; 8], 1);
        let mut dst = RowMajorMatrix::new(alloc::vec![MaybeUninit::uninit(); 8], 1);

        coset_dft_oop(
            &src.as_view(),
            &mut dst.as_view_mut(),
            &twiddles,
            0,
            &|_, _| {},
        );
    }

    #[test]
    fn coset_dft_idft_roundtrip() {
        let dft = Radix2DitParallel::<F>::default();
        let shift = F::GENERATOR;
        let mut rng = SmallRng::seed_from_u64(42);
        let original = RowMajorMatrix::<F>::rand(&mut rng, 16, 3);

        let evals = dft.coset_dft_batch(original.clone(), shift);
        let recovered = dft.coset_idft_batch(evals.to_row_major_matrix(), shift);

        assert_eq!(original, recovered);
    }

    #[test]
    fn coset_dft_matches_default_trait() {
        let dft = Radix2DitParallel::<F>::default();
        let shift = F::two_adic_generator(4) * F::GENERATOR;
        let mut rng = SmallRng::seed_from_u64(7);
        let mat = RowMajorMatrix::<F>::rand(&mut rng, 16, 4);

        let override_result = dft
            .coset_dft_batch(mat.clone(), shift)
            .to_row_major_matrix();

        let mut shifted = mat;
        crate::util::coset_shift_cols(&mut shifted, shift);
        let default_result = dft.dft_batch(shifted).to_row_major_matrix();

        assert_eq!(override_result, default_result);
    }
}