rscrypto 0.8.0

use core::mem::MaybeUninit;

use super::{
  BLOCK_LEN, Blake3, CHUNK_LEN, CV_STACK_LEN, ChunkState, CvBytes, OUT_LEN, OutputState, control,
  hash_full_chunks_cvs_scoped, hash_one_full_chunk_cv, hash_power_of_two_subtree_roots_parallel_rayon, join, kernels,
  pow2_floor, reduce_power_of_two_chunk_cvs_any, single_chunk_output, words8_to_le_bytes, words16_from_le_bytes_64,
};
use crate::{hashes::crypto::blake3::kernels::Kernel, traits::ct};

struct ParallelBatchSpec<'a> {
  input: &'a [u8],
  base_counter: u64,
  batch_chunks: usize,
  commit_chunks: usize,
  threads: usize,
  keep_last_full_chunk: bool,
}

#[derive(Clone, Default)]
pub(super) struct ParallelBatchScratch {
  roots: alloc::vec::Vec<[u32; 8]>,
  leaf_cvs: alloc::vec::Vec<[u32; 8]>,
  zeroize: bool,
}

impl ParallelBatchScratch {
  #[inline]
  fn clear(&mut self) {
    if self.zeroize {
      ct::zeroize_words_no_fence(self.roots.as_flattened_mut());
      ct::zeroize_words_no_fence(self.leaf_cvs.as_flattened_mut());
      core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);
    }
    self.roots.clear();
    self.leaf_cvs.clear();
    self.zeroize = false;
  }
}

impl Drop for ParallelBatchScratch {
  fn drop(&mut self) {
    self.clear();
  }
}

#[cfg(feature = "diag")]
#[doc(hidden)]
#[unsafe(no_mangle)]
#[inline(never)]
pub fn diag_zeroize_blake3_parallel_scratch(input: [u32; 8]) -> u32 {
  let mut scratch = ParallelBatchScratch::default();
  scratch.zeroize = true;
  scratch.roots.push(input);
  core::hint::black_box(scratch.roots[0][0])
}

impl Blake3 {
  #[inline]
  fn streaming_parallel_threads(
    &self,
    input_bytes: usize,
    admission_full_chunks: usize,
    commit_full_chunks: usize,
  ) -> Option<usize> {
    control::streaming_parallel_threads_for_flags(
      self.chunk_state.flags,
      input_bytes,
      admission_full_chunks,
      commit_full_chunks,
    )
  }

  #[cold]
  #[inline(never)]
  fn commit_parallel_batch(&mut self, batch: ParallelBatchSpec<'_>, scratch: &mut ParallelBatchScratch) {
    #[inline]
    fn push_stack(stack: &mut [MaybeUninit<[u32; 8]>; CV_STACK_LEN], len: &mut usize, mut cv: [u32; 8], zeroize: bool) {
      stack[*len].write(cv);
      *len += 1;
      if zeroize {
        ct::zeroize_words(&mut cv);
      }
    }

    #[inline]
    fn pop_stack(stack: &mut [MaybeUninit<[u32; 8]>; CV_STACK_LEN], len: &mut usize, zeroize: bool) -> [u32; 8] {
      *len -= 1;
      // SAFETY: `len` tracks the number of initialized entries.
      let cv = unsafe { stack[*len].assume_init_read() };
      if zeroize {
        // SAFETY: the slot held the `Copy` CV read above and remains
        // initialized until the next push overwrites it.
        ct::zeroize_words(unsafe { stack[*len].assume_init_mut() });
      }
      cv
    }

    let batch_input = batch.input;
    let base_counter = batch.base_counter;
    let commit = batch.commit_chunks;
    let threads = batch.threads;
    let zeroize = self.chunk_state.flags & (super::KEYED_HASH | super::DERIVE_KEY_MATERIAL) != 0;

    let mut stack_len = self.cv_stack_len as usize;
    let mut counter = base_counter;
    let mut offset_chunks = 0usize;
    let mut remaining_commit = commit;

    scratch.clear();
    scratch.zeroize = zeroize;

    while remaining_commit != 0 {
      let mut size = pow2_floor(remaining_commit);

      let aligned_max = if counter == 0 {
        usize::MAX
      } else {
        let tz = counter.trailing_zeros() as usize;
        if tz >= (usize::BITS as usize) {
          usize::MAX
        } else {
          1usize << tz
        }
      };
      size = size.min(aligned_max).min(remaining_commit);
      debug_assert!(size.is_power_of_two());

      let bytes_base = offset_chunks * CHUNK_LEN;
      let subtree_bytes = size * CHUNK_LEN;
      let subtree_input = &batch_input[bytes_base..bytes_base + subtree_bytes];
      let bulk_kernel = kernels::kernel(self.bulk_kernel_id);

      let mut subtree_cv = if size >= threads && (counter == 0 || (counter & (size as u64 - 1)) == 0) {
        const MAX_SUBTREE_CHUNKS: usize = 1 << 12;

        let mut subtree_chunks = size / threads;
        subtree_chunks = subtree_chunks.max(1);
        subtree_chunks = pow2_floor(subtree_chunks);
        subtree_chunks = subtree_chunks.min(MAX_SUBTREE_CHUNKS).min(size);

        let roots_len = size / subtree_chunks;
        debug_assert!(roots_len.is_power_of_two());
        debug_assert_eq!(roots_len * subtree_chunks, size);

        scratch.roots.resize(roots_len, [0u32; 8]);
        hash_power_of_two_subtree_roots_parallel_rayon(super::SubtreeRootsRequest {
          kernel: bulk_kernel,
          key_words: self.key_words,
          flags: self.chunk_state.flags,
          base_counter: counter,
          input: subtree_input,
          subtree_chunks,
          out: &mut scratch.roots,
          threads_total: threads,
        });

        reduce_power_of_two_chunk_cvs_any(
          bulk_kernel,
          self.key_words,
          self.chunk_state.flags,
          &scratch.roots,
          threads,
        )
      } else {
        scratch.leaf_cvs.resize(size, [0u32; 8]);
        hash_full_chunks_cvs_scoped(
          bulk_kernel,
          self.key_words,
          self.chunk_state.flags,
          counter,
          subtree_input,
          &mut scratch.leaf_cvs,
          threads,
        );
        reduce_power_of_two_chunk_cvs_any(
          bulk_kernel,
          self.key_words,
          self.chunk_state.flags,
          &scratch.leaf_cvs,
          threads,
        )
      };

      counter = counter.wrapping_add(size as u64);
      let level = size.trailing_zeros();
      let mut total = counter >> level;
      let mut cv = subtree_cv;
      while total & 1 == 0 {
        let mut left = pop_stack(&mut self.cv_stack, &mut stack_len, zeroize);
        let mut next_cv =
          kernels::parent_cv_inline(self.bulk_kernel_id, left, cv, self.key_words, self.chunk_state.flags);
        if zeroize {
          ct::zeroize_words_no_fence(&mut left);
          ct::zeroize_words_no_fence(&mut cv);
          core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);
        }
        core::mem::swap(&mut cv, &mut next_cv);
        if zeroize {
          ct::zeroize_words(&mut next_cv);
        }
        total >>= 1;
      }
      push_stack(&mut self.cv_stack, &mut stack_len, cv, zeroize);
      if zeroize {
        ct::zeroize_words_no_fence(&mut cv);
        ct::zeroize_words_no_fence(&mut subtree_cv);
        core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);
      }

      offset_chunks = offset_chunks.strict_add(size);
      remaining_commit -= size;
    }

    self.cv_stack_len = stack_len as u8;

    let new_counter = base_counter.strict_add(batch.batch_chunks as u64);
    self.chunk_state = ChunkState::new(
      self.key_words,
      new_counter,
      self.chunk_state.flags,
      self.chunk_state.kernel_id,
    );
    if batch.keep_last_full_chunk {
      let last_chunk_idx = batch.batch_chunks.strict_sub(1);
      let last_start = last_chunk_idx.strict_mul(CHUNK_LEN);
      let last_end = batch.batch_chunks.strict_mul(CHUNK_LEN);
      let last = &batch_input[last_start..last_end];
      self.pending_chunk_cv = Some(hash_one_full_chunk_cv(
        kernels::kernel(self.bulk_kernel_id),
        self.key_words,
        self.chunk_state.flags,
        new_counter.strict_sub(1),
        last,
      ));
      self.pending_cv_chunks = 1;
    }
  }

  #[cold]
  #[inline(never)]
  pub(super) fn try_parallel_update_batch(&mut self, input: &[u8]) -> Option<usize> {
    const MAX_PASS_CHUNKS: usize = 1 << 16;

    if self.chunk_state.len() != 0 || input.len() <= CHUNK_LEN {
      return None;
    }
    let full_chunks = input.len() / CHUNK_LEN;
    if full_chunks <= 1 {
      return None;
    }

    let batch = core::cmp::min(full_chunks, MAX_PASS_CHUNKS);
    let base_counter = self.chunk_state.chunk_counter;

    let keep_last_full_chunk = input.len().is_multiple_of(CHUNK_LEN) && batch == full_chunks;
    let commit = if keep_last_full_chunk { batch - 1 } else { batch };
    if commit == 0 {
      return None;
    }

    let bytes = batch * CHUNK_LEN;
    let threads = self.streaming_parallel_threads(bytes, batch, commit)?;
    if threads <= 1 {
      return None;
    }

    let mut scratch = core::mem::take(&mut self.parallel_batch_scratch);
    self.commit_parallel_batch(
      ParallelBatchSpec {
        input: &input[..bytes],
        base_counter,
        batch_chunks: batch,
        commit_chunks: commit,
        threads,
        keep_last_full_chunk,
      },
      &mut scratch,
    );
    self.parallel_batch_scratch = scratch;
    Some(bytes)
  }
}

#[inline]
fn left_subtree_len_bytes(input_len: usize) -> usize {
  debug_assert!(input_len > CHUNK_LEN);
  let full_chunks = input_len / CHUNK_LEN;
  let has_partial = !input_len.is_multiple_of(CHUNK_LEN);
  let total_chunks = full_chunks.strict_add(usize::from(has_partial));
  debug_assert!(total_chunks >= 2);

  let left_chunks = if total_chunks.is_power_of_two() {
    total_chunks / 2
  } else {
    pow2_floor(total_chunks)
  };
  debug_assert!(left_chunks >= 1 && left_chunks < total_chunks);
  left_chunks.strict_mul(CHUNK_LEN)
}

#[inline]
fn compress_parents_parallel_bytes(
  kernel: Kernel,
  child_cvs: &[CvBytes],
  key_words: [u32; 8],
  flags: u32,
  out: &mut [CvBytes],
) -> usize {
  debug_assert!(child_cvs.len() >= 2);
  debug_assert!(out.len() >= child_cvs.len().div_ceil(2));

  let pairs = child_cvs.len() / 2;
  kernels::parent_cvs_many_from_bytes_inline(kernel.id, &child_cvs[..pairs * 2], key_words, flags, &mut out[..pairs]);
  if (child_cvs.len() & 1) == 1 {
    out[pairs] = child_cvs[child_cvs.len() - 1];
    pairs.strict_add(1)
  } else {
    pairs
  }
}

#[inline]
fn compress_subtree_wide_bytes<J: join::Join>(
  kernel: Kernel,
  key_words: [u32; 8],
  chunk_counter: u64,
  flags: u32,
  input: &[u8],
  out: &mut [CvBytes],
  par_budget: usize,
) -> usize {
  debug_assert!(!input.is_empty());

  let simd_degree = kernel.id.simd_degree();
  let max_leaf_bytes = simd_degree.strict_mul(CHUNK_LEN);
  if input.len() <= max_leaf_bytes {
    let chunks_exact = input.chunks_exact(CHUNK_LEN);
    let full_chunks = chunks_exact.len();
    debug_assert!(full_chunks <= simd_degree);
    debug_assert!(out.len() >= full_chunks.strict_add(usize::from(!chunks_exact.remainder().is_empty())));

    if full_chunks != 0 {
      // SAFETY: `input` has at least `full_chunks * CHUNK_LEN` bytes and
      // `out` has at least `full_chunks * OUT_LEN` bytes.
      unsafe {
        (kernel.hash_many_contiguous)(
          input.as_ptr(),
          full_chunks,
          &key_words,
          chunk_counter,
          flags,
          out.as_mut_ptr().cast::<u8>(),
        );
      }
    }

    let mut out_len = full_chunks;
    let rem = chunks_exact.remainder();
    if !rem.is_empty() {
      let mut cv_words = single_chunk_output(
        kernel,
        key_words,
        chunk_counter.strict_add(full_chunks as u64),
        flags,
        rem,
      )
      .chaining_value();
      out[out_len] = words8_to_le_bytes(&cv_words);
      out_len = out_len.strict_add(1);
      if flags & (super::KEYED_HASH | super::DERIVE_KEY_MATERIAL) != 0 {
        ct::zeroize_words(&mut cv_words);
      }
    }

    return out_len;
  }

  debug_assert!(out.len() >= simd_degree.max(2));
  let left_len = left_subtree_len_bytes(input.len());
  let (left, right) = input.split_at(left_len);
  let right_chunk_counter = chunk_counter.strict_add((left.len() / CHUNK_LEN) as u64);

  const MAX_SIMD_DEGREE: usize = 16;
  let mut cv_array = [[0u8; OUT_LEN]; 2 * MAX_SIMD_DEGREE];

  let simd = simd_degree;
  let degree = if simd == 1 && left.len() == CHUNK_LEN {
    1
  } else {
    simd.max(2)
  };
  let (left_out, right_out) = cv_array.split_at_mut(degree);

  let next_budget = par_budget.saturating_sub(1);
  let (left_n, right_n) = if par_budget == 0 {
    (
      compress_subtree_wide_bytes::<join::SerialJoin>(kernel, key_words, chunk_counter, flags, left, left_out, 0),
      compress_subtree_wide_bytes::<join::SerialJoin>(
        kernel,
        key_words,
        right_chunk_counter,
        flags,
        right,
        right_out,
        0,
      ),
    )
  } else {
    J::join(
      || compress_subtree_wide_bytes::<J>(kernel, key_words, chunk_counter, flags, left, left_out, next_budget),
      || {
        compress_subtree_wide_bytes::<J>(
          kernel,
          key_words,
          right_chunk_counter,
          flags,
          right,
          right_out,
          next_budget,
        )
      },
    )
  };

  debug_assert_eq!(left_n, degree);
  debug_assert!(right_n >= 1 && right_n <= left_n);

  let out_len = if left_n == 1 {
    out[0] = left_out[0];
    out[1] = right_out[0];
    2
  } else {
    let num_children = left_n.strict_add(right_n);
    compress_parents_parallel_bytes(kernel, &cv_array[..num_children], key_words, flags, out)
  };
  if flags & (super::KEYED_HASH | super::DERIVE_KEY_MATERIAL) != 0 {
    ct::zeroize(cv_array.as_flattened_mut());
  }
  out_len
}

#[inline]
fn compress_subtree_to_parent_node_bytes<J: join::Join>(
  kernel: Kernel,
  key_words: [u32; 8],
  chunk_counter: u64,
  flags: u32,
  input: &[u8],
  par_budget: usize,
) -> [u8; BLOCK_LEN] {
  debug_assert!(input.len() > CHUNK_LEN);

  const MAX_SIMD_DEGREE: usize = 16;
  let mut cv_array = [[0u8; OUT_LEN]; MAX_SIMD_DEGREE];
  let mut num_cvs = compress_subtree_wide_bytes::<J>(
    kernel,
    key_words,
    chunk_counter,
    flags,
    input,
    &mut cv_array,
    par_budget,
  );
  debug_assert!(num_cvs >= 2);

  let mut out_array = [[0u8; OUT_LEN]; MAX_SIMD_DEGREE / 2];
  while num_cvs > 2 {
    let cv_slice = &cv_array[..num_cvs];
    let new_n = compress_parents_parallel_bytes(kernel, cv_slice, key_words, flags, &mut out_array);
    cv_array[..new_n].copy_from_slice(&out_array[..new_n]);
    num_cvs = new_n;
  }

  let mut out = [0u8; BLOCK_LEN];
  out[..OUT_LEN].copy_from_slice(&cv_array[0]);
  out[OUT_LEN..].copy_from_slice(&cv_array[1]);
  if flags & (super::KEYED_HASH | super::DERIVE_KEY_MATERIAL) != 0 {
    ct::zeroize_no_fence(cv_array.as_flattened_mut());
    ct::zeroize_no_fence(out_array.as_flattened_mut());
    core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);
  }
  out
}

#[inline]
pub(super) fn root_output_oneshot_join_parallel(
  kernel: Kernel,
  mut key_words: [u32; 8],
  flags: u32,
  input: &[u8],
  threads: usize,
) -> OutputState {
  debug_assert!(input.len() > CHUNK_LEN);
  debug_assert!(threads > 1);

  // Cap Rayon recursion depth to approximately match `threads` leaves.
  let depth = (usize::BITS - 1 - threads.leading_zeros()) as usize;
  let budget = depth.max(1);

  let mut parent_block =
    compress_subtree_to_parent_node_bytes::<join::RayonJoin>(kernel, key_words, 0, flags, input, budget);
  let mut block_words = words16_from_le_bytes_64(&parent_block);
  let output = OutputState {
    kernel_id: kernel.id,
    input_chaining_value: key_words,
    block_words,
    counter: 0,
    block_len: BLOCK_LEN as u32,
    flags: super::PARENT | flags,
  };
  if flags & (super::KEYED_HASH | super::DERIVE_KEY_MATERIAL) != 0 {
    ct::zeroize_no_fence(&mut parent_block);
    ct::zeroize_words_no_fence(&mut block_words);
    ct::zeroize_words_no_fence(&mut key_words);
    core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);
  }
  output
}