#***************************************************************************
# Written by Nir Drucker and Shay Gueron
# AWS Cryptographic Algorithms Group
# (ndrucker@amazon.com, gueron@amazon.com)
#
# Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#  
# Licensed under the Apache License, Version 2.0 (the "License").
# You may not use this file except in compliance with the License.
# A copy of the License is located at
#  
#     http://www.apache.org/licenses/LICENSE-2.0
#  
# or in the "license" file accompanying this file. This file is distributed 
# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either 
# express or implied. See the License for the specific language governing 
# permissions and limitations under the License.
# The license is detailed in the file LICENSE.txt, and applies to this file.
#***************************************************************************

.data

.align 16, 0x90
MASK1:
.long 0x0c0f0e0d,0x0c0f0e0d,0x0c0f0e0d,0x0c0f0e0d
CON1:
.long 1,1,1,1

.set k256_size, 32

.text

################################################################################
# void aes256_key_expansion(OUT aes256_ks_t* ks, IN const uint8_t* key);
# The output parameter must be 16 bytes aligned!
#
#Linux ABI
#define out %rdi
#define in  %rsi

#define CON      %xmm0
#define MASK_REG %xmm1

#define IN0      %xmm2
#define IN1      %xmm3

#define TMP1     %xmm4
#define TMP2     %xmm5

#define ZERO     %xmm15

.macro ROUND1 in0 in1
    add $k256_size, out
    vpshufb     MASK_REG, \in1, TMP2
    vaesenclast CON,  TMP2, TMP2
    vpslld      $1,   CON,  CON
    vpslldq     $4,   \in0, TMP1
    vpxor       TMP1, \in0, \in0
    vpslldq     $4,   TMP1, TMP1
    vpxor       TMP1, \in0, \in0
    vpslldq     $4,   TMP1, TMP1       
    vpxor       TMP1, \in0, \in0
    vpxor       TMP2, \in0, \in0
    vmovdqa     \in0, (out)
.endm

.macro ROUND2
   vpshufd     $0xff, IN0,  TMP2
   vaesenclast ZERO,  TMP2, TMP2
   vpslldq     $4,    IN1,  TMP1
   vpxor       TMP1,  IN1,  IN1
   vpslldq     $4,    TMP1, TMP1
   vpxor       TMP1,  IN1,  IN1
   vpslldq     $4,    TMP1, TMP1
   vpxor       TMP1,  IN1,  IN1
   vpxor       TMP2,  IN1,  IN1
   vmovdqa     IN1,   16(out)
.endm

.type   aes256_key_expansion,@function
.hidden aes256_key_expansion
.globl  aes256_key_expansion
aes256_key_expansion:
   vmovdqu (in),   IN0
   vmovdqu 16(in), IN1
   vmovdqa IN0,    (out)
   vmovdqa IN1,  16(out)

   vmovdqa CON1(%rip), CON
   vmovdqa MASK1(%rip), MASK_REG

   vpxor ZERO, ZERO, ZERO

   mov $6, %ax
.loop256:

   ROUND1 IN0 IN1
   dec %ax
   ROUND2
   jne .loop256

   ROUND1 IN0 IN1

   ret
.size aes256_key_expansion, .-aes256_key_expansion