philbin 1.0.1

#![cfg(target_arch = "aarch64")]

use crate::arch::{HasNeon, Simd, with_simd};
use crate::base::block::{
  Array128, Block, gen_shared_block128, gen_shared_block256,
  gen_shared_block512,
};
use std::arch::aarch64::*;
use std::ops::{BitAnd, BitXor};
use zerocopy::transmute;

/// A 128-bit SIMD vector using NEON.
#[derive(Clone, Copy)]
pub struct Neon1x128<S: Simd> {
  val: uint8x16_t,
  simd: S,
}

impl<S: Simd> Block for Neon1x128<S>
where
  S: HasNeon,
{
  gen_shared_block128!(S);

  #[inline(always)]
  fn aes_encrypt_round(state: Self, round_key: Self) -> Self {
    with_simd! {
      fn helper(
        token: impl HasNeon,
        state: uint8x16_t,
        round_key: uint8x16_t
      ) -> uint8x16_t {
        aes_single_round(state, round_key)
      }
    }

    Self {
      val: helper(state.simd, state.val, round_key.val),
      simd: state.simd,
    }
  }
}

impl<S: Simd> BitAnd for Neon1x128<S>
where
  S: HasNeon,
{
  type Output = Neon1x128<S>;

  #[inline(always)]
  fn bitand(self, rhs: Self) -> Self::Output {
    with_simd! {
      fn helper(
        token: impl HasNeon,
        first: uint8x16_t,
        second: uint8x16_t
      ) -> uint8x16_t {
        vandq_u8(first, second)
      }
    }

    Self {
      val: helper(self.simd, self.val, rhs.val),
      simd: self.simd,
    }
  }
}

impl<S: Simd> BitXor for Neon1x128<S>
where
  S: HasNeon,
{
  type Output = Neon1x128<S>;

  #[inline(always)]
  fn bitxor(self, rhs: Self) -> Self::Output {
    with_simd! {
      fn helper(
        token: impl HasNeon,
        first: uint8x16_t,
        second: uint8x16_t
      ) -> uint8x16_t {
        veorq_u8(first, second)
      }
    }

    Self {
      val: helper(self.simd, self.val, rhs.val),
      simd: self.simd,
    }
  }
}

/// A 256-bit SIMD vector using Neon.
#[derive(Clone, Copy)]
pub struct Neon2x128<S: Simd> {
  val: [uint8x16_t; 2],
  simd: S,
}

impl<S: Simd> Block for Neon2x128<S>
where
  S: HasNeon,
{
  gen_shared_block256!(S);

  fn xor_down(self) -> Array128 {
    with_simd! {
      fn helper(
        token: impl HasNeon,
        input: [uint8x16_t; 2]
      ) -> Array128 {
        transmute!(veorq_u8(input[0], input[1]))
      }
    }

    helper(self.simd, self.val)
  }

  fn aes_encrypt_round(state: Self, round_key: Self) -> Self {
    with_simd! {
      fn helper(
        token: impl HasNeon,
        state: [uint8x16_t; 2],
        round_key: [uint8x16_t; 2],
      ) -> [uint8x16_t; 2] {
        [
          aes_single_round(state[0], round_key[0]),
          aes_single_round(state[1], round_key[1]),
        ]
      }
    }

    Self {
      val: helper(state.simd, state.val, round_key.val),
      simd: state.simd,
    }
  }
}

impl<S: Simd> BitAnd for Neon2x128<S>
where
  S: HasNeon,
{
  type Output = Neon2x128<S>;

  #[inline(always)]
  fn bitand(self, rhs: Self) -> Self::Output {
    with_simd! {
      fn helper(
        token: impl HasNeon,
        first: [uint8x16_t; 2],
        second: [uint8x16_t; 2],
      ) -> [uint8x16_t; 2]{
        [
          vandq_u8(first[0], second[0]),
          vandq_u8(first[1], second[1]),
        ]
      }
    }

    Self {
      val: helper(self.simd, self.val, rhs.val),
      simd: self.simd,
    }
  }
}

impl<S: Simd> BitXor for Neon2x128<S>
where
  S: HasNeon,
{
  type Output = Neon2x128<S>;

  #[inline(always)]
  fn bitxor(self, rhs: Self) -> Self::Output {
    with_simd! {
      fn helper(
        token: impl HasNeon,
        first: [uint8x16_t; 2],
        second: [uint8x16_t; 2],
      ) -> [uint8x16_t; 2]{
        [
          veorq_u8(first[0], second[0]),
          veorq_u8(first[1], second[1]),
        ]
      }
    }

    Self {
      val: helper(self.simd, self.val, rhs.val),
      simd: self.simd,
    }
  }
}

/// A 512-bit SIMD vector using Neon.
#[derive(Clone, Copy)]
pub struct Neon4x128<S: Simd> {
  val: [uint8x16_t; 4],
  simd: S,
}

impl<S: Simd> Block for Neon4x128<S>
where
  S: HasNeon,
{
  gen_shared_block512!(S);

  fn xor_down(self) -> Array128 {
    with_simd! {
      fn helper(
        token: impl HasNeon,
        input: [uint8x16_t; 4]
      ) -> Array128 {
        // 512 bit vec can be seen as 4x128: (a, b, c, d)
        // (a ^ c) ^ (b ^ d) == a ^ b ^ c ^ d
        let temp1 = veorq_u8(input[0], input[1]);
        let temp2 = veorq_u8(input[2], input[3]);
        let out = veorq_u8(temp1, temp2);
        transmute!(out)
      }
    }

    helper(self.simd, self.val)
  }

  fn aes_encrypt_round(state: Self, round_key: Self) -> Self {
    with_simd! {
      fn helper(
        token: impl HasNeon,
        state: [uint8x16_t; 4],
        round_key: [uint8x16_t; 4],
      ) -> [uint8x16_t; 4] {
        [
          aes_single_round(state[0], round_key[0]),
          aes_single_round(state[1], round_key[1]),
          aes_single_round(state[2], round_key[2]),
          aes_single_round(state[3], round_key[3]),
        ]
      }
    }

    Self {
      val: helper(state.simd, state.val, round_key.val),
      simd: state.simd,
    }
  }
}

impl<S: Simd> BitAnd for Neon4x128<S>
where
  S: HasNeon,
{
  type Output = Neon4x128<S>;

  #[inline(always)]
  fn bitand(self, rhs: Self) -> Self::Output {
    with_simd! {
      fn helper(
        token: impl HasNeon,
        first: [uint8x16_t; 4],
        second: [uint8x16_t; 4],
      ) -> [uint8x16_t; 4]{
        [
          vandq_u8(first[0], second[0]),
          vandq_u8(first[1], second[1]),
          vandq_u8(first[2], second[2]),
          vandq_u8(first[3], second[3]),
        ]
      }
    }

    Self {
      val: helper(self.simd, self.val, rhs.val),
      simd: self.simd,
    }
  }
}

impl<S: Simd> BitXor for Neon4x128<S>
where
  S: HasNeon,
{
  type Output = Neon4x128<S>;

  #[inline(always)]
  fn bitxor(self, rhs: Self) -> Self::Output {
    with_simd! {
      fn helper(
        token: impl HasNeon,
        first: [uint8x16_t; 4],
        second: [uint8x16_t; 4],
      ) -> [uint8x16_t; 4]{
        [
          veorq_u8(first[0], second[0]),
          veorq_u8(first[1], second[1]),
          veorq_u8(first[2], second[2]),
          veorq_u8(first[3], second[3]),
        ]
      }
    }

    Self {
      val: helper(self.simd, self.val, rhs.val),
      simd: self.simd,
    }
  }
}

#[inline]
#[target_feature(enable = "aes")]
pub fn aes_single_round(
  state: uint8x16_t,
  round_key: uint8x16_t,
) -> uint8x16_t {
  // ARM's AES instructions are... complicated.
  //
  // Intel's AESENC (`aes_single_round`) instruction does this:
  //   ShiftRows → SubBytes → MixColumns → AddRoundKey
  //
  // Awesome, great, A++, tyvm Intel.
  //
  // ARM's AESE (`vaeseq_u8`) instruction does this:
  //   AddRoundKey → SubBytes → ShiftRows
  //
  // Notice that AddRoundKey is first and that MixColumns is missing.
  // MixColumns is a separate AESMC (`vaesmcq_u8`) instruction.
  //
  // ARM's more granular approach is annoying, but fine.
  // AddRoundKey (XOR with key) being first instead of last is also fine when
  // implementing actual AES which strings together multiple rounds in a row,
  // but AEGIS doesn't quite work that way.
  //
  // So we end up emulating Intel's AESENC here:
  //
  //  vaeseq_u8(state, 0) -> vaesmcq_u8 -> XOR with `round_key`
  //
  // Passing all-zero for key in `vaeseq_u8` "nullifies" the initial
  // AddRoundKey (XOR with 0 does nothing). We manually XOR with the key at the
  // end.

  // `vdupq_n_u8(0)` is a single cycle register op with no memory dependency.
  // It cannot be optimized further!
  let zero = vdupq_n_u8(0);

  veorq_u8(vaesmcq_u8(vaeseq_u8(state, zero)), round_key)
}