content: -----------------------------------------------------------------------------
  -- |
  -- Module    : Documentation.SBV.Examples.Crypto.AES
  -- Copyright : (c) Levent Erkok
  -- License   : BSD3
  -- Maintainer: erkokl@gmail.com
  -- Stability : experimental
  --
  -- An implementation of AES (Advanced Encryption Standard), using SBV.
  -- For details on AES, see <http://en.wikipedia.org/wiki/Advanced_Encryption_Standard>.
  --
  -- We do a T-box implementation, which leads to good C code as we can take
  -- advantage of look-up tables. Note that we make virtually no attempt to
  -- optimize our Haskell code. The concern here is not with getting Haskell running
  -- fast at all. The idea is to program the T-Box implementation as naturally and clearly
  -- as possible in Haskell, and have SBV's code-generator generate fast C code automatically.
  -- Therefore, we merely use ordinary Haskell lists as our data-structures, and do not
  -- bother with any unboxing or strictness annotations. Thus, we achieve the separation
  -- of concerns: Correctness via clarity and simplicity and proofs on the Haskell side,
  -- performance by relying on SBV's code generator. If necessary, the generated code
  -- can be FFI'd back into Haskell to complete the loop.
  --
  -- All 3 valid key sizes (128, 192, and 256) as required by the FIPS-197 standard
  -- are supported.
  -----------------------------------------------------------------------------
  
  {-# LANGUAGE CPP              #-}
  {-# LANGUAGE DataKinds        #-}
  {-# LANGUAGE ParallelListComp #-}
  {-# LANGUAGE TypeApplications #-}
  
  #if MIN_VERSION_base(4,19,0)
  {-# OPTIONS_GHC -Wall -Werror -Wno-incomplete-uni-patterns -Wno-x-partial #-}
  #else
  {-# OPTIONS_GHC -Wall -Werror -Wno-incomplete-uni-patterns #-}
  #endif
  
  module Documentation.SBV.Examples.Crypto.AES where
  
  import Control.Monad (void, when)
  
  import Data.SBV
  import Data.SBV.Tools.CodeGen
  import Data.SBV.Tools.Polynomial
  
  import Data.List (transpose)
  import Data.Maybe (fromJust)
  
  import Data.Proxy
  
  import Numeric (showHex)
  
  import Test.QuickCheck hiding (verbose)
  
  #ifdef DOCTEST
  -- $setup
  -- >>> import Data.SBV
  #endif
  
  -----------------------------------------------------------------------------
  -- * Formalizing GF(2^8)
  -----------------------------------------------------------------------------
  
  -- | An element of the Galois Field 2^8, which are essentially polynomials with
  -- maximum degree 7. They are conveniently represented as values between 0 and 255.
  type GF28 = SWord 8
  
  -- | Multiplication in GF(2^8). This is simple polynomial multiplication, followed
  -- by the irreducible polynomial @x^8+x^4+x^3+x^1+1@. We simply use the 'pMult'
  -- function exported by SBV to do the operation.
  gf28Mult :: GF28 -> GF28 -> GF28
  gf28Mult x y = pMult (x, y, [8, 4, 3, 1, 0])
  
  -- | Exponentiation by a constant in GF(2^8). The implementation uses the usual
  -- square-and-multiply trick to speed up the computation.
  gf28Pow :: GF28 -> Int -> GF28
  gf28Pow n = pow
    where sq x  = x `gf28Mult` x
          pow 0    = 1
          pow i
           | odd i = n `gf28Mult` sq (pow (i `shiftR` 1))
           | True  = sq (pow (i `shiftR` 1))
  
  -- | Computing inverses in GF(2^8). By the mathematical properties of GF(2^8)
  -- and the particular irreducible polynomial used @x^8+x^5+x^3+x^1+1@, it
  -- turns out that raising to the 254 power gives us the multiplicative inverse.
  -- Of course, we can prove this using SBV:
  --
  -- >>> prove $ \x -> x ./= 0 .=> x `gf28Mult` gf28Inverse x .== 1
  -- Q.E.D.
  --
  -- Note that we exclude @0@ in our theorem, as it does not have a
  -- multiplicative inverse.
  gf28Inverse :: GF28 -> GF28
  gf28Inverse x = x `gf28Pow` 254
  
  -----------------------------------------------------------------------------
  -- * Implementing AES
  -----------------------------------------------------------------------------
  
  -----------------------------------------------------------------------------
  -- ** Types and basic operations
  -----------------------------------------------------------------------------
  -- | AES state. The state consists of four 32-bit words, each of which is in turn treated
  -- as four GF28's, i.e., 4 bytes. The T-Box implementation keeps the four-bytes together
  -- for efficient representation.
  type State = [SWord 32]
  
  -- | The key, which can be 128, 192, or 256 bits. Represented as a sequence of 32-bit words.
  type Key = [SWord 32]
  
  -- | The key schedule. AES executes in rounds, and it treats first and last round keys slightly
  -- differently than the middle ones. We reflect that choice by being explicit about it in our type.
  -- The length of the middle list of keys depends on the key-size, which in turn determines
  -- the number of rounds.
  type KS = (Key, [Key], Key)
  
  -- | Rotating a state row by a fixed amount to the right.
  rotR :: [GF28] -> Int -> [GF28]
  rotR [a, b, c, d] 1 = [d, a, b, c]
  rotR [a, b, c, d] 2 = [c, d, a, b]
  rotR [a, b, c, d] 3 = [b, c, d, a]
  rotR xs           i = error $ "rotR: Unexpected input: " ++ show (xs, i)
  
  -----------------------------------------------------------------------------
  -- ** The key schedule
  -----------------------------------------------------------------------------
  
  -- | Definition of round-constants, as specified in Section 5.2 of the AES standard.
  roundConstants :: [GF28]
  roundConstants = 0 : [ gf28Pow 2 (k-1) | k <- [1 .. ] ]
  
  -- | The @InvMixColumns@ transformation, as described in Section 5.3.3 of the standard. Note
  -- that this transformation is only used explicitly during key-expansion in the T-Box implementation
  -- of AES.
  invMixColumns :: State -> State
  invMixColumns state = map fromBytes $ transpose $ mmult (map toBytes state)
   where dot f   = foldr1 xor . zipWith ($) f
         mmult n = [map (dot r) n | r <- [ [mE, mB, mD, m9]
                                         , [m9, mE, mB, mD]
                                         , [mD, m9, mE, mB]
                                         , [mB, mD, m9, mE]
                                         ]]
         -- table-lookup versions of gf28Mult with the constants used in invMixColumns
         mE = select mETable 0
         mB = select mBTable 0
         mD = select mDTable 0
         m9 = select m9Table 0
         mETable = map (gf28Mult 0xE . literal) [0..255]
         mBTable = map (gf28Mult 0xB . literal) [0..255]
         mDTable = map (gf28Mult 0xD . literal) [0..255]
         m9Table = map (gf28Mult 0x9 . literal) [0..255]
  
  -- | Key expansion. Starting with the given key, returns an infinite sequence of
  -- words, as described by the AES standard, Section 5.2, Figure 11.
  keyExpansion :: Int -> Key -> [Key]
  keyExpansion nk key = chop4 keys
     where keys :: [SWord 32]
           keys = key ++ [nextWord i prev old | i <- [nk ..] | prev <- drop (nk-1) keys | old <- keys]
  
           nextWord :: Int -> SWord 32 -> SWord 32 -> SWord 32
           nextWord i prev old
             | i `mod` nk == 0           = old `xor` subWordRcon (prev `rotateL` 8) (roundConstants !! (i `div` nk))
             | i `mod` nk == 4 && nk > 6 = old `xor` subWordRcon prev 0
             | True                      = old `xor` prev
  
           subWordRcon :: SWord 32 -> GF28 -> SWord 32
           subWordRcon w rc = fromBytes [a `xor` rc, b, c, d]
              where [a, b, c, d] = map sbox $ toBytes w
  
  -----------------------------------------------------------------------------
  -- ** The S-box transformation
  -----------------------------------------------------------------------------
  
  -- | The values of the AES S-box table. Note that we describe the S-box programmatically
  -- using the mathematical construction given in Section 5.1.1 of the standard. However,
  -- the code-generation will turn this into a mere look-up table, as it is just a
  -- constant table, all computation being done at \"compile-time\".
  sboxTable :: [GF28]
  sboxTable = [xformByte (gf28Inverse (literal b)) | b <- [0 .. 255]]
    where xformByte :: GF28 -> GF28
          xformByte b = foldr xor 0x63 [b `rotateR` i | i <- [0, 4, 5, 6, 7]]
  
  -- | The sbox transformation. We simply select from the sbox table. Note that we
  -- are obliged to give a default value (here @0@) to be used if the index is out-of-bounds
  -- as required by SBV's 'select' function. However, that will never happen since
  -- the table has all 256 elements in it.
  sbox :: GF28 -> GF28
  sbox = select sboxTable 0
  
  -----------------------------------------------------------------------------
  -- ** The inverse S-box transformation
  -----------------------------------------------------------------------------
  
  -- | The values of the inverse S-box table. Again, the construction is programmatic.
  unSBoxTable :: [GF28]
  unSBoxTable = [gf28Inverse (xformByte (literal b)) | b <- [0 .. 255]]
    where xformByte :: GF28 -> GF28
          xformByte b = foldr xor 0x05 [b `rotateR` i | i <- [2, 5, 7]]
  
  -- | The inverse s-box transformation.
  unSBox :: GF28 -> GF28
  unSBox = select unSBoxTable 0
  
  -- | Prove that the 'sbox' and 'unSBox' are inverses. We have:
  --
  -- >>> prove sboxInverseCorrect
  -- Q.E.D.
  --
  sboxInverseCorrect :: GF28 -> SBool
  sboxInverseCorrect x = unSBox (sbox x) .== x .&& sbox (unSBox x) .== x
  
  -----------------------------------------------------------------------------
  -- ** AddRoundKey transformation
  -----------------------------------------------------------------------------
  
  -- | Adding the round-key to the current state. We simply exploit the fact
  -- that addition is just xor in implementing this transformation.
  addRoundKey :: Key -> State -> State
  addRoundKey = zipWith xor
  
  -----------------------------------------------------------------------------
  -- ** Tables for T-Box encryption
  -----------------------------------------------------------------------------
  
  -- | T-box table generation function for encryption
  t0Func :: GF28 -> [GF28]
  t0Func a = [s `gf28Mult` 2, s, s, s `gf28Mult` 3] where s = sbox a
  
  -- | First look-up table used in encryption
  t0 :: GF28 -> SWord 32
  t0 = select t0Table 0 where t0Table = [fromBytes (t0Func a)          | a <- map literal [0..255]]
  
  -- | Second look-up table used in encryption
  t1 :: GF28 -> SWord 32
  t1 = select t1Table 0 where t1Table = [fromBytes (t0Func a `rotR` 1) | a <- map literal [0..255]]
  
  -- | Third look-up table used in encryption
  t2 :: GF28 -> SWord 32
  t2 = select t2Table 0 where t2Table = [fromBytes (t0Func a `rotR` 2) | a <- map literal [0..255]]
  
  -- | Fourth look-up table used in encryption
  t3 :: GF28 -> SWord 32
  t3 = select t3Table 0 where t3Table = [fromBytes (t0Func a `rotR` 3) | a <- map literal [0..255]]
  
  -----------------------------------------------------------------------------
  -- ** Tables for T-Box decryption
  -----------------------------------------------------------------------------
  
  -- | T-box table generating function for decryption
  u0Func :: GF28 -> [GF28]
  u0Func a = [s `gf28Mult` 0xE, s `gf28Mult` 0x9, s `gf28Mult` 0xD, s `gf28Mult` 0xB] where s = unSBox a
  
  -- | First look-up table used in decryption
  u0 :: GF28 -> SWord 32
  u0 = select t0Table 0 where t0Table = [fromBytes (u0Func a)          | a <- map literal [0..255]]
  
  -- | Second look-up table used in decryption
  u1 :: GF28 -> SWord 32
  u1 = select t1Table 0 where t1Table = [fromBytes (u0Func a `rotR` 1) | a <- map literal [0..255]]
  
  -- | Third look-up table used in decryption
  u2 :: GF28 -> SWord 32
  u2 = select t2Table 0 where t2Table = [fromBytes (u0Func a `rotR` 2) | a <- map literal [0..255]]
  
  -- | Fourth look-up table used in decryption
  u3 :: GF28 -> SWord 32
  u3 = select t3Table 0 where t3Table = [fromBytes (u0Func a `rotR` 3) | a <- map literal [0..255]]
  
  -----------------------------------------------------------------------------
  -- ** AES rounds
  -----------------------------------------------------------------------------
  
  -- | Generic round function. Given the function to perform one round, a key-schedule,
  -- and a starting state, it performs the AES rounds.
  doRounds :: (Bool -> State -> Key -> State) -> KS -> State -> State
  doRounds rnd (ikey, rkeys, fkey) sIn = rnd True (last rs) fkey
    where s0 = ikey `addRoundKey` sIn
          rs = s0 : [rnd False s k | s <- rs | k <- rkeys ]
  
  -- | One encryption round. The first argument indicates whether this is the final round
  -- or not, in which case the construction is slightly different.
  aesRound :: Bool -> State -> Key -> State
  aesRound isFinal s key = d `addRoundKey` key
    where d = map (f isFinal) [0..3]
          a = map toBytes s
          f True j = fromBytes [ sbox (a !! ((j+0) `mod` 4) !! 0)
                               , sbox (a !! ((j+1) `mod` 4) !! 1)
                               , sbox (a !! ((j+2) `mod` 4) !! 2)
                               , sbox (a !! ((j+3) `mod` 4) !! 3)
                               ]
          f False j = e0 `xor` e1 `xor` e2 `xor` e3
                where e0 = t0 (a !! ((j+0) `mod` 4) !! 0)
                      e1 = t1 (a !! ((j+1) `mod` 4) !! 1)
                      e2 = t2 (a !! ((j+2) `mod` 4) !! 2)
                      e3 = t3 (a !! ((j+3) `mod` 4) !! 3)
  
  -- | One decryption round. Similar to the encryption round, the first argument
  -- indicates whether this is the final round or not.
  aesInvRound :: Bool -> State -> Key -> State
  aesInvRound isFinal s key = d `addRoundKey` key
    where d = map (f isFinal) [0..3]
          a = map toBytes s
          f True j = fromBytes [ unSBox (a !! ((j+0) `mod` 4) !! 0)
                               , unSBox (a !! ((j+3) `mod` 4) !! 1)
                               , unSBox (a !! ((j+2) `mod` 4) !! 2)
                               , unSBox (a !! ((j+1) `mod` 4) !! 3)
                               ]
          f False j = e0 `xor` e1 `xor` e2 `xor` e3
                where e0 = u0 (a !! ((j+0) `mod` 4) !! 0)
                      e1 = u1 (a !! ((j+3) `mod` 4) !! 1)
                      e2 = u2 (a !! ((j+2) `mod` 4) !! 2)
                      e3 = u3 (a !! ((j+1) `mod` 4) !! 3)
  
  -----------------------------------------------------------------------------
  -- * AES API
  -----------------------------------------------------------------------------
  
  -- | Key schedule. Given a 128, 192, or 256 bit key, expand it to get key-schedules
  -- for encryption and decryption. The key is given as a sequence of 32-bit words.
  -- (4 elements for 128-bits, 6 for 192, and 8 for 256.) Compare this function to 'aesInvKeySchedule'
  -- which can calculate the key-expansion for decryption on the fly, as opposed to calculating
  -- the forward key-expansion first.
  aesKeySchedule :: Key -> (KS, KS)
  aesKeySchedule key
    | nk `elem` [4, 6, 8]
    = (encKS, decKS)
    | True
    = error "aesKeySchedule: Invalid key size"
    where nk = length key
          nr = nk + 6
          encKS@(f, m, l) = (head rKeys, take (nr-1) (tail rKeys), rKeys !! nr)
          decKS = (l, map invMixColumns (reverse m), f)
          rKeys = keyExpansion nk key
  
  -- | Block encryption. The first argument is the plain-text, which must have
  -- precisely 4 elements, for a total of 128-bits of input. The second
  -- argument is the key-schedule to be used, obtained by a call to 'aesKeySchedule'.
  -- The output will always have 4 32-bit words, which is the cipher-text.
  aesEncrypt :: [SWord 32] -> KS -> [SWord 32]
  aesEncrypt pt encKS
    | length pt == 4
    = doRounds aesRound encKS pt
    | True
    = error "aesEncrypt: Invalid plain-text size"
  
  -- | Block decryption. The arguments are the same as in 'aesEncrypt', except
  -- the first argument is the cipher-text and the output is the corresponding
  -- plain-text.
  aesDecrypt :: [SWord 32] -> KS -> [SWord 32]
  aesDecrypt ct decKS
    | length ct == 4
    = doRounds aesInvRound decKS ct
    | True
    = error "aesDecrypt: Invalid cipher-text size"
  
  -----------------------------------------------------------------------------
  -- * On-the-fly decryption
  -- ${ontheflyintro}
  -----------------------------------------------------------------------------
  {- $ontheflyintro
     While regular encryption can be fused with key-generation, the standard method of AES
     decryption has to perform the key-expansion before decryption starts. This can be undesirable
     as it necessarily serializes the action of key-expansion before decryption. An
     alternative is to do on-the-fly decryption: We can expand the key in reverse, and thus
     need not save the key-schedule. One downside of this approach, however, is that we need
     to keep the "unwound" key: That is, instead of the common key used for encryption and
     decryption, we need to hold on to the final value of key-expansion, so it can be run
     in reverse. In this section, we implement on-the-fly decryption using this idea.
  -}
  
  -- | Inverse key expansion. Starting from the final round key, unwinds key generation operation
  -- to construct keys for the previous rounds. Used in on-the-fly decryption.
  invKeyExpansion :: Int -> Key -> [Key]
  invKeyExpansion nk rkey = map reverse (chop4 keys)
     where keys :: [SWord 32]
           keys = rkey ++ [invNextWord i prev old | i <- reverse [0 .. remaining - 1] | prev <- drop 1 keys | old <- keys]
  
           totalWords = 4 * (nk + 6 + 1)
           remaining  = totalWords - nk
  
           invNextWord :: Int -> SWord 32 -> SWord 32 -> SWord 32
           invNextWord i prev old
             | i `mod` nk == 0           = old `xor` subWordRcon (prev `rotateL` 8) (roundConstants !! (1 + i `div` nk))
             | i `mod` nk == 4 && nk > 6 = old `xor` subWordRcon prev 0
             | True                      = old `xor` prev
  
           subWordRcon :: SWord 32 -> GF28 -> SWord 32
           subWordRcon w rc = fromBytes [a `xor` rc, b, c, d]
              where [a, b, c, d] = map sbox $ toBytes w
  
  -- | AES inverse key schedule. Starting from the last-round key, construct the sequence of keys
  -- that can be used for doing on-the-fly decryption. Compare this function to 'aesKeySchedule' which
  -- returns both encryption and decryption schedules: In this case, we don't calculate the encryption
  -- sequence, hence we can fuse this function with the decryption operation.
  aesInvKeySchedule :: Key -> KS
  aesInvKeySchedule key
    | nk `elem` [4, 6, 8]
    = decKS
    | True
    = error "aesInvKeySchedule: Invalid key size"
    where nk = length key
          nr = nk + 6
          decKS = (head rKeys, take (nr-1) (tail rKeys), rKeys !! nr)
          rKeys = invKeyExpansion nk key
  
  -- | Block decryption, starting from the unwound key. That is, start from the final key.
  -- Also; we don't use the T-box implementation. Just pure AES inverse cipher.
  aesDecryptUnwoundKey :: [SWord 32] -> KS -> [SWord 32]
  aesDecryptUnwoundKey ct decKS
    | length ct == 4
    = doRounds aesInvRoundRegular decKS ct
    | True
    = error "aesDecrypt: Invalid cipher-text size"
    where aesInvRoundRegular isFinal s key = u
            where u :: State
                  u = map (f isFinal) [0 .. 3]
                    where a   = map toBytes s
                          kbs = map toBytes key
                          f True j = fromBytes [ unSBox (a !! ((j+0) `mod` 4) !! 0)
                                               , unSBox (a !! ((j+3) `mod` 4) !! 1)
                                               , unSBox (a !! ((j+2) `mod` 4) !! 2)
                                               , unSBox (a !! ((j+1) `mod` 4) !! 3)
                                               ] `xor` (key !! j)
                          f False j = e0 `xor` e1 `xor` e2 `xor` e3
                                where e0 = otfU0 $ unSBox (a !! ((j+0) `mod` 4) !! 0) `xor` (kbs !! j !! 0)
                                      e1 = otfU1 $ unSBox (a !! ((j+3) `mod` 4) !! 1) `xor` (kbs !! j !! 1)
                                      e2 = otfU2 $ unSBox (a !! ((j+2) `mod` 4) !! 2) `xor` (kbs !! j !! 2)
                                      e3 = otfU3 $ unSBox (a !! ((j+1) `mod` 4) !! 3) `xor` (kbs !! j !! 3)
  
                  otfU0Func b = [b `gf28Mult` 0xE, b `gf28Mult` 0x9, b `gf28Mult` 0xD, b `gf28Mult` 0xB]
                  otfU0 = select t0Table 0 where t0Table = [fromBytes (otfU0Func a)          | a <- map literal [0..255]]
                  otfU1 = select t1Table 0 where t1Table = [fromBytes (otfU0Func a `rotR` 1) | a <- map literal [0..255]]
                  otfU2 = select t2Table 0 where t2Table = [fromBytes (otfU0Func a `rotR` 2) | a <- map literal [0..255]]
                  otfU3 = select t3Table 0 where t3Table = [fromBytes (otfU0Func a `rotR` 3) | a <- map literal [0..255]]
  
  -----------------------------------------------------------------------------
  -- * Test vectors
  -----------------------------------------------------------------------------
  
  -- | Common plain text for test vectors
  commonPT :: [SWord 32]
  commonPT = [0x00112233, 0x44556677, 0x8899aabb, 0xccddeeff]
  
  -- | Key for 128-bit encryption test
  aes128Key :: Key
  aes128Key = [0x00010203, 0x04050607, 0x08090a0b, 0x0c0d0e0f]
  
  -- | Key for 192-bit encryption test
  aes192Key :: Key
  aes192Key = aes128Key ++ [0x10111213, 0x14151617]
  
  -- | Key for 256-bit encryption test
  aes256Key :: Key
  aes256Key = aes192Key ++ [0x18191a1b, 0x1c1d1e1f]
  
  -- | Expected cipher-text for 128-bit encryption
  aes128CT :: [SWord 32]
  aes128CT = [0x69c4e0d8, 0x6a7b0430, 0xd8cdb780, 0x70b4c55a]
  
  -- | Expected cipher-text for 192-bit encryption
  aes192CT :: [SWord 32]
  aes192CT = [0xdda97ca4, 0x864cdfe0, 0x6eaf70a0, 0xec0d7191]
  
  -- | Expected cipher-text for 256-bit encryption
  aes256CT :: [SWord 32]
  aes256CT = [0x8ea2b7ca, 0x516745bf, 0xeafc4990, 0x4b496089]
  
  -- | Calculate the 128-bit final-round key from on-the-fly decryption key schedule
  aes128InvKey :: Key
  aes128InvKey = extractFinalKey aes128Key
  
  -- | Calculate the 192-bit final-round key from on-the-fly decryption key schedule
  aes192InvKey :: Key
  aes192InvKey = extractFinalKey aes192Key
  
  -- | Calculate the 192-bit final-round key from on-the-fly decryption key schedule. Compare this
  -- to 'aes192InvKey': Typically we just need the final 6-blocks, but it is advantageous to have
  -- the entire last 8-blocks even for 192-bit keys. That is,  e store the final 256-bits of key-expansion
  -- for speed purposes for both 192 and 256 bit versions. (But only the final 128 bits for the 128-bit version.)
  aes192InvKeyExtended :: Key
  aes192InvKeyExtended = extractFinalKeyExtended aes192Key
  
  -- | Calculate the 256-bit final-round key from on-the-fly decryption key schedule
  aes256InvKey :: Key
  aes256InvKey = extractFinalKey aes256Key
  
  -- | Extract the final key for on-the-fly decryption. This will extract exactly the number of blocks we need.
  extractFinalKey :: [SWord 32] -> [SWord 32]
  extractFinalKey initKey = take nk (extractFinalKeyExtended initKey)
    where nk = length initKey
  
  -- | Extract the extended key for on-the-fly decryption. This will extract 4-blocks for 128-bit decryption,
  -- but 256 bit for both 192 and 256-bit variants
  extractFinalKeyExtended :: [SWord 32] -> [SWord 32]
  extractFinalKeyExtended initKey = take feed (concatMap reverse (chop4 (take feed roundKeys)))
    where nk             = length initKey
          feed | nk == 4 = 4
               | True    = 8
  
          ((f, m, l), _) = aesKeySchedule initKey
          roundKeys      = l ++ concat (reverse m) ++ f
  
  -----------------------------------------------------------------------------
  -- ** 128-bit enc/dec test
  -----------------------------------------------------------------------------
  
  -- | 128-bit encryption test, from Appendix C.1 of the AES standard:
  --
  -- >>> map hex8 t128Enc
  -- ["69c4e0d8","6a7b0430","d8cdb780","70b4c55a"]
  t128Enc :: [SWord 32]
  t128Enc = aesEncrypt commonPT ks
    where (ks, _) = aesKeySchedule aes128Key
  
  -- | 128-bit decryption test, from Appendix C.1 of the AES standard:
  --
  -- >>> map hex8 t128Dec
  -- ["00112233","44556677","8899aabb","ccddeeff"]
  t128Dec :: [SWord 32]
  t128Dec = aesDecrypt aes128CT ks
    where (_, ks) = aesKeySchedule aes128Key
  
  -----------------------------------------------------------------------------
  -- ** 192-bit enc/dec test
  -----------------------------------------------------------------------------
  
  -- | 192-bit encryption test, from Appendix C.2 of the AES standard:
  --
  -- >>> map hex8 t192Enc
  -- ["dda97ca4","864cdfe0","6eaf70a0","ec0d7191"]
  t192Enc :: [SWord 32]
  t192Enc = aesEncrypt commonPT ks
    where (ks, _) = aesKeySchedule aes192Key
  
  -- | 192-bit decryption test, from Appendix C.2 of the AES standard:
  --
  -- >>> map hex8 t192Dec
  -- ["00112233","44556677","8899aabb","ccddeeff"]
  --
  t192Dec :: [SWord 32]
  t192Dec = aesDecrypt aes192CT ks
    where (_, ks) = aesKeySchedule aes192Key
  
  -----------------------------------------------------------------------------
  -- ** 256-bit enc/dec test
  -----------------------------------------------------------------------------
  
  -- | 256-bit encryption, from Appendix C.3 of the AES standard:
  --
  -- >>> map hex8 t256Enc
  -- ["8ea2b7ca","516745bf","eafc4990","4b496089"]
  t256Enc :: [SWord 32]
  t256Enc = aesEncrypt commonPT ks
    where (ks, _) = aesKeySchedule aes256Key
  
  -- | 256-bit decryption, from Appendix C.3 of the AES standard:
  --
  -- >>> map hex8 t256Dec
  -- ["00112233","44556677","8899aabb","ccddeeff"]
  t256Dec :: [SWord 32]
  t256Dec = aesDecrypt aes256CT ks
    where (_, ks) = aesKeySchedule aes256Key
  
  -- | Various tests for round-trip properties. We have:
  --
  -- >>> runAESTests False
  -- GOOD: Key generation AES128
  -- GOOD: Key generation AES192
  -- GOOD: Key generation AES256
  -- GOOD: Encryption     AES128
  -- GOOD: Decryption     AES128
  -- GOOD: Decryption-OTF AES128
  -- GOOD: Encryption     AES192
  -- GOOD: Decryption     AES192
  -- GOOD: Decryption-OTF AES192
  -- GOOD: Encryption     AES256
  -- GOOD: Decryption     AES256
  -- GOOD: Decryption-OTF AES256
  runAESTests :: Bool -> IO ()
  runAESTests runQC = do
                   testInvKeyExpansion
  
                   check "AES128" aes128Key aes128InvKey aes128CT
                   check "AES192" aes192Key aes192InvKey aes192CT
                   check "AES256" aes256Key aes256InvKey aes256CT
  
                   -- Quick-check tests are rather slow. So only run when requested.
                   when runQC $ do
                     putStrLn "Quick-check AES128 roundtrip" >> quickCheck roundTrip128
                     putStrLn "Quick-check AES192 roundtrip" >> quickCheck roundTrip192
                     putStrLn "Quick-check AES256 roundtrip" >> quickCheck roundTrip256
  
    where check :: String -> Key -> Key -> [SWord 32] -> IO ()
          check what key invKey ctExpected = do eq ("Encryption     " ++ what) ctExpected ctGot
                                                eq ("Decryption     " ++ what) commonPT   ptGot
                                                eq ("Decryption-OTF " ++ what) commonPT   ptGotInv
             where (encKS, decKS) = aesKeySchedule key
                   ctGot          = aesEncrypt           commonPT   encKS
                   ptGot          = aesDecrypt           ctExpected decKS
                   ptGotInv       = aesDecryptUnwoundKey ctExpected (aesInvKeySchedule invKey)
  
                   eq tag expected got
                     | length expected /= length got
                     = error $ unlines [ "BAD!: " ++ tag
                                       , "Comparing different sized lists:"
                                       , "Expected: " ++ show expected
                                       , "Got     : " ++ show got
                                       ]
                     | map extract expected == map extract got
                     = putStrLn $ "GOOD: " ++ tag
                     | True
                     = error $ unlines [ "BAD!: " ++ tag
                                       , "Expected: " ++ unwords (map hex8 expected)
                                       , "Got     : " ++ unwords (map hex8 got)
                                       ]
                    where extract :: SWord 32 -> Integer
                          extract = fromIntegral . fromJust . unliteral
  
          testInvKeyExpansion :: IO ()
          testInvKeyExpansion = do goTestInvKey "128" aes128Key
                                   goTestInvKey "192" aes192Key
                                   goTestInvKey "256" aes256Key
          goTestInvKey what k = do
            let nk = length k
                nr = nk + 6
  
                feed = case nk of
                         4 -> 4
                         _ -> 8
  
                ((f, m, l), _) = aesKeySchedule k
                required       = l ++ concat (reverse m) ++ f
                invKeySchedule = take (nr+1) $ invKeyExpansion nk (take nk (concatMap reverse (chop4 (take feed required))))
                obtained       = concat invKeySchedule
  
                expected = map (fromJust . unliteral) required
                result   = map (fromJust . unliteral) obtained
  
                sh i a b
                 | a == b = pad ++ show i ++ " " ++ disp a
                 | True   = pad ++ show i ++ " " ++ disp a ++ " |vs| " ++ disp b
                 where pad = if i < 10 then " " else ""
  
                disp = unwords . map (hex8 . literal)
  
                lexpected = length expected
                lresult   = length result
  
            when (lexpected /= lresult) $
               error $ what ++ ": BAD! Mismatching lengths: " ++ show (lexpected, lresult)
  
            let debugging = False
  
            if expected == result
               then if debugging
                       then putStrLn $ unlines $ ("Size " ++ what ++ ": Good") : zipWith3 sh [(0::Int)..] (chop4 expected) (chop4 result)
                       else putStrLn $ "GOOD: Key generation AES" ++ what
               else error    $ unlines $ ("Size " ++ what ++ ": BAD!") : zipWith3 sh [(0::Int)..] (chop4 expected) (chop4 result)
  
          roundTrip128 (i0, i1, i2, i3) (k0, k1, k2, k3)                 = roundTrip [i0, i1, i2, i3] [k0, k1, k2, k3]
          roundTrip192 (i0, i1, i2, i3) (k0, k1, k2, k3, k4, k5)         = roundTrip [i0, i1, i2, i3] [k0, k1, k2, k3, k4, k5]
          roundTrip256 (i0, i1, i2, i3) (k0, k1, k2, k3, k4, k5, k6, k7) = roundTrip [i0, i1, i2, i3] [k0, k1, k2, k3, k4, k5, k6, k7]
  
          roundTrip :: [SWord32] -> [SWord32] -> SBool
          roundTrip ptIn keyIn = pt .== pt' .&& pt .== pt''
             where pt  = map toSized ptIn
                   key = map toSized keyIn
  
                   (encKS, decKS) = aesKeySchedule key
                   ct   = aesEncrypt pt encKS
                   pt'  = aesDecrypt ct decKS
                   pt'' = aesDecryptUnwoundKey ct (aesInvKeySchedule (extractFinalKey key))
  
  -----------------------------------------------------------------------------
  -- * Verification
  -- ${verifIntro}
  -----------------------------------------------------------------------------
  {- $verifIntro
    While SMT based technologies can prove correct many small properties fairly quickly, it would
    be naive for them to automatically verify that our AES implementation is correct. (By correct,
    we mean decryption followed by encryption yielding the same result.) However, we can state
    this property precisely using SBV, and use quick-check to gain some confidence.
  -}
  
  -- | Correctness theorem for 128-bit AES. Ideally, we would run:
  --
  -- @
  --   prove aes128IsCorrect
  -- @
  --
  -- to get a proof automatically. Unfortunately, while SBV will successfully generate the proof
  -- obligation for this theorem and ship it to the SMT solver, it would be naive to expect the SMT-solver
  -- to finish that proof in any reasonable time with the currently available SMT solving technologies.
  -- Instead, we can issue:
  --
  -- @
  --   quickCheck aes128IsCorrect
  -- @
  --
  -- and get some degree of confidence in our code. Similar predicates can be easily constructed for 192, and
  -- 256 bit cases as well.
  aes128IsCorrect :: (SWord 32, SWord 32, SWord 32, SWord 32)  -- ^ plain-text words
                  -> (SWord 32, SWord 32, SWord 32, SWord 32)  -- ^ key-words
                  -> SBool                                     -- ^ True if round-trip gives us plain-text back
  aes128IsCorrect (i0, i1, i2, i3) (k0, k1, k2, k3) = pt .== pt'
     where pt  = [i0, i1, i2, i3]
           key = [k0, k1, k2, k3]
           (encKS, decKS) = aesKeySchedule key
           ct  = aesEncrypt pt encKS
           pt' = aesDecrypt ct decKS
  
  -----------------------------------------------------------------------------
  -- * Block encryption at full size
  -----------------------------------------------------------------------------
  
  -- | 128-bit encryption, that takes 128-bit values, instead of chunks. We have:
  --
  -- >>> hex8 $ aes128Enc 0x000102030405060708090a0b0c0d0e0f 0x00112233445566778899aabbccddeeff
  -- "69c4e0d86a7b0430d8cdb78070b4c55a"
  --
  -- You can also render this function as a stand-alone function using:
  --
  -- @
  --   sbv2smt (smtFunction "aes128Enc" aes128Enc)
  -- @
  aes128Enc :: SWord 128 -> SWord 128 -> SWord 128
  aes128Enc key pt = from32 $ aesEncrypt (to32 pt) ks
    where to32 :: SWord 128 -> [SWord 32]
          to32 x = [ bvExtract (Proxy @127) (Proxy @96) x
                   , bvExtract (Proxy  @95) (Proxy @64) x
                   , bvExtract (Proxy  @63) (Proxy @32) x
                   , bvExtract (Proxy  @31) (Proxy  @0) x
                   ]
  
          from32 :: [SWord 32] -> SWord 128
          from32 [a, b, c, d] = a # b # c # d
          from32 _ = error "nope"
  
          (ks, _)  = aesKeySchedule (to32 key)
  
  -----------------------------------------------------------------------------
  -- * Code generation
  -- ${codeGenIntro}
  -----------------------------------------------------------------------------
  {- $codeGenIntro
     We have emphasized that our T-Box implementation in Haskell was guided by clarity and correctness, not
     performance. Indeed, our implementation is hardly the fastest AES implementation in Haskell. However,
     we can use it to automatically generate straight-line C-code that can run fairly fast.
  
     For the purposes of illustration, we only show here how to generate code for a 128-bit AES block-encrypt
     function, that takes 8 32-bit words as an argument. The first 4 are the 128-bit input, and the final
     four are the 128-bit key. The impact of this is that the generated function would expand the key for
     each block of encryption, a needless task unless we change the key in every block. In a more serious application,
     we would instead generate code for both the 'aesKeySchedule' and the 'aesEncrypt' functions, thus reusing the
     key-schedule over many applications of the encryption call. (Unfortunately doing this is rather cumbersome right
     now, since Haskell does not support fixed-size lists.)
  -}
  
  -- | Code generation for 128-bit AES encryption.
  --
  -- The following sample from the generated code-lines show how T-Boxes are rendered as C arrays:
  --
  -- @
  --   static const SWord32 table1[] = {
  --       0xc66363a5UL, 0xf87c7c84UL, 0xee777799UL, 0xf67b7b8dUL,
  --       0xfff2f20dUL, 0xd66b6bbdUL, 0xde6f6fb1UL, 0x91c5c554UL,
  --       0x60303050UL, 0x02010103UL, 0xce6767a9UL, 0x562b2b7dUL,
  --       0xe7fefe19UL, 0xb5d7d762UL, 0x4dababe6UL, 0xec76769aUL,
  --       ...
  --       }
  -- @
  --
  -- The generated program has 5 tables (one sbox table, and 4-Tboxes), all converted to fast C arrays. Here
  -- is a sample of the generated straightline C-code:
  --
  -- @
  --   const SWord8  s1915 = (SWord8) s1912;
  --   const SWord8  s1916 = table0[s1915];
  --   const SWord16 s1917 = (((SWord16) s1914) << 8) | ((SWord16) s1916);
  --   const SWord32 s1918 = (((SWord32) s1911) << 16) | ((SWord32) s1917);
  --   const SWord32 s1919 = s1844 ^ s1918;
  --   const SWord32 s1920 = s1903 ^ s1919;
  -- @
  --
  -- The GNU C-compiler does a fine job of optimizing this straightline code to generate a fairly efficient C implementation.
  cgAES128BlockEncrypt :: IO ()
  cgAES128BlockEncrypt = compileToC Nothing "aes128BlockEncrypt" $ do
          pt  <- cgInputArr 4 "pt"        -- plain-text as an array of 4 Word32's
          key <- cgInputArr 4 "key"       -- key as an array of 4 Word32s
  
          -- Use the test values from Appendix C.1 of the AES standard as the driver values
          cgSetDriverValues $ map (fromIntegral . fromJust . unliteral) $ commonPT ++ aes128Key
  
          let (encKs, _) = aesKeySchedule key
          cgOutputArr "ct" $ aesEncrypt pt encKs
  
  -----------------------------------------------------------------------------
  -- * C-library generation
  -- ${libraryIntro}
  -----------------------------------------------------------------------------
  {- $libraryIntro
     The 'cgAES128BlockEncrypt' example shows how to generate code for 128-bit AES encryption. As the generated
     function performs encryption on a given block, it performs key expansion as necessary. However, this is
     not quite practical: We would like to expand the key only once, and encrypt the stream of plain-text blocks using
     the same expanded key (potentially using some crypto-mode), until we decide to change the key. In this
     section, we show how to use SBV to instead generate a library of functions that can be used in such a scenario.
     The generated library is a typical @.a@ archive, that can be linked using the C-compiler as usual.
  -}
  
  -- | Components of the AES implementation that the library is generated from. For each case, we provide
  -- the driver values from the AES test-vectors.
  aesLibComponents :: Int -> [(String, [Integer], SBVCodeGen ())]
  aesLibComponents sz = [ ("aes" ++ show sz ++ "KeySchedule",    keyDriverVals,    keySchedule)
                        , ("aes" ++ show sz ++ "BlockEncrypt",   encDriverVals,    enc)
                        , ("aes" ++ show sz ++ "BlockDecrypt",   decDriverVals,    dec)
                        , ("aes" ++ show sz ++ "InvKeySchedule", invKeyDriverVals, invKeySchedule)
                        , ("aes" ++ show sz ++ "OTFDecrypt",     invDecDriverVals, otfDec)
                        ]
    where badSize = error $ "aesLibComponents: Size must be one of 128, 192, or 256; received: " ++ show sz
  
          -- key-schedule
          nk
           | sz == 128 = 4
           | sz == 192 = 6
           | sz == 256 = 8
           | True      = badSize
  
          -- We get 4*(nr+1) keys, where nr = nk + 6
          nr = nk + 6
          xk = 4 * (nr + 1)
  
          (keyDriverVals, invKeyDriverVals, encDriverVals, decDriverVals, invDecDriverVals)
             | sz == 128 = (keyDriver aes128Key, keyDriver aes128InvKey, encDriver commonPT aes128Key, decDriver aes128CT aes128Key, invDecDriver aes128CT aes128InvKey)
             | sz == 192 = (keyDriver aes192Key, keyDriver aes192InvKey, encDriver commonPT aes192Key, decDriver aes192CT aes192Key, invDecDriver aes192CT aes192InvKey)
             | sz == 256 = (keyDriver aes256Key, keyDriver aes256InvKey, encDriver commonPT aes256Key, decDriver aes256CT aes256Key, invDecDriver aes256CT aes256InvKey)
             | True      = badSize
             where keyDriver       key = map cvt $ concatMap reverse (chop4 key)
                   encDriver    pt key = map cvt $ pt ++ flatten (fst (aesKeySchedule    key))
                   decDriver    ct key = map cvt $ ct ++ flatten (snd (aesKeySchedule    key))
                   invDecDriver ct key = map cvt $ ct ++ flatten      (aesInvKeySchedule key)
  
                   flatten (f, mid, l) = f ++ concat mid ++ l
                   cvt = fromIntegral . fromJust . unliteral
  
          keySchedule = do key <- cgInputArr nk "key"     -- key
                           let (encKS, decKS) = aesKeySchedule key
                           cgOutputArr "encKS" (ksToXKey encKS)
                           cgOutputArr "decKS" (ksToXKey decKS)
  
          invKeySchedule = do key <- cgInputArr nk "key"     -- key
                              let decKS = aesInvKeySchedule (concatMap reverse (chop4 key))
                              cgOutputArr "decKS" (ksToXKey decKS)
  
          -- encryption
          enc = do pt   <- cgInputArr 4  "pt"    -- plain-text
                   xkey <- cgInputArr xk "xkey"  -- expanded key
                   cgOutputArr "ct" $ aesEncrypt pt (xkeyToKS xkey)
  
          -- decryption
          dec = do pt   <- cgInputArr 4  "ct"    -- cipher-text
                   xkey <- cgInputArr xk "xkey"  -- expanded key
                   cgOutputArr "pt" $ aesDecrypt pt (xkeyToKS xkey)
  
          -- on-the-fly decryption
          otfDec = do ct   <- cgInputArr 4  "ct"    -- cipher-text
                      xkey <- cgInputArr xk "xkey"  -- expanded key
                      cgOutputArr "pt" $ aesDecryptUnwoundKey ct (xkeyToKS xkey)
  
          -- Transforming back and forth from our KS type to a flat array used by the generated C code
          -- Turn a series of expanded keys to our internal KS type
          xkeyToKS :: [SWord 32] -> KS
          xkeyToKS xs = (f, m, l)
             where f  = take 4 xs                             -- first round key
                   m  = chop4 (take (xk - 8) (drop 4 xs))     -- middle rounds
                   l  = drop (xk - 4) xs                      -- last round key
  
          -- Turn a KS to a series of expanded key words
          ksToXKey :: KS -> [SWord 32]
          ksToXKey (f, m, l) = f ++ concat m ++ l
  
  -- | Generate code for AES functionality; given the key size.
  cgAESLibrary :: Int -> Maybe FilePath -> IO ()
  cgAESLibrary sz mbd
    | sz `elem` [128, 192, 256] = void $ compileToCLib mbd nm [(fnm, configure dvals f) | (fnm, dvals, f) <- aesLibComponents sz]
    | True                      = error $ "cgAESLibrary: Size must be one of 128, 192, or 256, received: " ++ show sz
    where nm = "aes" ++ show sz ++ "Lib"
  
          configure dvals code = cgSetDriverValues dvals >> code
  
  -- | Generate a C library, containing functions for performing 128-bit enc/dec/key-expansion.
  -- A note on performance: In a very rough speed test, the generated code was able to do
  -- 6.3 million block encryptions per second on a decent MacBook Pro. On the same machine, OpenSSL
  -- reports 8.2 million block encryptions per second. So, the generated code is about 25% slower
  -- as compared to the highly optimized OpenSSL implementation. (Note that the speed test was done
  -- somewhat simplistically, so these numbers should be considered very rough estimates.)
  cgAES128Library :: IO ()
  cgAES128Library = cgAESLibrary 128 Nothing
  
  --------------------------------------------------------------------------------------------
  -- | For doctest purposes only
  hex8 :: (SymVal a, Show a, Integral a) => SBV a -> String
  hex8 v = replicate (8 - length s) '0' ++ s
    where s = flip showHex "" . fromJust . unliteral $ v
  
  -- | Chunk in groups of 4. (This function must be in some standard library, where?)
  chop4 :: [a] -> [[a]]
  chop4 [] = []
  chop4 xs = let (f, r) = splitAt 4 xs in f : chop4 r
  
  {- HLint ignore aesRound             "Use head" -}
  {- HLint ignore aesInvRound          "Use head" -}
  {- HLint ignore aesDecryptUnwoundKey "Use head" -}
  