philbin 1.0.1

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

// NOTE: We use `#[inline]` directives very sparingly in this module as there
// are no assembly intrinsics used here and we want to avoid code bloat.
// Definitely avoid any `#[inline(always)]`.
//
// In general, let the compiler decide what should be inlined for the
// software-only implementation.

/// A 128-bit software-emulated "SIMD" vector.
#[derive(Clone, Copy)]
pub struct Fallback128<S: Simd> {
  val: Array128,
  // We don't actually need this field, but:
  //  - The shared code in the `gen_shared_block*` macros expects it.
  //  - Block needs the generic param, so at best we'd have to use
  //    `PhantomData<S>` to avoid an unconstrained generic param in the impl
  //    blocks.
  // `Simd` impl structs are all zero bytes so there's no cost to having it.
  simd: S,
}

impl<S: Simd> Block for Fallback128<S> {
  gen_shared_block128!(S);

  fn aes_encrypt_round(state: Self, round_key: Self) -> Self {
    Self {
      val: aes_single_round(state.val, round_key.val),
      simd: state.simd,
    }
  }
}

impl<S: Simd> BitAnd for Fallback128<S> {
  type Output = Fallback128<S>;

  fn bitand(self, rhs: Self) -> Self::Output {
    Self {
      val: and128(self.val, rhs.val),
      simd: self.simd,
    }
  }
}

impl<S: Simd> BitXor for Fallback128<S> {
  type Output = Fallback128<S>;

  fn bitxor(self, rhs: Self) -> Self::Output {
    Self {
      val: xor128(self.val, rhs.val),
      simd: self.simd,
    }
  }
}

/// A 256-bit software-emulated "SIMD" vector.
#[derive(Clone, Copy)]
pub struct Fallback256<S: Simd> {
  val: [Array128; 2],
  simd: S,
}

impl<S: Simd> Block for Fallback256<S> {
  gen_shared_block256!(S);

  fn xor_down(self) -> Array128 {
    xor128(self.val[0], self.val[1])
  }

  fn aes_encrypt_round(state: Self, round_key: Self) -> Self {
    Self {
      val: [
        aes_single_round(state.val[0], round_key.val[0]),
        aes_single_round(state.val[1], round_key.val[1]),
      ],
      simd: state.simd,
    }
  }
}

impl<S: Simd> BitAnd for Fallback256<S> {
  type Output = Fallback256<S>;

  fn bitand(self, rhs: Self) -> Self::Output {
    Self {
      val: [
        and128(self.val[0], rhs.val[0]),
        and128(self.val[1], rhs.val[1]),
      ],
      simd: self.simd,
    }
  }
}

impl<S: Simd> BitXor for Fallback256<S> {
  type Output = Fallback256<S>;

  fn bitxor(self, rhs: Self) -> Self::Output {
    Self {
      val: [
        xor128(self.val[0], rhs.val[0]),
        xor128(self.val[1], rhs.val[1]),
      ],
      simd: self.simd,
    }
  }
}

/// A 512-bit software-emulated "SIMD" vector.
#[derive(Clone, Copy)]
pub struct Fallback512<S: Simd> {
  val: [Array128; 4],
  simd: S,
}

impl<S: Simd> Block for Fallback512<S> {
  gen_shared_block512!(S);

  fn xor_down(self) -> Array128 {
    // 512 bit vec can be seen as 4x128: (a, b, c, d)
    // (a ^ c) ^ (b ^ d) == a ^ b ^ c ^ d
    let temp1 = xor128(self.val[0], self.val[1]);
    let temp2 = xor128(self.val[2], self.val[3]);
    xor128(temp1, temp2)
  }

  fn aes_encrypt_round(state: Self, round_key: Self) -> Self {
    Self {
      val: [
        aes_single_round(state.val[0], round_key.val[0]),
        aes_single_round(state.val[1], round_key.val[1]),
        aes_single_round(state.val[2], round_key.val[2]),
        aes_single_round(state.val[3], round_key.val[3]),
      ],
      simd: state.simd,
    }
  }
}

impl<S: Simd> BitAnd for Fallback512<S> {
  type Output = Fallback512<S>;

  fn bitand(self, rhs: Self) -> Self::Output {
    Self {
      val: [
        and128(self.val[0], rhs.val[0]),
        and128(self.val[1], rhs.val[1]),
        and128(self.val[2], rhs.val[2]),
        and128(self.val[3], rhs.val[3]),
      ],
      simd: self.simd,
    }
  }
}

impl<S: Simd> BitXor for Fallback512<S> {
  type Output = Fallback512<S>;

  fn bitxor(self, rhs: Self) -> Self::Output {
    Self {
      val: [
        xor128(self.val[0], rhs.val[0]),
        xor128(self.val[1], rhs.val[1]),
        xor128(self.val[2], rhs.val[2]),
        xor128(self.val[3], rhs.val[3]),
      ],
      simd: self.simd,
    }
  }
}

#[inline]
fn xor128(first: Array128, second: Array128) -> Array128 {
  std::array::from_fn(|i| first[i] ^ second[i])
}

#[inline]
fn and128(first: Array128, second: Array128) -> Array128 {
  std::array::from_fn(|i| first[i] & second[i])
}

fn aes_single_round(state: Array128, round_key: Array128) -> Array128 {
  let mut block_state = state.into();
  aes::hazmat::cipher_round(&mut block_state, &round_key.into());
  block_state.into()
}