1use crate::{CVWords, IncrementCounter, BLOCK_LEN, OUT_LEN};
2
3// Unsafe because this may only be called on platforms supporting AVX-512.
4pub unsafe fn compress_in_place(
5 cv: &mut CVWords,
6 block: &[u8; BLOCK_LEN],
7 block_len: u8,
8 counter: u64,
9 flags: u8,
10) {
11 ffi::blake3_compress_in_place_avx512(cv.as_mut_ptr(), block.as_ptr(), block_len, counter, flags)
12}
13
14// Unsafe because this may only be called on platforms supporting AVX-512.
15pub unsafe fn compress_xof(
16 cv: &CVWords,
17 block: &[u8; BLOCK_LEN],
18 block_len: u8,
19 counter: u64,
20 flags: u8,
21) -> [u8; 64] {
22 let mut out: [u8; 64] = [0u8; 64];
23 ffi::blake3_compress_xof_avx512(
24 cv.as_ptr(),
25 block.as_ptr(),
26 block_len,
27 counter,
28 flags,
29 out.as_mut_ptr(),
30 );
31 out
32}
33
34// Unsafe because this may only be called on platforms supporting AVX-512.
35pub unsafe fn hash_many<const N: usize>(
36 inputs: &[&[u8; N]],
37 key: &CVWords,
38 counter: u64,
39 increment_counter: IncrementCounter,
40 flags: u8,
41 flags_start: u8,
42 flags_end: u8,
43 out: &mut [u8],
44) {
45 // The Rust hash_many implementations do bounds checking on the `out`
46 // array, but the C implementations don't. Even though this is an unsafe
47 // function, assert the bounds here.
48 assert!(out.len() >= inputs.len() * OUT_LEN);
49 ffi::blake3_hash_many_avx512(
50 inputs.as_ptr() as *const *const u8,
51 num_inputs:inputs.len(),
52 N / BLOCK_LEN,
53 key.as_ptr(),
54 counter,
55 increment_counter.yes(),
56 flags,
57 flags_start,
58 flags_end,
59 out.as_mut_ptr(),
60 )
61}
62
63// Unsafe because this may only be called on platforms supporting AVX-512.
64#[cfg(unix)]
65pub unsafe fn xof_many(
66 cv: &CVWords,
67 block: &[u8; BLOCK_LEN],
68 block_len: u8,
69 counter: u64,
70 flags: u8,
71 out: &mut [u8],
72) {
73 debug_assert_eq!(0, out.len() % BLOCK_LEN, "whole blocks only");
74 ffi::blake3_xof_many_avx512(
75 cv.as_ptr(),
76 block.as_ptr(),
77 block_len,
78 counter,
79 flags,
80 out.as_mut_ptr(),
81 outblocks:out.len() / BLOCK_LEN,
82 );
83}
84
85pub mod ffi {
86 extern "C" {
87 pub fn blake3_compress_in_place_avx512(
88 cv: *mut u32,
89 block: *const u8,
90 block_len: u8,
91 counter: u64,
92 flags: u8,
93 );
94 pub fn blake3_compress_xof_avx512(
95 cv: *const u32,
96 block: *const u8,
97 block_len: u8,
98 counter: u64,
99 flags: u8,
100 out: *mut u8,
101 );
102 pub fn blake3_hash_many_avx512(
103 inputs: *const *const u8,
104 num_inputs: usize,
105 blocks: usize,
106 key: *const u32,
107 counter: u64,
108 increment_counter: bool,
109 flags: u8,
110 flags_start: u8,
111 flags_end: u8,
112 out: *mut u8,
113 );
114 #[cfg(unix)]
115 pub fn blake3_xof_many_avx512(
116 cv: *const u32,
117 block: *const u8,
118 block_len: u8,
119 counter: u64,
120 flags: u8,
121 out: *mut u8,
122 outblocks: usize,
123 );
124 }
125}
126
127#[cfg(test)]
128mod test {
129 use super::*;
130
131 #[test]
132 fn test_compress() {
133 if !crate::platform::avx512_detected() {
134 return;
135 }
136 crate::test::test_compress_fn(compress_in_place, compress_xof);
137 }
138
139 #[test]
140 fn test_hash_many() {
141 if !crate::platform::avx512_detected() {
142 return;
143 }
144 crate::test::test_hash_many_fn(hash_many, hash_many);
145 }
146
147 #[cfg(unix)]
148 #[test]
149 fn test_xof_many() {
150 if !crate::platform::avx512_detected() {
151 return;
152 }
153 crate::test::test_xof_many_fn(xof_many);
154 }
155}
156