# Generated by tools/generators/sha3_gen.py; do not edit by hand.
import Base
import ../keccak/types.bend as T
import ../keccak/lane.bend as L
import ../keccak/permutation.bend as P

# SHA3-256 (FIPS 202) on byte lists, on the Keccak-f[1600] of
# src/crypto/keccak/permutation.bend (two rounds per step, proved equal to
# FIPS 202's permutation in proofs/crypto/keccak/permutation.bend). The
# message is padded (0x06 .. 0x80), read as little-endian lanes, and absorbed
# 17 lanes (136 bytes) at a time; the digest is the first four lanes.
# Proved equal to spec/crypto/sha3.bend in proofs/crypto/sha3/.

def pack(a: U32, b: U32, c: U32, d: U32) -> U32:
  U32.or(U32.or(U32.or(U32.and(a, 255), U32.shln(U32.and(b, 255), 8n)),
    U32.shln(U32.and(c, 255), 16n)), U32.shln(U32.and(d, 255), 24n))

def lanes(bytes: List<&2, U32>) -> List<&2, T.Lane>:
  match bytes:
    case a <> b <> c <> d <> e <> f <> g <> h <> rest:
      T.W{pack(a, b, c, d), pack(e, f, g, h)} <> lanes(rest)
    case _:
      Nil{}

# 0x06, zeros, 0x80 up to the end of the block (0x86 when one byte is left).
def tail(q: Nat) -> List<&2, U32>:
  match q:
    case 0n:
      Nil{}
    case 1n:
      [134]
    case 2n+k:
      6 <> List.append(&2, U32, List.replicate(U32, k, 0), [128])

def suffix(+n: Nat) -> List<&2, U32>:
  tail(Nat.sub(136n, Nat.mod(n, 136n)))

def zero() -> T.State:
  T.S{T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}}

# XOR one block into the leading 17 lanes.
def inject(s: T.State, l0: T.Lane, l1: T.Lane, l2: T.Lane, l3: T.Lane, l4: T.Lane, l5: T.Lane, l6: T.Lane, l7: T.Lane, l8: T.Lane, l9: T.Lane, l10: T.Lane, l11: T.Lane, l12: T.Lane, l13: T.Lane, l14: T.Lane, l15: T.Lane, l16: T.Lane) -> T.State:
  match s:
    case T.S{a0, a1, a2, a3, a4, a5, a6, a7, a8, a9, a10, a11, a12, a13, a14, a15, a16, a17, a18, a19, a20, a21, a22, a23, a24}:
      T.S{L.xor(a0, l0), L.xor(a1, l1), L.xor(a2, l2), L.xor(a3, l3), L.xor(a4, l4), L.xor(a5, l5), L.xor(a6, l6), L.xor(a7, l7), L.xor(a8, l8), L.xor(a9, l9), L.xor(a10, l10), L.xor(a11, l11), L.xor(a12, l12), L.xor(a13, l13), L.xor(a14, l14), L.xor(a15, l15), L.xor(a16, l16), a17, a18, a19, a20, a21, a22, a23, a24}

# Absorb every whole block of 17 lanes with the given round count (24).
def absorb(ls: List<&2, T.Lane>, +rounds: Nat, s: T.State) -> T.State:
  match ls:
    case l0 <> l1 <> l2 <> l3 <> l4 <> l5 <> l6 <> l7 <> l8 <> l9 <> l10 <> l11 <> l12 <> l13 <> l14 <> l15 <> l16 <> rest:
      absorb(rest, rounds, P.rounds(rounds, 0n, inject(s, l0, l1, l2, l3, l4, l5, l6, l7, l8, l9, l10, l11, l12, l13, l14, l15, l16)))
    case _:
      s

def le_bytes(+x: U32, tail: List<&2, U32>) -> List<&2, U32>:
  U32.and(x, 255) <> U32.and(U32.shrn(x, 8n), 255) <> U32.and(U32.shrn(x, 16n), 255) <> U32.shrn(x, 24n) <> tail

# The 32-byte digest: lanes 0..3, little-endian.
def digest_bytes(s: T.State) -> List<&2, U32>:
  match s:
    case T.S{T.W{l0, h0}, T.W{l1, h1}, T.W{l2, h2}, T.W{l3, h3}, a4, a5, a6, a7, a8, a9, a10, a11, a12, a13, a14, a15, a16, a17, a18, a19, a20, a21, a22, a23, a24}:
      le_bytes(l0, le_bytes(h0, le_bytes(l1, le_bytes(h1, le_bytes(l2, le_bytes(h2, le_bytes(l3, le_bytes(h3, Nil{}))))))))

def hash_padded(bytes: List<&2, U32>, +rounds: Nat) -> List<&2, U32>:
  digest_bytes(absorb(lanes(bytes), rounds, zero()))

def keccak(+bytes: List<&2, U32>, +rounds: Nat) -> List<&2, U32>:
  hash_padded(List.append(&2, U32, bytes, suffix(List.length(&2, U32, bytes))), rounds)

def sha3_256(bytes: List<&2, U32>) -> List<&2, U32>:
  keccak(bytes, 24n)