Skip to content

Commit 36348fe

Browse files
coyotte508coyotte508
committed
add blake3-wasm and xetchunk-wasm
1 parent 12f9e97 commit 36348fe

File tree

16 files changed

+918
-8
lines changed

16 files changed

+918
-8
lines changed

‎packages/blake3-wasm/assembly/blake3.ts‎

Lines changed: 361 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,361 @@
1+
// Import AssemblyScript types
2+
import type { usize, u32, u8, u64 } from "assemblyscript";
3+
import { StaticArray } from "assemblyscript";
4+
5+
// Constants from the reference implementation
6+
const OUT_LEN: usize = 32;
7+
// const KEY_LEN: usize = 32;
8+
const BLOCK_LEN: usize = 64;
9+
const CHUNK_LEN: usize = 1024;
10+
11+
const CHUNK_START: u32 = 1 << 0;
12+
const CHUNK_END: u32 = 1 << 1;
13+
const PARENT: u32 = 1 << 2;
14+
const ROOT: u32 = 1 << 3;
15+
//const KEYED_HASH: u32 = 1 << 4;
16+
//const DERIVE_KEY_CONTEXT: u32 = 1 << 5;
17+
// const DERIVE_KEY_MATERIAL: u32 = 1 << 6;
18+
19+
const IV: StaticArray<u32> = [
20+
0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19,
21+
];
22+
23+
const MSG_PERMUTATION: StaticArray<usize> = [2, 6, 3, 10, 7, 0, 4, 13, 1, 11, 12, 5, 9, 14, 15, 8];
24+
25+
// The mixing function, G, which mixes either a column or a diagonal.
26+
function g(state: StaticArray<u32>, a: usize, b: usize, c: usize, d: usize, mx: u32, my: u32): void {
27+
state[a] = state[a] + state[b] + mx;
28+
state[d] = rotl32(state[d] ^ state[a], 16);
29+
state[c] = state[c] + state[d];
30+
state[b] = rotl32(state[b] ^ state[c], 12);
31+
state[a] = state[a] + state[b] + my;
32+
state[d] = rotl32(state[d] ^ state[a], 8);
33+
state[c] = state[c] + state[d];
34+
state[b] = rotl32(state[b] ^ state[c], 7);
35+
}
36+
37+
// Rotate left by n bits
38+
function rotl32(x: u32, n: u32): u32 {
39+
return (x << n) | (x >>> (32 - n));
40+
}
41+
42+
function round(state: StaticArray<u32>, m: StaticArray<u32>): void {
43+
// Mix the columns.
44+
g(state, 0, 4, 8, 12, m[0], m[1]);
45+
g(state, 1, 5, 9, 13, m[2], m[3]);
46+
g(state, 2, 6, 10, 14, m[4], m[5]);
47+
g(state, 3, 7, 11, 15, m[6], m[7]);
48+
// Mix the diagonals.
49+
g(state, 0, 5, 10, 15, m[8], m[9]);
50+
g(state, 1, 6, 11, 12, m[10], m[11]);
51+
g(state, 2, 7, 8, 13, m[12], m[13]);
52+
g(state, 3, 4, 9, 14, m[14], m[15]);
53+
}
54+
55+
function permute(m: StaticArray<u32>): void {
56+
const permuted = new StaticArray<u32>(16);
57+
for (let i = 0; i < 16; i++) {
58+
permuted[i] = m[MSG_PERMUTATION[i]];
59+
}
60+
for (let i = 0; i < 16; i++) {
61+
m[i] = permuted[i];
62+
}
63+
}
64+
65+
function compress(
66+
chaining_value: StaticArray<u32>,
67+
block_words: StaticArray<u32>,
68+
counter: u64,
69+
block_len: u32,
70+
flags: u32
71+
): StaticArray<u32> {
72+
const counter_low = counter as u32;
73+
const counter_high = (counter >> 32) as u32;
74+
const state = new StaticArray<u32>(16);
75+
76+
// Initialize state
77+
for (let i = 0; i < 8; i++) {
78+
state[i] = chaining_value[i];
79+
state[i + 8] = IV[i];
80+
}
81+
state[12] = counter_low;
82+
state[13] = counter_high;
83+
state[14] = block_len;
84+
state[15] = flags;
85+
86+
const block = new StaticArray<u32>(16);
87+
for (let i = 0; i < 16; i++) {
88+
block[i] = block_words[i];
89+
}
90+
91+
// Apply rounds
92+
round(state, block);
93+
permute(block);
94+
round(state, block);
95+
permute(block);
96+
round(state, block);
97+
permute(block);
98+
round(state, block);
99+
permute(block);
100+
round(state, block);
101+
permute(block);
102+
round(state, block);
103+
permute(block);
104+
round(state, block);
105+
106+
// Final mixing
107+
for (let i = 0; i < 8; i++) {
108+
state[i] ^= state[i + 8];
109+
state[i + 8] ^= chaining_value[i];
110+
}
111+
112+
return state;
113+
}
114+
115+
function words_from_little_endian_bytes(bytes: Uint8Array, words: StaticArray<u32>): void {
116+
for (let i = 0; i < words.length; i++) {
117+
const offset = i * 4;
118+
words[i] = bytes[offset] | (bytes[offset + 1] << 8) | (bytes[offset + 2] << 16) | (bytes[offset + 3] << 24);
119+
}
120+
}
121+
122+
export class Blake3Hasher {
123+
private chunk_state: ChunkState;
124+
private key_words: StaticArray<u32>;
125+
private cv_stack: StaticArray<StaticArray<u32>>;
126+
private cv_stack_len: u8;
127+
private flags: u32;
128+
129+
constructor() {
130+
this.key_words = new StaticArray<u32>(8);
131+
for (let i = 0; i < 8; i++) {
132+
this.key_words[i] = IV[i];
133+
}
134+
this.chunk_state = new ChunkState(this.key_words, 0, 0);
135+
this.cv_stack = new StaticArray<StaticArray<u32>>(54);
136+
for (let i = 0; i < 54; i++) {
137+
this.cv_stack[i] = new StaticArray<u32>(8);
138+
}
139+
this.cv_stack_len = 0;
140+
this.flags = 0;
141+
}
142+
143+
update(input: Uint8Array): void {
144+
let inputPos = 0;
145+
while (inputPos < input.length) {
146+
if (this.chunk_state.len() == CHUNK_LEN) {
147+
const chunk_cv = this.chunk_state.output().chaining_value();
148+
const total_chunks = this.chunk_state.chunk_counter + 1;
149+
this.add_chunk_chaining_value(chunk_cv, total_chunks);
150+
this.chunk_state = new ChunkState(this.key_words, total_chunks, this.flags);
151+
}
152+
153+
const want = CHUNK_LEN - this.chunk_state.len();
154+
const take = min(want, input.length - inputPos);
155+
this.chunk_state.update(input.subarray(inputPos, inputPos + take));
156+
inputPos += take;
157+
}
158+
}
159+
160+
finalize(out: Uint8Array): void {
161+
let output = this.chunk_state.output();
162+
let parent_nodes_remaining = this.cv_stack_len;
163+
164+
while (parent_nodes_remaining > 0) {
165+
parent_nodes_remaining--;
166+
output = parent_output(
167+
this.cv_stack[parent_nodes_remaining],
168+
output.chaining_value(),
169+
this.key_words,
170+
this.flags
171+
);
172+
}
173+
174+
output.root_output_bytes(out);
175+
}
176+
177+
private add_chunk_chaining_value(new_cv: StaticArray<u32>, total_chunks: u64): void {
178+
let mut_new_cv = new_cv;
179+
let mut_total_chunks = total_chunks;
180+
181+
while ((mut_total_chunks & 1) == 0) {
182+
mut_new_cv = parent_cv(this.pop_stack(), mut_new_cv, this.key_words, this.flags);
183+
mut_total_chunks >>= 1;
184+
}
185+
186+
this.push_stack(mut_new_cv);
187+
}
188+
189+
private push_stack(cv: StaticArray<u32>): void {
190+
for (let i = 0; i < 8; i++) {
191+
this.cv_stack[this.cv_stack_len][i] = cv[i];
192+
}
193+
this.cv_stack_len++;
194+
}
195+
196+
private pop_stack(): StaticArray<u32> {
197+
this.cv_stack_len--;
198+
return this.cv_stack[this.cv_stack_len];
199+
}
200+
}
201+
202+
class ChunkState {
203+
chaining_value: StaticArray<u32>;
204+
chunk_counter: u64;
205+
block: Uint8Array;
206+
block_len: u8;
207+
blocks_compressed: u8;
208+
flags: u32;
209+
210+
constructor(key_words: StaticArray<u32>, chunk_counter: u64, flags: u32) {
211+
this.chaining_value = new StaticArray<u32>(8);
212+
for (let i = 0; i < 8; i++) {
213+
this.chaining_value[i] = key_words[i];
214+
}
215+
this.chunk_counter = chunk_counter;
216+
this.block = new Uint8Array(BLOCK_LEN);
217+
this.block_len = 0;
218+
this.blocks_compressed = 0;
219+
this.flags = flags;
220+
}
221+
222+
len(): usize {
223+
return BLOCK_LEN * this.blocks_compressed + this.block_len;
224+
}
225+
226+
start_flag(): u32 {
227+
return this.blocks_compressed == 0 ? CHUNK_START : 0;
228+
}
229+
230+
update(input: Uint8Array): void {
231+
let inputPos = 0;
232+
while (inputPos < input.length) {
233+
if (this.block_len == BLOCK_LEN) {
234+
const block_words = new StaticArray<u32>(16);
235+
words_from_little_endian_bytes(this.block, block_words);
236+
const compressed = compress(
237+
this.chaining_value,
238+
block_words,
239+
this.chunk_counter,
240+
BLOCK_LEN,
241+
this.flags | this.start_flag()
242+
);
243+
for (let i = 0; i < 8; i++) {
244+
this.chaining_value[i] = compressed[i];
245+
}
246+
this.blocks_compressed++;
247+
this.block = new Uint8Array(BLOCK_LEN);
248+
this.block_len = 0;
249+
}
250+
251+
const want = BLOCK_LEN - this.block_len;
252+
const take = min(want, input.length - inputPos);
253+
for (let i = 0; i < take; i++) {
254+
this.block[this.block_len + i] = input[inputPos + i];
255+
}
256+
this.block_len += take;
257+
inputPos += take;
258+
}
259+
}
260+
261+
output(): Output {
262+
const block_words = new StaticArray<u32>(16);
263+
words_from_little_endian_bytes(this.block, block_words);
264+
return new Output(
265+
this.chaining_value,
266+
block_words,
267+
this.chunk_counter,
268+
this.block_len,
269+
this.flags | this.start_flag() | CHUNK_END
270+
);
271+
}
272+
}
273+
274+
class Output {
275+
input_chaining_value: StaticArray<u32>;
276+
block_words: StaticArray<u32>;
277+
counter: u64;
278+
block_len: u32;
279+
flags: u32;
280+
281+
constructor(
282+
input_chaining_value: StaticArray<u32>,
283+
block_words: StaticArray<u32>,
284+
counter: u64,
285+
block_len: u32,
286+
flags: u32
287+
) {
288+
this.input_chaining_value = input_chaining_value;
289+
this.block_words = block_words;
290+
this.counter = counter;
291+
this.block_len = block_len;
292+
this.flags = flags;
293+
}
294+
295+
chaining_value(): StaticArray<u32> {
296+
const compressed = compress(this.input_chaining_value, this.block_words, this.counter, this.block_len, this.flags);
297+
const result = new StaticArray<u32>(8);
298+
for (let i = 0; i < 8; i++) {
299+
result[i] = compressed[i];
300+
}
301+
return result;
302+
}
303+
304+
root_output_bytes(out: Uint8Array): void {
305+
let output_block_counter: u64 = 0;
306+
for (let i = 0; i < out.length; i += 2 * OUT_LEN) {
307+
const words = compress(
308+
this.input_chaining_value,
309+
this.block_words,
310+
output_block_counter,
311+
this.block_len,
312+
this.flags | ROOT
313+
);
314+
const out_block = out.subarray(i, i + 2 * OUT_LEN);
315+
for (let j = 0; j < words.length; j++) {
316+
const word = words[j];
317+
const offset = j * 4;
318+
if (offset < out_block.length) {
319+
out_block[offset] = word & 0xff;
320+
if (offset + 1 < out_block.length) {
321+
out_block[offset + 1] = (word >> 8) & 0xff;
322+
if (offset + 2 < out_block.length) {
323+
out_block[offset + 2] = (word >> 16) & 0xff;
324+
if (offset + 3 < out_block.length) {
325+
out_block[offset + 3] = (word >> 24) & 0xff;
326+
}
327+
}
328+
}
329+
}
330+
}
331+
output_block_counter++;
332+
}
333+
}
334+
}
335+
336+
function parent_output(
337+
left_child_cv: StaticArray<u32>,
338+
right_child_cv: StaticArray<u32>,
339+
key_words: StaticArray<u32>,
340+
flags: u32
341+
): Output {
342+
const block_words = new StaticArray<u32>(16);
343+
for (let i = 0; i < 8; i++) {
344+
block_words[i] = left_child_cv[i];
345+
block_words[i + 8] = right_child_cv[i];
346+
}
347+
return new Output(key_words, block_words, 0, BLOCK_LEN, PARENT | flags);
348+
}
349+
350+
function parent_cv(
351+
left_child_cv: StaticArray<u32>,
352+
right_child_cv: StaticArray<u32>,
353+
key_words: StaticArray<u32>,
354+
flags: u32
355+
): StaticArray<u32> {
356+
return parent_output(left_child_cv, right_child_cv, key_words, flags).chaining_value();
357+
}
358+
359+
function min(a: usize, b: usize): usize {
360+
return a < b ? a : b;
361+
}

‎packages/blake3-wasm/assembly/index.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
// Re-export everything from blake3.ts
2+
export * from "./blake3";

‎packages/blake3-wasm/assembly/tsconfig.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"extends": "../node_modules/.pnpm/[email protected]/node_modules/assemblyscript/std/assembly.json",
3+
"include": ["./**/*.ts"]
4+
}

‎packages/blake3-wasm/package.json‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
{
2+
"name": "@huggingface/blake3-wasm",
3+
"version": "0.0.1",
4+
"scripts": {
5+
"build:debug": "asc assembly/index.ts --target debug",
6+
"build:release": "asc assembly/index.ts --target release",
7+
"build": "npm run build:debug && npm run build:release",
8+
"test": "node tests"
9+
},
10+
"dependencies": {
11+
"assemblyscript": "^0.27.36"
12+
},
13+
"type": "module",
14+
"exports": {
15+
".": {
16+
"import": "./build/release.js",
17+
"types": "./build/release.d.ts"
18+
}
19+
},
20+
"devDependencies": {
21+
"assemblyscript": "^0.27.36"
22+
}
23+
}

0 commit comments

Comments
 (0)