scalar-crypto: aesks1i clarifications

Trying to make the behavior of aesks1i easier to understand, and more
obviously correspond to the specification.

- Ensure that the aes_decode_rcon function only accepts valid values, plus
  add comments.

- Re-name operands to aesks1i rcon -> rnum to be in line with the specification

- Re-structure the Sail code for clarity based on jrtc27's suggestion.

2 files changed, 20 insertions, 16 deletions

diff --git a/model/riscv_insts_zkn.sail b/model/riscv_insts_zkn.sail
index 5cda5a3..89f81fa 100644
--- a/model/riscv_insts_zkn.sail
+++ b/model/riscv_insts_zkn.sail
@@ -237,8 +237,8 @@ union clause ast = AES64ES   : (regidx, regidx, regidx)
 union clause ast = AES64DSM  : (regidx, regidx, regidx)
 union clause ast = AES64DS   : (regidx, regidx, regidx)
 
-mapping clause encdec = AES64KS1I (rcon, rs1, rd) if (haveZkne() | haveZknd()) & (sizeof(xlen) == 64) & (rcon <_u 0xB)
-  <-> 0b00 @ 0b11000 @ 0b1 @ rcon @ rs1 @ 0b001 @ rd @ 0b0010011
+mapping clause encdec = AES64KS1I (rnum, rs1, rd) if (haveZkne() | haveZknd()) & (sizeof(xlen) == 64) & (rnum <_u 0xB)
+  <-> 0b00 @ 0b11000 @ 0b1 @ rnum @ rs1 @ 0b001 @ rd @ 0b0010011
 
 mapping clause encdec = AES64IM (rs1, rd) if haveZknd() & sizeof(xlen) == 64
   <-> 0b00 @ 0b11000 @ 0b00000 @ rs1 @ 0b001 @ rd @ 0b0010011
@@ -279,13 +279,16 @@ mapping clause assembly = AES64DSM (rs2, rs1, rd)
 mapping clause assembly = AES64DS (rs2, rs1, rd)
   <-> "aes64ds" ^ spc() ^ reg_name(rd) ^ sep() ^ reg_name(rs1) ^ sep() ^ reg_name(rs2)
 
-function clause execute (AES64KS1I(rcon, rs1, rd)) = {
+/* Note: The decoding for this instruction ensures that `rnum` is always in
+   the range 0x0..0xA. See the encdec clause for AES64KS1I.
+   The rum == 0xA case is used specifically for the AES-256 KeySchedule */
+function clause execute (AES64KS1I(rnum, rs1, rd)) = {
   assert(sizeof(xlen) == 64);
-  let rs1_hi  : bits(32) = X(rs1)[63..32];
-  let rc      : bits(32) = aes_decode_rcon(rcon);
-  let rotated : bits(32) = if (rcon == 0xA) then rs1_hi else (rs1_hi >>> 8);
-  let post_sb : bits(32) = aes_subword_fwd(rotated);
-  X(rd) = (post_sb ^ rc) @ (post_sb ^ rc);
+  let prev     : bits(32) = X(rs1)[63..32];
+  let subwords : bits(32) = aes_subword_fwd(prev);
+  let result   : bits(32) = if (rnum == 0xA) then subwords
+                            else (subwords >>> 8) ^ aes_decode_rcon(rnum);
+  X(rd) = result @ result;
   RETIRE_SUCCESS
 }
 
diff --git a/model/riscv_types_kext.sail b/model/riscv_types_kext.sail
index 9c0db6d..2c49706 100644
--- a/model/riscv_types_kext.sail
+++ b/model/riscv_types_kext.sail
@@ -68,8 +68,15 @@ function aes_mixcolumn_inv(x) = {
   b3 @ b2 @ b1 @ b0 /* Return value */
 }
 
-val aes_decode_rcon : bits(4) -> bits(32)
+/* Turn a round number into a round constant for AES. Note that the
+   AES64KS1I instruction is defined such that the r argument is always
+   in the range 0x0..0xA. Values of rnum outside the range 0x0..0xA
+   do not decode to the AES64KS1I instruction. The 0xA case is used
+   specifically for the AES-256 KeySchedule, and this function is never
+   called in that case. */
+val aes_decode_rcon : bits(4) -> bits(32) effect {escape}
 function aes_decode_rcon(r) = {
+  assert(r <_u 0xA);
   match r {
     0x0 => 0x00000001,
     0x1 => 0x00000002,
@@ -80,13 +87,7 @@ function aes_decode_rcon(r) = {
     0x6 => 0x00000040,
     0x7 => 0x00000080,
     0x8 => 0x0000001b,
-    0x9 => 0x00000036,
-    0xA => 0x00000000,
-    0xB => 0x00000000,
-    0xC => 0x00000000,
-    0xD => 0x00000000,
-    0xE => 0x00000000,
-    0xF => 0x00000000
+    0x9 => 0x00000036
   }
 }