// Copyright (c) 2023 Kim Shrier. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
// Package sha3 implements the 512, 384, 256, and 224
// bit hash algorithms and the 128 and 256 bit
// extended output functions as defined in
// https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.202.pdf.
// Last updated: August 2015
module sha3

import encoding.binary
import math.bits

// when the state is 1600 bits, a lane is 64 bits
type Lane = u64

// the state is a 5 x 5 array of lanes
struct State {
mut:
	a [5][5]Lane
}

// to_bytes converts the state into a byte array
//
// A 1600 bit state fits into 200 bytes.
//
// The state consists of 25 u64 values arranged in a
// 5 x 5 array.
//
// The state is not modified by this method.
@[direct_array_access]
fn (s State) to_bytes() []u8 {
	mut byte_array := []u8{len: 200}
	mut index := 0

	for y in 0 .. 5 {
		for x in 0 .. 5 {
			unsafe {
				binary.little_endian_put_u64_at(mut byte_array, s.a[x][y], index)
			}
			index += 8
		}
	}

	return byte_array
}

// from_bytes sets the state from an array of bytes
//
// A 1600 bit state fits into 200 bytes.  It is expected
// that the input byte array is 200 bytes long.
//
// All 25 u64 values are set from the input byte array.
@[direct_array_access]
fn (mut s State) from_bytes(byte_array []u8) {
	mut index := 0

	for y in 0 .. 5 {
		for x in 0 .. 5 {
			s.a[x][y] = binary.little_endian_u64_at(byte_array, index)
			index += 8
		}
	}
}

// xor_bytes xor's an array of bytes into the state
//
// This is how new data gets absorbed into the sponge.
//
// It is expected that the length of the byte array is
// the same as the rate.  The rate is how many bytes
// can be absorbed at one permutation of the state.
//
// The rate must be less than the size of the state
// in bytes.
@[direct_array_access]
fn (mut s State) xor_bytes(byte_array []u8, rate int) {
	mut index := 0

	for y in 0 .. 5 {
		for x in 0 .. 5 {
			s.a[x][y] ^= binary.little_endian_u64_at(byte_array, index)
			index += 8

			if index >= rate {
				return
			}
		}
	}
}

// kaccak_p_1600_24 performs 24 rounnds on a 1600 bit state.
//
@[direct_array_access]
fn (mut s State) kaccak_p_1600_24() {
	mut b := [5][5]Lane{}
	for round_index in 0 .. 24 {
		// theta
		// C[x] = A[x,0] xor A[x,1] xor A[x,2] xor A[x,3] xor A[x,4],   for x in 0…4
		// D[x] = C[x-1] xor rot(C[x+1],1),                             for x in 0…4
		// A[x,y] = A[x,y] xor D[x],                           for (x,y) in (0…4,0…4)
		c0 := s.a[0][0] ^ s.a[0][1] ^ s.a[0][2] ^ s.a[0][3] ^ s.a[0][4]
		c1 := s.a[1][0] ^ s.a[1][1] ^ s.a[1][2] ^ s.a[1][3] ^ s.a[1][4]
		c2 := s.a[2][0] ^ s.a[2][1] ^ s.a[2][2] ^ s.a[2][3] ^ s.a[2][4]
		c3 := s.a[3][0] ^ s.a[3][1] ^ s.a[3][2] ^ s.a[3][3] ^ s.a[3][4]
		c4 := s.a[4][0] ^ s.a[4][1] ^ s.a[4][2] ^ s.a[4][3] ^ s.a[4][4]

		d0 := c4 ^ bits.rotate_left_64(c1, 1)
		d1 := c0 ^ bits.rotate_left_64(c2, 1)
		d2 := c1 ^ bits.rotate_left_64(c3, 1)
		d3 := c2 ^ bits.rotate_left_64(c4, 1)
		d4 := c3 ^ bits.rotate_left_64(c0, 1)

		// vfmt off
		s.a[0][0] ^= d0 s.a[0][1] ^= d0 s.a[0][2] ^= d0 s.a[0][3] ^= d0 s.a[0][4] ^= d0
		s.a[1][0] ^= d1 s.a[1][1] ^= d1 s.a[1][2] ^= d1 s.a[1][3] ^= d1 s.a[1][4] ^= d1
		s.a[2][0] ^= d2 s.a[2][1] ^= d2 s.a[2][2] ^= d2	s.a[2][3] ^= d2	s.a[2][4] ^= d2
		s.a[3][0] ^= d3	s.a[3][1] ^= d3	s.a[3][2] ^= d3	s.a[3][3] ^= d3	s.a[3][4] ^= d3
		s.a[4][0] ^= d4	s.a[4][1] ^= d4	s.a[4][2] ^= d4	s.a[4][3] ^= d4	s.a[4][4] ^= d4
		// vfmt on

		// rho and pi
		// B[y,2*x+3*y] = rot(A[x,y], r[x,y]),                 for (x,y) in (0…4,0…4)
		b[0][0] = s.a[0][0]
		b[0][1] = bits.rotate_left_64(s.a[3][0], 28)
		b[0][2] = bits.rotate_left_64(s.a[1][0], 1)
		b[0][3] = bits.rotate_left_64(s.a[4][0], 27)
		b[0][4] = bits.rotate_left_64(s.a[2][0], 62)

		b[1][0] = bits.rotate_left_64(s.a[1][1], 44)
		b[1][1] = bits.rotate_left_64(s.a[4][1], 20)
		b[1][2] = bits.rotate_left_64(s.a[2][1], 6)
		b[1][3] = bits.rotate_left_64(s.a[0][1], 36)
		b[1][4] = bits.rotate_left_64(s.a[3][1], 55)

		b[2][0] = bits.rotate_left_64(s.a[2][2], 43)
		b[2][1] = bits.rotate_left_64(s.a[0][2], 3)
		b[2][2] = bits.rotate_left_64(s.a[3][2], 25)
		b[2][3] = bits.rotate_left_64(s.a[1][2], 10)
		b[2][4] = bits.rotate_left_64(s.a[4][2], 39)

		b[3][0] = bits.rotate_left_64(s.a[3][3], 21)
		b[3][1] = bits.rotate_left_64(s.a[1][3], 45)
		b[3][2] = bits.rotate_left_64(s.a[4][3], 8)
		b[3][3] = bits.rotate_left_64(s.a[2][3], 15)
		b[3][4] = bits.rotate_left_64(s.a[0][3], 41)

		b[4][0] = bits.rotate_left_64(s.a[4][4], 14)
		b[4][1] = bits.rotate_left_64(s.a[2][4], 61)
		b[4][2] = bits.rotate_left_64(s.a[0][4], 18)
		b[4][3] = bits.rotate_left_64(s.a[3][4], 56)
		b[4][4] = bits.rotate_left_64(s.a[1][4], 2)

		// chi
		// A[x,y] = B[x,y] xor ((not B[x+1,y]) and B[x+2,y]),  for (x,y) in (0…4,0…4)
		for y in 0 .. 5 {
			s.a[0][y] = b[0][y] ^ (~b[1][y] & b[2][y])
			s.a[1][y] = b[1][y] ^ (~b[2][y] & b[3][y])
			s.a[2][y] = b[2][y] ^ (~b[3][y] & b[4][y])
			s.a[3][y] = b[3][y] ^ (~b[4][y] & b[0][y])
			s.a[4][y] = b[4][y] ^ (~b[0][y] & b[1][y])
		}

		// iota
		// A[0,0] = A[0,0] xor RC
		s.a[0][0] ^= iota_round_constants[round_index]
	}
}

// iota_round_constants are precomputed xor masks for the iota function
const iota_round_constants = [u64(0x0000000000000001), 0x0000000000008082, 0x800000000000808a,
	0x8000000080008000, 0x000000000000808b, 0x0000000080000001, 0x8000000080008081,
	0x8000000000008009, 0x000000000000008a, 0x0000000000000088, 0x0000000080008009,
	0x000000008000000a, 0x000000008000808b, 0x800000000000008b, 0x8000000000008089,
	0x8000000000008003, 0x8000000000008002, 0x8000000000000080, 0x000000000000800a,
	0x800000008000000a, 0x8000000080008081, 0x8000000000008080, 0x0000000080000001,
	0x8000000080008008]

fn (s State) str() string {
	mut output := '\n             y = 0            y = 1            y = 2            y = 3            y = 4\n'
	for x in 0 .. 5 {
		output += 'x = ${x}: ${s.a[x][0]:016x} ${s.a[x][1]:016x} ${s.a[x][2]:016x} ${s.a[x][3]:016x} ${s.a[x][4]:016x}\n'
	}

	return output
}