Skip to main content

runar_compiler_rust/codegen/
sha256.rs

1//! SHA-256 compression codegen for Bitcoin Script.
2//!
3//! Port of packages/runar-compiler/src/passes/sha256-codegen.ts.
4//!
5//! emitSha256Compress: [state(32), block(64)] -> [newState(32)]
6//!
7//! Optimized architecture (inspired by twostack/tstokenlib):
8//!   - All 32-bit words stored as 4-byte little-endian during computation.
9//!   - Bitwise ops (AND, OR, XOR, INVERT) are endian-agnostic on equal-length arrays.
10//!   - ROTR uses OP_RSHIFT/OP_LSHIFT (native BE byte-array shifts).
11//!   - Batched addN for T1 (5 addends) converts all to numeric once, adds, converts back.
12//!   - BE<->LE conversion only at input unpack / output pack.
13
14use super::stack::{PushValue, StackOp};
15
16use std::sync::OnceLock;
17
18// SHA-256 round constants
19const K: [u32; 64] = [
20    0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5,
21    0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,
22    0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3,
23    0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174,
24    0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc,
25    0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,
26    0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7,
27    0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967,
28    0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13,
29    0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85,
30    0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3,
31    0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,
32    0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5,
33    0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
34    0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208,
35    0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2,
36];
37
38/// Encode a uint32 as 4-byte little-endian.
39fn u32_to_le(n: u32) -> Vec<u8> {
40    vec![
41        (n & 0xff) as u8,
42        ((n >> 8) & 0xff) as u8,
43        ((n >> 16) & 0xff) as u8,
44        ((n >> 24) & 0xff) as u8,
45    ]
46}
47
48// =========================================================================
49// Emitter with depth tracking
50// =========================================================================
51
52struct Emitter {
53    ops: Vec<StackOp>,
54    depth: i64,
55    alt_depth: i64,
56}
57
58impl Emitter {
59    fn new(initial_depth: i64) -> Self {
60        Emitter {
61            ops: Vec::new(),
62            depth: initial_depth,
63            alt_depth: 0,
64        }
65    }
66
67    fn e_raw(&mut self, sop: StackOp) {
68        self.ops.push(sop);
69    }
70
71    fn oc(&mut self, code: &str) {
72        self.ops.push(StackOp::Opcode(code.to_string()));
73    }
74
75    fn push_i(&mut self, v: i128) {
76        self.ops.push(StackOp::Push(PushValue::Int(v)));
77        self.depth += 1;
78    }
79
80    fn push_b(&mut self, v: Vec<u8>) {
81        self.ops.push(StackOp::Push(PushValue::Bytes(v)));
82        self.depth += 1;
83    }
84
85    fn dup(&mut self) {
86        self.ops.push(StackOp::Dup);
87        self.depth += 1;
88    }
89
90    fn drop(&mut self) {
91        self.ops.push(StackOp::Drop);
92        self.depth -= 1;
93    }
94
95    fn swap(&mut self) {
96        self.ops.push(StackOp::Swap);
97    }
98
99    fn over(&mut self) {
100        self.ops.push(StackOp::Over);
101        self.depth += 1;
102    }
103
104    fn rot(&mut self) {
105        self.ops.push(StackOp::Rot);
106    }
107
108    fn pick(&mut self, d: usize) {
109        if d == 0 {
110            self.dup();
111            return;
112        }
113        if d == 1 {
114            self.over();
115            return;
116        }
117        self.push_i(d as i128);
118        self.ops.push(StackOp::Pick { depth: d });
119        // push_i added 1, pick removes the depth literal but adds the picked value = net 0
120    }
121
122    fn roll(&mut self, d: usize) {
123        if d == 0 {
124            return;
125        }
126        if d == 1 {
127            self.swap();
128            return;
129        }
130        if d == 2 {
131            self.rot();
132            return;
133        }
134        self.push_i(d as i128);
135        self.ops.push(StackOp::Roll { depth: d });
136        self.depth -= 1; // push_i added 1, roll removes depth literal and item = net -1
137    }
138
139    fn to_alt(&mut self) {
140        self.oc("OP_TOALTSTACK");
141        self.depth -= 1;
142        self.alt_depth += 1;
143    }
144
145    fn from_alt(&mut self) {
146        self.oc("OP_FROMALTSTACK");
147        self.depth += 1;
148        self.alt_depth -= 1;
149    }
150
151    fn bin_op(&mut self, code: &str) {
152        self.oc(code);
153        self.depth -= 1;
154    }
155
156    fn uni_op(&mut self, code: &str) {
157        self.oc(code);
158    }
159
160    fn dup2(&mut self) {
161        self.oc("OP_2DUP");
162        self.depth += 2;
163    }
164
165    fn split(&mut self) {
166        self.oc("OP_SPLIT");
167        // splits: consumes 2 (value + position), produces 2 = net 0
168    }
169
170    fn split4(&mut self) {
171        self.push_i(4);
172        self.split();
173    }
174
175    fn assert_depth(&self, expected: i64, msg: &str) {
176        assert_eq!(
177            self.depth, expected,
178            "SHA256 codegen: {}. Expected depth {}, got {}",
179            msg, expected, self.depth
180        );
181    }
182
183    // --- Byte reversal (only for BE<->LE conversion at boundaries) ---
184
185    /// Reverse 4 bytes on TOS: [abcd] -> [dcba]. Net: 0. 12 ops.
186    fn reverse_bytes4(&mut self) {
187        self.push_i(1);
188        self.split();
189        self.push_i(1);
190        self.split();
191        self.push_i(1);
192        self.split();
193        self.swap();
194        self.bin_op("OP_CAT");
195        self.swap();
196        self.bin_op("OP_CAT");
197        self.swap();
198        self.bin_op("OP_CAT");
199    }
200
201    // --- LE <-> Numeric conversions ---
202
203    /// Convert 4-byte LE to unsigned script number. [le4] -> [num]. Net: 0. 3 ops.
204    fn le2num(&mut self) {
205        self.push_b(vec![0x00]); // unsigned padding
206        self.bin_op("OP_CAT");
207        self.uni_op("OP_BIN2NUM");
208    }
209
210    /// Convert script number to 4-byte LE (truncates to 32 bits). [num] -> [le4]. Net: 0. 5 ops.
211    fn num2le(&mut self) {
212        self.push_i(5);
213        self.bin_op("OP_NUM2BIN"); // 5-byte LE
214        self.push_i(4);
215        self.split(); // [4-byte LE, overflow+sign]
216        self.drop(); // discard overflow byte
217    }
218
219    // --- LE arithmetic ---
220
221    /// [a(LE), b(LE)] -> [(a+b mod 2^32)(LE)]. Net: -1. 13 ops.
222    fn add32(&mut self) {
223        self.le2num();
224        self.swap();
225        self.le2num();
226        self.bin_op("OP_ADD");
227        self.num2le();
228    }
229
230    /// Add N LE values. [v0..vN-1] (vN-1=TOS) -> [sum(LE)]. Net: -(N-1).
231    fn add_n(&mut self, n: usize) {
232        if n < 2 {
233            return;
234        }
235        self.le2num();
236        for _ in 1..n {
237            self.swap();
238            self.le2num();
239            self.bin_op("OP_ADD");
240        }
241        self.num2le();
242    }
243
244    // --- ROTR/SHR using OP_LSHIFT/OP_RSHIFT (native BE byte-array shifts) ---
245
246    /// ROTR(x, n) on BE 4-byte value. [x_BE] -> [rotated_BE]. Net: 0. 7 ops.
247    fn rotr_be(&mut self, n: usize) {
248        self.dup(); // [x, x]
249        self.push_i(n as i128);
250        self.bin_op("OP_RSHIFT"); // [x, x>>n]
251        self.swap(); // [x>>n, x]
252        self.push_i((32 - n) as i128);
253        self.bin_op("OP_LSHIFT"); // [x>>n, x<<(32-n)]
254        self.bin_op("OP_OR"); // [ROTR result]
255    }
256
257    /// SHR(x, n) on BE 4-byte value. [x_BE] -> [shifted_BE]. Net: 0. 2 ops.
258    fn shr_be(&mut self, n: usize) {
259        self.push_i(n as i128);
260        self.bin_op("OP_RSHIFT");
261    }
262
263    // --- SHA-256 sigma functions ---
264
265    /// big_sigma0(a) = ROTR(2)^ROTR(13)^ROTR(22). [a(LE)] -> [S0(LE)]. Net: 0.
266    fn big_sigma0(&mut self) {
267        self.reverse_bytes4(); // LE -> BE
268        self.dup();
269        self.dup();
270        self.rotr_be(2);
271        self.swap();
272        self.rotr_be(13);
273        self.bin_op("OP_XOR");
274        self.swap();
275        self.rotr_be(22);
276        self.bin_op("OP_XOR");
277        self.reverse_bytes4(); // BE -> LE
278    }
279
280    /// big_sigma1(e) = ROTR(6)^ROTR(11)^ROTR(25). [e(LE)] -> [S1(LE)]. Net: 0.
281    fn big_sigma1(&mut self) {
282        self.reverse_bytes4();
283        self.dup();
284        self.dup();
285        self.rotr_be(6);
286        self.swap();
287        self.rotr_be(11);
288        self.bin_op("OP_XOR");
289        self.swap();
290        self.rotr_be(25);
291        self.bin_op("OP_XOR");
292        self.reverse_bytes4();
293    }
294
295    /// small_sigma0(x) = ROTR(7)^ROTR(18)^SHR(3). [x(LE)] -> [s0(LE)]. Net: 0.
296    fn small_sigma0(&mut self) {
297        self.reverse_bytes4();
298        self.dup();
299        self.dup();
300        self.rotr_be(7);
301        self.swap();
302        self.rotr_be(18);
303        self.bin_op("OP_XOR");
304        self.swap();
305        self.shr_be(3);
306        self.bin_op("OP_XOR");
307        self.reverse_bytes4();
308    }
309
310    /// small_sigma1(x) = ROTR(17)^ROTR(19)^SHR(10). [x(LE)] -> [s1(LE)]. Net: 0.
311    fn small_sigma1(&mut self) {
312        self.reverse_bytes4();
313        self.dup();
314        self.dup();
315        self.rotr_be(17);
316        self.swap();
317        self.rotr_be(19);
318        self.bin_op("OP_XOR");
319        self.swap();
320        self.shr_be(10);
321        self.bin_op("OP_XOR");
322        self.reverse_bytes4();
323    }
324
325    /// Ch(e,f,g) = (e&f)^(~e&g). [e, f, g] (g=TOS), all LE -> [Ch(LE)]. Net: -2.
326    fn ch(&mut self) {
327        self.rot();
328        self.dup();
329        self.uni_op("OP_INVERT");
330        self.rot();
331        self.bin_op("OP_AND");
332        self.to_alt();
333        self.bin_op("OP_AND");
334        self.from_alt();
335        self.bin_op("OP_XOR");
336    }
337
338    /// Maj(a,b,c) = (a&b)|(c&(a^b)). [a, b, c] (c=TOS), all LE -> [Maj(LE)]. Net: -2.
339    fn maj(&mut self) {
340        self.to_alt();
341        self.dup2();
342        self.bin_op("OP_AND");
343        self.to_alt();
344        self.bin_op("OP_XOR");
345        self.from_alt();
346        self.swap();
347        self.from_alt();
348        self.bin_op("OP_AND");
349        self.bin_op("OP_OR");
350    }
351
352    /// Convert N x BE words on TOS to LE, preserving stack order.
353    fn be_words_to_le(&mut self, n: usize) {
354        for _ in 0..n {
355            self.reverse_bytes4();
356            self.to_alt();
357        }
358        for _ in 0..n {
359            self.from_alt();
360        }
361    }
362
363    /// Convert 8 x BE words on TOS to LE AND reverse order.
364    fn be_words_to_le_reversed8(&mut self) {
365        for i in (0..8).rev() {
366            self.roll(i);
367            self.reverse_bytes4();
368            self.to_alt();
369        }
370        for _ in 0..8 {
371            self.from_alt();
372        }
373    }
374}
375
376// =========================================================================
377// One compression round
378// =========================================================================
379
380/// Emit one compression round. Stack: [W0..W63, a,b,c,d,e,f,g,h] (a=TOS, all LE). Net: 0.
381fn emit_round(em: &mut Emitter, t: usize) {
382    let d0 = em.depth;
383
384    // --- T1 = Sigma1(e) + Ch(e,f,g) + h + K[t] + W[t] ---
385    em.pick(4); // e copy (+1)
386    em.big_sigma1(); // Sigma1(e) (0)
387
388    em.pick(5);
389    em.pick(7);
390    em.pick(9); // e, f, g copies (+3)
391    em.ch(); // Ch(e,f,g) (-2) -> net +2
392
393    em.pick(9); // h copy (+1) -> net +3
394    em.push_b(u32_to_le(K[t])); // K[t] as LE (+1) -> net +4
395    em.pick(75 - t); // W[t] copy (+1) -> net +5
396
397    em.add_n(5); // T1 = sum of 5 (-4) -> net +1
398
399    // --- T2 = Sigma0(a) + Maj(a,b,c) ---
400    em.dup();
401    em.to_alt(); // save T1 copy to alt
402
403    em.pick(1); // a copy (+1) -> net +2
404    em.big_sigma0(); // Sigma0(a) (0)
405
406    em.pick(2);
407    em.pick(4);
408    em.pick(6); // a, b, c copies (+3) -> net +5
409    em.maj(); // Maj(a,b,c) (-2) -> net +3
410    em.add32(); // T2 = Sigma0 + Maj (-1) -> net +2
411
412    // --- Register update ---
413    em.from_alt(); // T1 copy from alt (+1) -> net +3
414
415    em.swap();
416    em.add32(); // new_a = T1 + T2 (-1) -> net +2
417
418    em.swap();
419    em.roll(5); // d to top
420    em.add32(); // new_e = d + T1 (-1) -> net +1
421
422    em.roll(8);
423    em.drop(); // drop h (-1) -> net 0
424
425    // Rotate: [ne,na,a,b,c,e,f,g] -> [na,a,b,c,ne,e,f,g]
426    em.swap();
427    em.roll(4);
428    em.roll(4);
429    em.roll(4);
430    em.roll(3);
431
432    em.assert_depth(d0, &format!("compress: after round {}", t));
433}
434
435// =========================================================================
436// Reusable compress ops generator
437// =========================================================================
438
439fn generate_compress_ops() -> Vec<StackOp> {
440    let mut em = Emitter::new(2); // pretend state+block are the only items
441
442    // Phase 1: Save init state to alt, unpack block into 16 LE words
443    em.swap();
444    em.dup();
445    em.to_alt();
446    em.to_alt();
447    em.assert_depth(1, "compress: after state save");
448
449    for _ in 0..15 {
450        em.split4();
451    }
452    em.assert_depth(16, "compress: after block unpack");
453    em.be_words_to_le(16);
454    em.assert_depth(16, "compress: after block LE convert");
455
456    // Phase 2: W expansion
457    for _t in 16..64 {
458        em.over();
459        em.small_sigma1();
460        em.pick(6 + 1);
461        em.pick(14 + 2);
462        em.small_sigma0();
463        em.pick(15 + 3);
464        em.add_n(4);
465    }
466    em.assert_depth(64, "compress: after W expansion");
467
468    // Phase 3: Unpack state into 8 LE working vars
469    em.from_alt();
470    for _ in 0..7 {
471        em.split4();
472    }
473    em.assert_depth(72, "compress: after state unpack");
474    em.be_words_to_le_reversed8();
475    em.assert_depth(72, "compress: after state LE convert");
476
477    // Phase 4: 64 compression rounds
478    for t in 0..64 {
479        emit_round(&mut em, t);
480    }
481
482    // Phase 5: Add initial state, pack result
483    em.from_alt();
484    em.assert_depth(73, "compress: before final add");
485
486    for _ in 0..7 {
487        em.split4();
488    }
489    em.be_words_to_le_reversed8();
490    em.assert_depth(80, "compress: after init unpack");
491
492    for i in 0..8 {
493        em.roll(8 - i);
494        em.add32();
495        em.to_alt();
496    }
497    em.assert_depth(64, "compress: after final add");
498
499    em.from_alt();
500    em.reverse_bytes4();
501    for _ in 1..8 {
502        em.from_alt();
503        em.reverse_bytes4();
504        em.swap();
505        em.bin_op("OP_CAT");
506    }
507    em.assert_depth(65, "compress: after pack");
508
509    for _ in 0..64 {
510        em.swap();
511        em.drop();
512    }
513    em.assert_depth(1, "compress: final");
514
515    em.ops
516}
517
518// Cache the ops since they're identical every time
519static COMPRESS_OPS: OnceLock<Vec<StackOp>> = OnceLock::new();
520
521fn get_compress_ops() -> &'static Vec<StackOp> {
522    COMPRESS_OPS.get_or_init(generate_compress_ops)
523}
524
525// =========================================================================
526// Public entry points
527// =========================================================================
528
529/// Emit SHA-256 compression in Bitcoin Script.
530/// Stack on entry: [..., state(32 BE), block(64 BE)]
531/// Stack on exit:  [..., newState(32 BE)]
532pub fn emit_sha256_compress(emit: &mut dyn FnMut(StackOp)) {
533    for op in get_compress_ops() {
534        emit(op.clone());
535    }
536}
537
538/// Emit SHA-256 finalization in Bitcoin Script.
539/// Stack on entry: [..., state(32 BE), remaining(var len BE), msgBitLen(bigint)]
540/// Stack on exit:  [..., hash(32 BE)]
541pub fn emit_sha256_finalize(emit: &mut dyn FnMut(StackOp)) {
542    let mut em = Emitter::new(3); // state + remaining + msgBitLen
543
544    // ---- Step 1: Convert msgBitLen to 8-byte BE ----
545    em.push_i(9);
546    em.bin_op("OP_NUM2BIN"); // 9-byte LE
547    em.push_i(8);
548    em.split(); // [8-byte LE, sign byte]
549    em.drop(); // [8-byte LE]
550    // Reverse 8 bytes to BE: split(4), reverse each half, cat
551    em.push_i(4);
552    em.split(); // [lo4_LE, hi4_LE]
553    em.reverse_bytes4(); // [lo4_LE, hi4_rev]
554    em.swap();
555    em.reverse_bytes4(); // [hi4_rev, lo4_rev]
556    em.bin_op("OP_CAT"); // [bitLenBE(8)]
557    em.to_alt(); // save bitLenBE to alt
558    em.assert_depth(2, "finalize: after bitLen conversion");
559
560    // ---- Step 2: Pad remaining ----
561    em.push_b(vec![0x80]);
562    em.bin_op("OP_CAT"); // [state, remaining||0x80]
563
564    // Get padded length
565    em.oc("OP_SIZE");
566    em.depth += 1; // [state, padded, paddedLen]
567
568    // Branch: 1 block (paddedLen <= 56) or 2 blocks (paddedLen > 56)
569    em.dup();
570    em.push_i(57);
571    em.bin_op("OP_LESSTHAN"); // paddedLen < 57?
572
573    em.oc("OP_IF");
574    em.depth -= 1; // consume flag
575
576    // ---- 1-block path: pad to 56 bytes ----
577    em.push_i(56);
578    em.swap();
579    em.bin_op("OP_SUB"); // zeroCount = 56 - paddedLen
580    em.push_i(0);
581    em.swap();
582    em.bin_op("OP_NUM2BIN"); // zero bytes
583    em.bin_op("OP_CAT"); // [state, padded(56 bytes)]
584    em.from_alt(); // bitLenBE from alt
585    em.bin_op("OP_CAT"); // [state, block1(64 bytes)]
586    // Splice sha256Compress ops
587    let compress_ops = get_compress_ops();
588    for op in compress_ops {
589        em.e_raw(op.clone());
590    }
591    em.depth = 1; // after compress: 1 result
592
593    em.oc("OP_ELSE");
594    em.depth = 3; // reset to branch entry: [state, padded, paddedLen]
595
596    // ---- 2-block path: pad to 120 bytes ----
597    em.push_i(120);
598    em.swap();
599    em.bin_op("OP_SUB"); // zeroCount = 120 - paddedLen
600    em.push_i(0);
601    em.swap();
602    em.bin_op("OP_NUM2BIN"); // zero bytes
603    em.bin_op("OP_CAT"); // [state, padded(120 bytes)]
604    em.from_alt(); // bitLenBE from alt
605    em.bin_op("OP_CAT"); // [state, fullPadded(128 bytes)]
606
607    // Split into 2 blocks
608    em.push_i(64);
609    em.split(); // [state, block1(64), block2(64)]
610    em.to_alt(); // save block2
611
612    // First compress: [state, block1]
613    for op in compress_ops {
614        em.e_raw(op.clone());
615    }
616    em.depth = 1; // after first compress: [midState]
617
618    // Second compress: [midState, block2]
619    em.from_alt(); // [midState, block2]
620    for op in compress_ops {
621        em.e_raw(op.clone());
622    }
623    em.depth = 1; // after second compress: [result]
624
625    em.oc("OP_ENDIF");
626    // Both paths leave 1 item (result) on stack
627    em.assert_depth(1, "finalize: final");
628
629    for op in em.ops {
630        emit(op);
631    }
632}