AES divides the input into chunks of 128 bits. 128 bits is 16 bytes. 16 bytes can be arranged in a 4×4 matrix and are represented in this two dimensional form by AES. This matrix is called the “state”.

Using cryptol, the AES algorithm is described with the following functions

AESRound : (RoundKey, State) -> State

AESRound (rk, s) = AddRoundKey (rk, MixColumns (ShiftRows (SubBytes s)))

The final round leaves out the MixColumns transform.

AESFinalRound : (RoundKey, State) -> State
AESFinalRound (rk, s) = AddRoundKey (rk, ShiftRows (SubBytes s))

The inverse transform is as follows

AESInvRound : (RoundKey, State) -> State

AESInvRound (rk, s) =
InvMixColumns (AddRoundKey (rk, InvSubBytes (InvShiftRows s)))

AESFinalInvRound : (RoundKey, State) -> State

AESFinalInvRound (rk, s) = AddRoundKey (rk, InvSubBytes (InvShiftRows s))