warriors/trisha/scripts/generate_sha256_baseline.py

#!/usr/bin/env python3
"""Emit a direct FIPS180-4 SHA256 RAM implementation, independent of Trident.

The baseline owns RAM50000..50301. Its public ABI has first arguments deepest
and final result h7 on top. Message words must include SHA256 padding.
No compiler output, baseline candidate output or reference digest is an input.
"""
from pathlib import Path

IV = [0x6A09E667, 0xBB67AE85, 0x3C6EF372, 0xA54FF53A,
      0x510E527F, 0x9B05688C, 0x1F83D9AB, 0x5BE0CD19]
K = [1116352408, 1899447441, 3049323471, 3921009573, 961987163, 1508970993, 2453635748, 2870763221, 3624381080, 310598401, 607225278, 1426881987, 1925078388, 2162078206, 2614888103, 3248222580, 3835390401, 4022224774, 264347078, 604807628, 770255983, 1249150122, 1555081692, 1996064986, 2554220882, 2821834349, 2952996808, 3210313671, 3336571891, 3584528711, 113926993, 338241895, 666307205, 773529912, 1294757372, 1396182291, 1695183700, 1986661051, 2177026350, 2456956037, 2730485921, 2820302411, 3259730800, 3345764771, 3516065817, 3600352804, 4094571909, 275423344, 430227734, 506948616, 659060556, 883997877, 958139571, 1322822218, 1537002063, 1747873779, 1955562222, 2024104815, 2227730452, 2361852424, 2428436474, 2756734187, 3204031479, 3329325298]  # FIPS180-4 section4.2.2
STATE, ORIGINAL, WORDS, TEMP1, TEMP2 = 50000, 50008, 50100, 50300, 50301
lines = ["// Generated by scripts/generate_sha256_baseline.py from FIPS180-4.",
         "// RAM50000..50301 reserved; canonical source argument/result ordering."]


def emit(*instructions):
    lines.extend("    " + instruction for instruction in instructions)


def label(name):
    lines.extend(["", name + ":"])


def load(address):
    emit(f"push {address}", "read_mem 1", "pop 1")


def store(address):
    emit(f"push {address}", "write_mem 1", "pop 1")


def narrow():
    emit("split", "swap 1", "pop 1")


def rotate(address, bits):
    emit(f"push {1 << bits}")
    load(address)
    emit("div_mod", f"push {1 << (32-bits)}", "mul", "add")


def shift(address, bits):
    emit(f"push {1 << bits}")
    load(address)
    emit("div_mod", "pop 1")


def sigma(address, rotations, shifted=None):
    for index, bits in enumerate(rotations):
        rotate(address, bits)
        if index:
            emit("xor")
    if shifted is not None:
        shift(address, shifted)
        emit("xor")


def initialize():
    for index, value in enumerate(IV):
        emit(f"push {value}", "dup 0")
        store(STATE + index)
        store(ORIGINAL + index)


label("__init")
for value in IV:
    emit(f"push {value}")
emit("return")

label("__compress")
for index in reversed(range(16)):
    store(WORDS + index)
for index in reversed(range(8)):
    emit("dup 0")
    store(STATE + index)
    store(ORIGINAL + index)
emit("call __sha256_compression", "return")

label("__double_sha256_single_block")
for index in reversed(range(16)):
    store(WORDS + index)
initialize()
emit("call __sha256_compression")
for index in reversed(range(8)):
    store(WORDS + index)
for index in range(8, 16):
    value = 0x80000000 if index == 8 else 256 if index == 15 else 0
    emit(f"push {value}")
    store(WORDS + index)
initialize()
emit("call __sha256_compression", "return")

label("__sha256_compression")
for index in range(16, 64):
    sigma(WORDS + index - 2, [17, 19], 10)
    load(WORDS + index - 7)
    emit("add")
    sigma(WORDS + index - 15, [7, 18], 3)
    emit("add")
    load(WORDS + index - 16)
    emit("add")
    narrow()
    store(WORDS + index)

for index, constant in enumerate(K):
    # T1 = h + Sigma1(e) + Ch(e,f,g) + K[t] + W[t].
    load(STATE + 7)
    sigma(STATE + 4, [6, 11, 25])
    emit("add")
    load(STATE + 4)
    load(STATE + 5)
    emit("and")
    load(STATE + 4)
    emit("push 4294967295", "xor")
    load(STATE + 6)
    emit("and", "xor", "add", f"push {constant}", "add")
    load(WORDS + index)
    emit("add")
    narrow()
    store(TEMP1)
    # T2 = Sigma0(a) + Maj(a,b,c).
    sigma(STATE, [2, 13, 22])
    for left, right in [(0, 1), (0, 2), (1, 2)]:
        load(STATE + left)
        load(STATE + right)
        emit("and")
        if (left, right) != (0, 1):
            emit("xor")
    emit("add")
    narrow()
    store(TEMP2)
    for destination, source in [(7, 6), (6, 5), (5, 4)]:
        load(STATE + source)
        store(STATE + destination)
    load(STATE + 3)
    load(TEMP1)
    emit("add")
    narrow()
    store(STATE + 4)
    for destination, source in [(3, 2), (2, 1), (1, 0)]:
        load(STATE + source)
        store(STATE + destination)
    load(TEMP1)
    load(TEMP2)
    emit("add")
    narrow()
    store(STATE)
for index in range(8):
    load(STATE + index)
    load(ORIGINAL + index)
    emit("add")
    narrow()
emit("return")

if __name__ == "__main__":
    assert len(K) == 64
    path = Path(__file__).resolve().parents[1] / "baselines/triton/std/crypto/sha256.tasm"
    path.write_text("\n".join(lines) + "\n")

    # Complete single-block double hashes, including padding-boundary and carry
    # cases. The oracle is Python hashlib, independently of both implementations.
    import hashlib
    import struct
    fixtures = path.parents[2] / 'fixtures'
    ram_driver = '''program sha256_ram
use std.crypto.sha256
use vm.core.convert
use vm.io.mem
fn main() {
    for i in 0..16 bounded 16 { mem.write(1000 + convert.as_field(i),pub_read()) }
    sha256.double_sha256_single_block_in_ram(1000,2000,3000)
    for i in 0..8 bounded 8 { pub_write(mem.read(2000 + convert.as_field(i))) }
}
'''
    struct_driver = 'program sha256_aggregate\nuse std.crypto.sha256\nfn main() {\n'+''.join(f'let w{i} = as_u32(pub_read())\n' for i in range(16))+'let d=sha256.double_sha256_single_block('+','.join(f'w{i}' for i in range(16))+')\n'+''.join(f'pub_write(as_field(d.h{i}))\n' for i in range(8))+'}\n'
    prefix = 'read_io 1 '*16+'call __double_sha256_single_block swap 7 swap 1 swap 6 swap 1 swap 2 swap 5 swap 2 swap 3 swap 4 swap 3 write_io 5 write_io 3 halt'
    for name,message in [('abc',b'abc'),('empty',b''),('55bytes',bytes(range(55))),('carry',b'\xff'*55)]:
        padded=message+b'\x80'+b'\0'*(55-len(message))+struct.pack('>Q',8*len(message))
        words=list(struct.unpack('>16I',padded))
        expected=list(struct.unpack('>8I',hashlib.sha256(hashlib.sha256(message).digest()).digest()))
        directory=fixtures/f'sha256-{name}';directory.mkdir(exist_ok=True)
        (directory/'main.tri').write_text(ram_driver)
        (directory/'struct.tri').write_text(struct_driver)
        (directory/'vector.bench.toml').write_text(f'source = "main.tri"\nhand = "../../std/crypto/sha256.tasm"\ntarget = "triton"\ninput = {words}\noutput = {expected}\nmax_cycles = 1000000\nreference = "Full double-SHA256, caller-owned RAM API, FIPS180-4 single-block padding and independent hashlib.sha256 twice; all256 output bits. Struct API is separately regression-tested against the same vectors."\nhand_prefix = "{prefix}"\n')

    # A two-block message exercises compress_in_ram with a non-IV chaining
    # state on its second invocation, independently checked by hashlib.
    message=bytes(range(64))
    padded=message+b'\x80'+b'\0'*55+struct.pack('>Q',8*len(message))
    words=list(struct.unpack('>32I',padded))
    expected=list(struct.unpack('>8I',hashlib.sha256(message).digest()))
    directory=fixtures/'sha256-two-blocks';directory.mkdir(exist_ok=True)
    source='program sha256_compression_chain\nuse std.crypto.sha256\nuse vm.io.mem\nuse vm.core.convert\nfn main() {\nfor i in 0..32 bounded 32 { mem.write(1000 + convert.as_field(i),pub_read()) }\n'
    source+=''.join(f'mem.write(2000 + {i},{v})\n' for i,v in enumerate(IV))
    source+='sha256.compress_in_ram(2000,1000,3000)\nsha256.compress_in_ram(2000,1016,3000)\nfor i in 0..8 bounded 8 { pub_write(mem.read(2000 + convert.as_field(i))) }\n}\n'
    (directory/'main.tri').write_text(source)
    prefix='call __init '+'read_io 1 '*16+'call __compress '+'read_io 1 '*16+'call __compress swap 7 swap 1 swap 6 swap 1 swap 2 swap 5 swap 2 swap 3 swap 4 swap 3 write_io 5 write_io 3 halt'
    (directory/'vector.bench.toml').write_text(f'source = "main.tri"\nhand = "../../std/crypto/sha256.tasm"\ninput = {words}\noutput = {expected}\nmax_cycles = 1000000\nreference = "Two full SHA256 compression blocks with non-IV chaining state on second invocation; hashlib.sha256(bytes(range(64))), all256 output bits."\nhand_prefix = "{prefix}"\n')