diff --git a/Ascon.cry b/Ascon.cry
index fbdd44e..df98a57 100644
--- a/Ascon.cry
+++ b/Ascon.cry
@@ -15,75 +15,86 @@ 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 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, m} (fin m, fin r, r >= 1) => [m] -> ([m / r][r], [m % 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]
+/** 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, m} (m < r, fin r) => [m] -> [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]
+/**
+ * Combination of parse and pad that splits a bitstring into a sequence
+ * of integers using Cryptol's native big-endian representation.
+ */
+toBlocks : {r, m} (r >= 1, fin r, fin m) => [m] -> [m / r + 1][r]
 toBlocks M = map reverse (M1 # [pad M2])
   where
     (M1, M2) = parse M
 
 // 3. Ascon Permutations
 
+/** Predicate on valid round counts. */
 type constraint ValidRnd rnd = rnd <= 16
 
-/// Core permutation function parameterized by a round count.
+/** Core permutation function parameterized by a round count. */
 Ascon_p : {rnd} (ValidRnd rnd) => State -> State
 Ascon_p S = foldl round S (drop`{back=rnd} Const)
 
-/// Single round of the Ascon-p permutation.
+/** Single round of the Ascon-p permutation. */
 round : State -> [64] -> State
 round 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.
+/**
+ * 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ᵢ.
+// 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
+/**
+ * Table 5. The constants constᵢ to derive round constants of the Ascon
+ * permutations
+ */
 Const : [16][64]
 Const =
     [ 0x000000000000003c
@@ -112,7 +123,7 @@ pS S = transpose (map SBox (transpose S))
 SBox : [5] -> [5]
 SBox i = SBoxTable@i
 
-/// Table 6.
+/** Table 6. */
 SBoxTable : [32][5]
 SBoxTable =
     map drop
@@ -142,32 +153,32 @@ pL [S0, S1, S2, S3, S4] =
     ]
     where
         sigma : [64] -> [6] -> [6] -> [64]
-        sigma x i j = x ^ (x >>> i) ^ (x >>> j)
+        sigma x i j = x ^ x>>>i ^ x>>>j
 
-wordsToBits : {w,n} (fin w) => [n][w] -> [w*n]
+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 : {w, n} (fin w) => [w*n] -> [n][w]
 bitsToWords M = map reverse (split M)
 
 // 4. Authenticated Encryption Schema: Ascon-AEAD128
 
-/// Encryption using Ascon-AEAD128
-///
-/// Parameters:
-/// - Key
-/// - Nonce
-/// - Additional data
-/// - Plaintext
-///
-/// Returns:
-/// - Authenticated ciphertext
+/** Encryption using Ascon-AEAD128
+ *
+ * Parameters:
+ * - K: Key
+ * - N: Nonce
+ * - A: Additional data
+ * - P: Plaintext
+ *
+ * Returns:
+ * - Authenticated ciphertext
+ */
 AEAD128_encrypt : {a, p} (fin a, fin p) => [128] -> [128] -> [a] -> [p] -> [p + 128]
 AEAD128_encrypt K N A P = C # T
   where
-    // key and nonce as two 64-bit integers
-    [K0,K1] = bitsToWords K
-    [N0,N1] = bitsToWords N
+    [K0, K1] = bitsToWords K
+    [N0, N1] = bitsToWords N
 
     S0 = Ascon_p`{12} [Ascon_AEAD128_IV, K0, K1, N0, N1] ^ [0, 0, 0, K0, K1]
     SA = AddAD S0 A
@@ -179,24 +190,29 @@ AEAD128_encrypt K N A P = C # T
     ST = Ascon_p`{12} (last SCs ^ [0, 0, K0, K1, 0])
     T = ExtractT ST ^ K
 
-/// Decryption using Ascon-AEAD128
-///
-/// Parameters:
-/// - Key
-/// - Nonce
-/// - Additional data
-/// - Ciphertext
-///
-/// Returns:
-/// - Some plaintext on success
-/// - None when message authentication fails
+/** Ascon-AEAD128 decryption algorithm on bitstreams.
+ *
+ * Type parameters:
+ * - a: Bit-length of additional data
+ * - p: Bit-length of plaintext
+ *
+ * Parameters:
+ * - K: Key
+ * - N: Nonce
+ * - A: Additional data
+ * - C: Ciphertext
+ *
+ * Returns:
+ * - Some plaintext on authentication success
+ * - None on authentication failure
+ */
 AEAD128_decrypt : {a, p} (fin a, fin p) => [128] -> [128] -> [a] -> [p + 128] -> Option [p]
 AEAD128_decrypt K N A (Cs_ # Cl # T) =
-  if T == T' then Some P else trace "P" P None
+  if T == T' then Some P else None
   where
     // key and nonce as two 64-bit integers
-    [K0,K1] = bitsToWords K
-    [N0,N1] = bitsToWords N
+    [K0, K1] = bitsToWords K
+    [N0, N1] = bitsToWords N
 
     S0 = Ascon_p`{12} [Ascon_AEAD128_IV, K0, K1, N0, N1] ^ [0, 0, 0, K0, K1]
     SA = AddAD S0 A
@@ -204,21 +220,52 @@ AEAD128_decrypt K N A (Cs_ # Cl # T) =
     Cs = split`{p/128, 128} Cs_
 
     SCs # [SCl] = [SA] # [Ascon_p`{8} (AssignC s c) | s <- SCs | c <- Cs]
+    Masks = map ExtractC SCs
 
-    Plmask # SCl' = ExtractC SCl
+    Maskl # SCl' = ExtractC SCl
     Sl' = AssignC SCl (Cl # (0b1#0 ^ SCl'))
 
     Ps = zipWith (\x y -> ExtractC x ^ y) SCs Cs
-    Pl = Plmask ^ Cl
-    P = join Ps # Pl
+    P = join (Ps ^ Masks) # (Maskl ^ Cl)
 
     ST = Ascon_p`{12} (Sl' ^ [0, 0, K0, K1, 0])
     T' = ExtractT ST ^ K
 
+/** Ascon-AEAD128 encryption function on bytes.
+ *
+ * Type parameters:
+ * - a: Byte-length of additional data
+ * - p: Byte-length of plaintext
+ *
+ * Parameters:
+ * - K: Key
+ * - N: Nonce
+ * - A: Additional data
+ * - P: Plaintext
+ *
+ * Returns:
+ * - Authenticated ciphertext
+ */
 AEAD128_encrypt_bytes : {a, p} (fin a, fin p) => [16][8] -> [16][8] -> [a][8] -> [p][8] -> [p + 16][8]
 AEAD128_encrypt_bytes K N A P =
   bitsToWords (AEAD128_encrypt (wordsToBits K) (wordsToBits N) (wordsToBits A) (wordsToBits P))
 
+/** Ascon-AEAD128 decryption algorithm on bytes.
+ *
+ * Type parameters:
+ * - a: Bit-length of additional data
+ * - p: Bit-length of plaintext
+ *
+ * Parameters:
+ * - K: Key
+ * - N: Nonce
+ * - A: Additional data
+ * - C: Ciphertext
+ *
+ * Returns:
+ * - Some plaintext on authentication success
+ * - None on authentication failure
+ */
 AEAD128_decrypt_bytes : {a, p} (fin a, fin p) => [16][8] -> [16][8] -> [a][8] -> [p + 16][8] -> Option ([p][8])
 AEAD128_decrypt_bytes K N A C =
   case AEAD128_decrypt (wordsToBits K) (wordsToBits N) (wordsToBits A) (wordsToBits C) of
@@ -253,6 +300,12 @@ Ascon_AEAD128_IV = 0x00001000808c0001
 
 // 5. Hash and eXtendable-Output Functions (XOFs)
 
+/** Hash function construction
+ *
+ * Parameters:
+ * - Initialization vector
+ * - List of blocks as integers
+ */
 hashBlocks : {n} (fin n) => [64] -> [n][64] -> [inf]
 hashBlocks IV Ms = wordsToBits [head S | S <- iterate Ascon_p`{12} Sn]
     where
@@ -262,38 +315,90 @@ hashBlocks IV Ms = wordsToBits [head S | S <- iterate Ascon_p`{12} Sn]
 AbsorbBlock64 : State -> [64] -> State
 AbsorbBlock64 [s0, s1, s2, s3, s4] X = Ascon_p`{12} [X ^ s0, s1, s2, s3, s4]
 
-/// 5.1. Specification of Ascon-Hash256
-Hash256 : {n} (fin n) => [n] -> [256]
+// 5.1. Specification of Ascon-Hash256
+
+/** Ascon-Hash256 implementation on bitstreams.
+ *
+ * Type parameters:
+ * - m: Bit-length of message
+ *
+ * Parameters:
+ * - M: Message
+ */
+Hash256 : {m} (fin m) => [m] -> [256]
 Hash256 M = take (hashBlocks Hash256_IV (toBlocks M))
 
-Hash256_bytes : {n} (fin n) => [n][8] -> [32][8]
+/** Ascon-Hash256 implementation on bytes.
+ *
+ * Type parameters:
+ * - m: Byte-length of message
+ *
+ * Parameters:
+ * - M: Message
+ */
+Hash256_bytes : {m} (fin m) => [m][8] -> [32][8]
 Hash256_bytes M = bitsToWords (Hash256 (wordsToBits M))
 
+/** Ascon-Hash256 initialization vector */
 Hash256_IV : [64]
 Hash256_IV = 0x0000080100cc0002
 
 // 5.2. Specification of Ascon-XOF128
 
-XOF128 : {r, n} (fin n, fin r) => [n] -> [r]
+/** Ascon-XOF256 implementation on bitstreams.
+ *
+ * Type parameters:
+ * - r: Bit-length of digest
+ * - m: Bit-length of message
+ *
+ * Parameters:
+ * - M: Message
+ */
+XOF128 : {r, m} (fin m, fin r) => [m] -> [r]
 XOF128 M = take (hashBlocks XOF128_IV (toBlocks M))
 
+/** Ascon-XOF256 implementation on bytes. */
 XOF128_bytes : {r, n} (fin n, fin r) => [n][8] -> [r][8]
 XOF128_bytes M = bitsToWords (XOF128 (wordsToBits M))
 
+/** Ascon-XOF128 initialization vector */
 XOF128_IV : [64]
 XOF128_IV = 0x0000080000cc0003
 
 // 5.3. Specification of Ascon-CXOF128
 
-CXOF128 : {r, c, n} (fin n, fin r, fin c, 64 >= width c) => [c] -> [n] -> [r]
+/** Ascon-CXOF256 implementation on bitstreams.
+ *
+ * Type parameters:
+ * - r: Bit-length of digest
+ * - c: Bit-length of customization string
+ * - m: Bit-length of message
+ *
+ * Parameters:
+ * - Z: User-defined customization string
+ * - M: Message
+ */
+CXOF128 : {r, c, m} (fin m, fin r, fin c, 64 >= width c) => [c] -> [m] -> [r]
 CXOF128 Z M = take (hashBlocks CXOF128_IV Ms)
   where
     Ms = [`c]
        # toBlocks Z
        # toBlocks M
 
-CXOF128_bytes : {r, z, n} (fin n, fin r, 61 >= width z) => [z][8] -> [n][8] -> [r][8]
+/** Ascon-CXOF256 implementation on bytes.
+ *
+ * Type parameters:
+ * - r: Byte-length of digest
+ * - c: Byte-length of customization string
+ * - m: Byte-length of message
+ *
+ * Parameters:
+ * - Z: User-defined customization string
+ * - M: Message
+ */
+CXOF128_bytes : {r, z, m} (fin m, fin r, 61 >= width z) => [z][8] -> [m][8] -> [r][8]
 CXOF128_bytes Z M = bitsToWords (CXOF128 (wordsToBits Z) (wordsToBits M))
 
+/** Ascon-CXOF128 initialization vector */
 CXOF128_IV : [64]
 CXOF128_IV = 0x0000080000cc0004