ab_blake3/
single_block.rs

1//! BLAKE3 functions that process at most a single block.
2//!
3//! This module and submodules are copied with modifications from the official [`blake3`] crate, but
4//! are unlikely to be upstreamed.
5
6#[cfg(test)]
7mod tests;
8
9// TODO: Workaround for https://github.com/Rust-GPU/rust-gpu/issues/312
10#[cfg(not(target_arch = "spirv"))]
11use crate::platform::{le_bytes_from_words_32, words_from_le_bytes_32};
12// TODO: Workaround for https://github.com/Rust-GPU/rust-gpu/issues/312
13#[cfg(not(target_arch = "spirv"))]
14use crate::{
15    BLOCK_LEN, BlockBytes, DERIVE_KEY_CONTEXT, DERIVE_KEY_MATERIAL, KEY_LEN, KEYED_HASH, OUT_LEN,
16};
17use crate::{BlockWords, CHUNK_END, CHUNK_START, CVWords, IV, ROOT, portable};
18// TODO: Workaround for https://github.com/Rust-GPU/rust-gpu/issues/312
19#[cfg(not(target_arch = "spirv"))]
20use blake3::IncrementCounter;
21// TODO: Workaround for https://github.com/Rust-GPU/rust-gpu/issues/312
22#[cfg(not(target_arch = "spirv"))]
23use blake3::platform::Platform;
24
25/// Hash single block worth of values
26// TODO: Workaround for https://github.com/Rust-GPU/rust-gpu/issues/312
27#[cfg(not(target_arch = "spirv"))]
28#[inline(always)]
29fn hash_block(input: &[u8], key: CVWords, flags: u8) -> Option<[u8; OUT_LEN]> {
30    // If the whole subtree is one block, hash it directly with a ChunkState.
31    if input.len() > BLOCK_LEN {
32        return None;
33    }
34
35    let mut cv = key;
36
37    let mut block = [0; BLOCK_LEN];
38    block[..input.len()].copy_from_slice(input);
39    Platform::detect().compress_in_place(
40        &mut cv,
41        &block,
42        input.len() as u8,
43        0,
44        flags | CHUNK_START | CHUNK_END | ROOT,
45    );
46
47    Some(*le_bytes_from_words_32(&cv))
48}
49
50/// Hash multiple single block-sized inputs
51// TODO: Workaround for https://github.com/Rust-GPU/rust-gpu/issues/312
52#[cfg(not(target_arch = "spirv"))]
53#[inline(always)]
54fn hash_block_many_exact<const NUM_BLOCKS: usize>(
55    inputs: &[BlockBytes; NUM_BLOCKS],
56    // TODO: `&mut [MaybeUninit<[u8; OUT_LEN]>; N]` would make more sense, but doesn't match
57    //  `blake3` API
58    outputs: &mut [[u8; OUT_LEN]; NUM_BLOCKS],
59    key: CVWords,
60    flags: u8,
61) {
62    let platform = Platform::detect();
63
64    let (input_chunks, remaining_inputs) = inputs.as_chunks::<16>();
65    let (output_chunks, remaining_output_chunks) = outputs.as_chunks_mut::<16>();
66
67    for (inputs, outputs) in input_chunks.iter().zip(output_chunks) {
68        // TODO: This is a very awkward API, ideally we wouldn't have this array allocated inline
69        //  for no good reason
70        platform.hash_many(
71            &inputs.each_ref(),
72            &key,
73            0,
74            IncrementCounter::No,
75            flags | CHUNK_START | CHUNK_END | ROOT,
76            0,
77            0,
78            outputs.as_flattened_mut(),
79        );
80    }
81
82    for (input, output) in remaining_inputs.iter().zip(remaining_output_chunks) {
83        let mut cv = key;
84
85        platform.compress_in_place(
86            &mut cv,
87            input,
88            BLOCK_LEN as u8,
89            0,
90            flags | CHUNK_START | CHUNK_END | ROOT,
91        );
92
93        output.copy_from_slice(le_bytes_from_words_32(&cv))
94    }
95}
96
97/// Hashing function for at most single block worth of bytes.
98///
99/// Returns `None` if the input length exceeds one block.
100// TODO: Workaround for https://github.com/Rust-GPU/rust-gpu/issues/312
101#[cfg(not(target_arch = "spirv"))]
102#[inline]
103#[cfg_attr(feature = "no-panic", no_panic::no_panic)]
104pub fn single_block_hash(input: &[u8]) -> Option<[u8; OUT_LEN]> {
105    hash_block(input, *IV, 0)
106}
107
108/// Hashing function for many single-block inputs
109// TODO: Workaround for https://github.com/Rust-GPU/rust-gpu/issues/312
110#[cfg(not(target_arch = "spirv"))]
111#[inline]
112#[cfg_attr(feature = "no-panic", no_panic::no_panic)]
113pub fn single_block_hash_many_exact<const NUM_BLOCKS: usize>(
114    inputs: &[BlockBytes; NUM_BLOCKS],
115    // TODO: `&mut [MaybeUninit<[u8; OUT_LEN]>; N]` would make more sense, but doesn't match
116    //  `blake3` API
117    outputs: &mut [[u8; OUT_LEN]; NUM_BLOCKS],
118) {
119    hash_block_many_exact(inputs, outputs, *IV, 0)
120}
121
122/// The keyed hash function for at most single block worth of bytes.
123///
124/// Returns `None` if the input length exceeds one block.
125// TODO: Workaround for https://github.com/Rust-GPU/rust-gpu/issues/312
126#[cfg(not(target_arch = "spirv"))]
127#[inline]
128#[cfg_attr(feature = "no-panic", no_panic::no_panic)]
129pub fn single_block_keyed_hash(key: &[u8; KEY_LEN], input: &[u8]) -> Option<[u8; OUT_LEN]> {
130    let key_words = words_from_le_bytes_32(key);
131    hash_block(input, key_words, KEYED_HASH)
132}
133
134/// Keyed hash function for many single-block inputs
135// TODO: Workaround for https://github.com/Rust-GPU/rust-gpu/issues/312
136#[cfg(not(target_arch = "spirv"))]
137#[inline]
138#[cfg_attr(feature = "no-panic", no_panic::no_panic)]
139pub fn single_block_keyed_hash_many_exact<const NUM_BLOCKS: usize>(
140    key: &[u8; KEY_LEN],
141    inputs: &[BlockBytes; NUM_BLOCKS],
142    // TODO: `&mut [MaybeUninit<[u8; OUT_LEN]>; N]` would make more sense, but doesn't match
143    //  `blake3` API
144    outputs: &mut [[u8; OUT_LEN]; NUM_BLOCKS],
145) {
146    let key_words = words_from_le_bytes_32(key);
147    hash_block_many_exact(inputs, outputs, key_words, KEYED_HASH)
148}
149
150// The key derivation function for at most a single block worth of bytes.
151//
152// Returns `None` if either context or key material length exceed one block.
153// TODO: Workaround for https://github.com/Rust-GPU/rust-gpu/issues/312
154#[cfg(not(target_arch = "spirv"))]
155#[inline]
156#[cfg_attr(feature = "no-panic", no_panic::no_panic)]
157pub fn single_block_derive_key(context: &str, key_material: &[u8]) -> Option<[u8; OUT_LEN]> {
158    let context_key = hash_block(context.as_bytes(), *IV, DERIVE_KEY_CONTEXT)?;
159    let context_key_words = words_from_le_bytes_32(&context_key);
160    hash_block(key_material, context_key_words, DERIVE_KEY_MATERIAL)
161}
162
163/// Hashing function for at most single block worth of words using portable implementation.
164///
165/// This API operates on words and is GPU-friendly.
166///
167/// `num_bytes` specifies how many actual bytes are occupied by useful value in `input`. Bytes
168/// outside that must be set to `0`.
169///
170/// NOTE: If unused bytes are not set to `0` or an invalid number of bytes is specified, it'll
171/// simply result in an invalid hash.
172///
173/// [`words_from_le_bytes_32()`], [`words_from_le_bytes_64()`] and [`le_bytes_from_words_32()`] can
174/// be used to convert bytes to words and back if necessary.
175///
176/// [`words_from_le_bytes_64()`]: crate::words_from_le_bytes_64
177#[inline]
178#[cfg_attr(feature = "no-panic", no_panic::no_panic)]
179pub fn single_block_hash_portable_words(input: &BlockWords, num_bytes: u32) -> CVWords {
180    let mut cv = *IV;
181
182    portable::compress_in_place_u32(
183        &mut cv,
184        input,
185        num_bytes,
186        0,
187        (CHUNK_START | CHUNK_END | ROOT) as u32,
188    );
189
190    cv
191}
ab_blake3/single_block.rs

ab_blake3/
single_block.rs
