import sparx/uint16

fn xor(a: uint16.Uint16, b: uint16.Uint16) -> uint16.Uint16 {
  uint16.bitwise_exclusive_or(a, b)
}

pub fn rotl(x: uint16.Uint16, n: Int) -> uint16.Uint16 {
  uint16.bitwise_or(
    uint16.bitwise_shift_left(x, n),
    uint16.bitwise_shift_right(x, 16 - n),
  )
}

pub fn a(l: uint16.Uint16, r: uint16.Uint16) -> #(uint16.Uint16, uint16.Uint16) {
  let l = uint16.unsafe_from_int(uint16.to_int(rotl(l, 9)) + uint16.to_int(r))
  let r = xor(rotl(r, 2), l)
  #(l, r)
}

pub fn a_inv(
  l: uint16.Uint16,
  r: uint16.Uint16,
) -> #(uint16.Uint16, uint16.Uint16) {
  let r = xor(r, l)
  let r = rotl(r, 14)
  let l = uint16.unsafe_from_int(uint16.to_int(l) - uint16.to_int(r))
  let l = rotl(l, 7)
  #(l, r)
}

fn l_2_mix(
  x: #(uint16.Uint16, uint16.Uint16, uint16.Uint16, uint16.Uint16),
) -> #(uint16.Uint16, uint16.Uint16, uint16.Uint16, uint16.Uint16) {
  let #(x0, x1, x2, x3) = x
  let tmp = rotl(xor(x0, x1), 8)
  #(x0, x1, xor(xor(x0, tmp), x2), xor(xor(x1, tmp), x3))
}

fn l_2_swap(
  x: #(uint16.Uint16, uint16.Uint16, uint16.Uint16, uint16.Uint16),
) -> #(uint16.Uint16, uint16.Uint16, uint16.Uint16, uint16.Uint16) {
  let #(x0, x1, x2, x3) = x
  #(x2, x3, x0, x1)
}

pub fn l_2(
  x: #(uint16.Uint16, uint16.Uint16, uint16.Uint16, uint16.Uint16),
) -> #(uint16.Uint16, uint16.Uint16, uint16.Uint16, uint16.Uint16) {
  l_2_swap(l_2_mix(x))
}

pub fn l_2_inv(
  x: #(uint16.Uint16, uint16.Uint16, uint16.Uint16, uint16.Uint16),
) -> #(uint16.Uint16, uint16.Uint16, uint16.Uint16, uint16.Uint16) {
  l_2_mix(l_2_swap(x))
}

pub type Block128 =
  #(
    uint16.Uint16,
    uint16.Uint16,
    uint16.Uint16,
    uint16.Uint16,
    uint16.Uint16,
    uint16.Uint16,
    uint16.Uint16,
    uint16.Uint16,
  )

pub fn l_4(x: Block128) -> Block128 {
  let #(x0, x1, x2, x3, x4, x5, x6, x7) = x
  let tmp = rotl(xor(xor(x0, x1), xor(x2, x3)), 8)
  let x4 = xor(xor(x4, x2), tmp)
  let x5 = xor(xor(x5, x1), tmp)
  let x6 = xor(xor(x6, x0), tmp)
  let x7 = xor(xor(x7, x3), tmp)
  // SWAP(x[0], x[4]); SWAP(x[1], x[5]); SWAP(x[2], x[6]); SWAP(x[3], x[7])
  #(x4, x5, x6, x7, x0, x1, x2, x3)
}

pub fn l_4_inv(x: Block128) -> Block128 {
  let #(x0, x1, x2, x3, x4, x5, x6, x7) = x
  // Inverse swap first
  let #(x0, x1, x2, x3, x4, x5, x6, x7) = #(x4, x5, x6, x7, x0, x1, x2, x3)
  // Then apply mix
  let tmp = rotl(xor(xor(x0, x1), xor(x2, x3)), 8)
  let x4 = xor(xor(x4, x2), tmp)
  let x5 = xor(xor(x5, x1), tmp)
  let x6 = xor(xor(x6, x0), tmp)
  let x7 = xor(xor(x7, x3), tmp)
  #(x0, x1, x2, x3, x4, x5, x6, x7)
}