/**

Implementation of Ascon-Based Lightweight Cryptography

Reference:

Meltem Sönmez Turan, Kerry A. McKay, Donghoon Chang, Jinkeon Kang,
John Kelsey (2025) Ascon-Based Lightweight Cryptography Standards
for Constrained Devices. (National Institute of Standards and Technology,
Gaithersburg, MD), NIST Special Publication (SP) NIST SP 800-232.
https://doi.org/10.6028/NIST.SP.800-232

*/
module Ascon where

// 2.1. Auxiliary Functions

/// Parse function.
///
/// The parse(𝑋, 𝑟) function parses the input bitstring 𝑋 into a sequence
/// of blocks 𝑋₀, 𝑋₁, …, 𝑋̃ℓ, where 𝓁 ← ⌊|𝑋|/𝑟⌋ (i.e., 𝑋 ← 𝑋₀ ∥ 𝑋₁ ∥ … ∥ 𝑋̃ℓ).
/// The 𝑋ᵢ blocks for 0 ≤ i ≤ 𝓁 − 1 each have a bit length 𝑟, whereas
/// 0 ≤ |𝑋̃ℓ| ≤ 𝑟 − 1 (see Algorithm 1). When |𝑋| mod 𝑟 = 0, the final
/// block is empty (i.e., |𝑋̃ℓ| = 0).
parse : {r, n} (fin n, fin r, r >= 1) => [n] -> ([n/r][r], [n%r])
parse (M_ # Ml) = (split M_, Ml)

/// Padding rule.
///
/// The function pad(𝑋, 𝑟) appends the bit 1 to the bitstring 𝑋, followed
/// bythe bitstring 0ʲ, where 𝑗 is equal to (−|𝑋|−1) mod 𝑟. The length of
/// the output bitstring is a multiple of 𝑟 (see Algorithm 2). For examples
/// of padding when representing the data as 64-bit unsigned integers, see
/// Appendix A.2.
pad : {r, n} (n < r, fin r) => [n] -> [r]
pad M = M # 0b1 # 0

/// Combination of parse and pad that splits a bitstring into a sequence
/// of integers using Cryptol's native big-endian representation.
toBlocks : {r, n} (r >= 1, fin r, fin n) => [n] -> [n / r + 1][r]
toBlocks M = map reverse (M1 # [pad M2])
  where
    (M1, M2) = parse M

// 3. Ascon Permutations

type constraint ValidRnd rnd = rnd <= 16

/// Core permutation function parameterized by a round count.
Ascon_p : {rnd} (ValidRnd rnd) => State -> State
Ascon_p S = foldl p`{rnd} S (drop`{back=rnd} Const)

/// Single round of the Ascon-p permutation.
p : {rnd} (ValidRnd rnd) => State -> [64] -> State
p S ci = pL (pS (pC S ci))

// 3.1. Internal State

/// The permutations operate on the 320-bit state 𝑆, which is represented
/// as five 64-bit words denoted as 𝑆ᵢ for 0 ≤ 𝑖 ≤ 4:
///
///     𝑆 = 𝑆₀ ‖ 𝑆₁ ‖ 𝑆₂ ‖ 𝑆₃ ‖ 𝑆₄
///
/// Let 𝑠ᵢⱼ represent the 𝑗th bit of 𝑆ᵢ, 0 ≤ 𝑗 < 64. In this specification
/// of the Ascon permutation, each state word represents a 64-bit unsigned
/// integer, where the least significant bit is the rightmost bit. Details
/// on other representations of the state can be found in Appendix A.
type State = [5][64]

// 3.2. Constant-Addition Layer pC

/// The constant cᵢ of round 𝑖 of the Ascon permutation Ascon-p[𝑟𝑛𝑑]
/// (instantiated with 𝑟𝑛𝑑 rounds) for 𝑟𝑛𝑑 ≤ 16 and 0 ≤ 𝑖 < 𝑟𝑛𝑑 − 1
/// is defined as
///
///     cᵢ = const[16 − 𝑟𝑛𝑑 + 𝑖]
///
/// where const[0],…,const[15] are defined in Table 5. The constant-addition
/// layer 𝑃𝑐 adds a 64-bit round constant cᵢ to 𝑆₂ in round 𝑖, for i ≥ 0,
///
///     𝑆₂ = 𝑆₂ ⊕ cᵢ.
pC : State -> [64] -> State
pC [S0, S1, S2, S3, S4] ci = [S0, S1, S2 ^ ci, S3, S4]

/// Table 5. The constants constᵢ to derive round constants of the Ascon
/// permutations
Const : [16][64]
Const =
    [ 0x000000000000003c
    , 0x000000000000002d
    , 0x000000000000001e
    , 0x000000000000000f
    , 0x00000000000000f0
    , 0x00000000000000e1
    , 0x00000000000000d2
    , 0x00000000000000c3
    , 0x00000000000000b4
    , 0x00000000000000a5
    , 0x0000000000000096
    , 0x0000000000000087
    , 0x0000000000000078
    , 0x0000000000000069
    , 0x000000000000005a
    , 0x000000000000004b
    ]

// 3.3. Substitution Layer pS

pS : State -> State
pS S = transpose (map SBox (transpose S))

SBox : [5] -> [5]
SBox i = SBoxTable@i

/// Table 6.
SBoxTable : [32][5]
SBoxTable =
    map drop
    [ 0x04, 0x0b, 0x1f, 0x14, 0x1a, 0x15, 0x09, 0x02
    , 0x1b, 0x05, 0x08, 0x12, 0x1d, 0x03, 0x06, 0x1c
    , 0x1e, 0x13, 0x07, 0x0e, 0x00, 0x0d, 0x11, 0x18
    , 0x10, 0x0c, 0x01, 0x19, 0x16, 0x0a, 0x0f, 0x17
    ]

property SBoxCorrect x0 x1 x2 x3 x4 = [y0, y1, y2, y3, y4] == SBox [x0, x1, x2, x3, x4]
           where
             y0 = x4&&x1 ^ x3 ^ x2&&x1 ^ x2 ^ x1&&x0 ^ x1 ^ x0
             y1 = x4 ^ x3&&x2 ^ x3&&x1 ^ x3 ^ x2&&x1 ^ x2 ^ x1 ^ x0
             y2 = x4&&x3 ^ x4 ^ x2 ^ x1 ^ 1
             y3 = x4&&x0 ^ x4 ^ x3&&x0 ^ x3 ^ x2 ^ x1 ^ x0
             y4 = x4&&x1 ^ x4 ^ x3 ^ x1&&x0 ^ x1

// 3.4. Linear Diffusion Layer pL

pL : State -> State
pL [S0, S1, S2, S3, S4] =
    [ sigma S0 19 28
    , sigma S1 61 39
    , sigma S2  1  6
    , sigma S3 10 17
    , sigma S4  7 41
    ]
    where
        sigma : [64] -> [6] -> [6] -> [64]
        sigma x i j = x ^ (x >>> i) ^ (x >>> j)

wordsToBits : {w,n} (fin w) => [n][w] -> [w*n]
wordsToBits M = join (map reverse M)

bitsToWords : {w,n} (fin w) => [w*n] -> [n][w]
bitsToWords M = map reverse (split M)