Skip to main content

ab_riscv_primitives/instructions/rv32/zk/zkn/
zkne.rs

1//! RV32 Zkne extension
2
3#[cfg(test)]
4mod tests;
5
6use crate::instructions::Instruction;
7use crate::instructions::rv32::zk::zkn::zknd::Rv32AesBs;
8use crate::registers::general_purpose::Register;
9use ab_riscv_macros::instruction;
10use core::fmt;
11
12/// RISC-V RV32 Zkne instructions (AES encryption)
13#[instruction]
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum Rv32ZkneInstruction<Reg> {
16    /// AES final round encryption step: SubBytes on one byte of rs2, rotated to the byte lane
17    /// selected by bs, XOR'd into rs1.
18    ///
19    /// `rd = rs1 ^ rol32(SBOX[(rs2 >> (bs*8)) & 0xff] as u32, bs*8)`
20    Aes32Esi {
21        rd: Reg,
22        rs1: Reg,
23        rs2: Reg,
24        bs: Rv32AesBs,
25    },
26    /// AES middle round encryption step: SubBytes + partial MixColumns on one byte of rs2, rotated
27    /// to the byte lane selected by bs, XOR'd into rs1.
28    ///
29    /// `rd = rs1 ^ rol32(MixColByte(SBOX[(rs2 >> (bs*8)) & 0xff]), bs*8)`
30    Aes32Esmi {
31        rd: Reg,
32        rs1: Reg,
33        rs2: Reg,
34        bs: Rv32AesBs,
35    },
36}
37
38/// Encoding layout (R-type, opcode 0x33, funct3 0x0):
39///
40/// ```text
41/// [31:30] bs       - 2-bit byte select
42/// [29:25] funct5   - 0b10001 (aes32esi) / 0b10011 (aes32esmi)
43/// [24:20] rs2
44/// [19:15] rs1
45/// [14:12] funct3   - 0b000
46/// [11:7]  rd
47/// [6:0]   opcode   - 0b0110011 (OP)
48/// ```
49///
50/// Ratified match/mask values (from riscv-opcodes):
51///   MATCH_AES32ESI  = 0x22000033, MASK_AES32ESI  = 0x3e00707f
52///   MATCH_AES32ESMI = 0x26000033, MASK_AES32ESMI = 0x3e00707f
53///
54/// `rd` and `rs1` are independent fields. The assembler convention places
55/// the accumulator in both rd and rs1 (the `rt` pattern), but the hardware
56/// does not require rd == rs1 and the decoder must not enforce it.
57#[instruction]
58impl<Reg> const Instruction for Rv32ZkneInstruction<Reg>
59where
60    Reg: [const] Register<Type = u32>,
61{
62    type Reg = Reg;
63
64    #[inline(always)]
65    fn try_decode(instruction: u32) -> Option<Self> {
66        let opcode = (instruction & 0b111_1111) as u8;
67        let rd_bits = ((instruction >> 7) & 0x1f) as u8;
68        let funct3 = ((instruction >> 12) & 0b111) as u8;
69        let rs1_bits = ((instruction >> 15) & 0x1f) as u8;
70        let rs2_bits = ((instruction >> 20) & 0x1f) as u8;
71        let funct5 = ((instruction >> 25) & 0b1_1111) as u8;
72        let bs_bits = ((instruction >> 30) & 0b11) as u8;
73
74        // R-type OP opcode only
75        if opcode != 0b0110011 {
76            None?;
77        }
78        if funct3 != 0b000 {
79            None?;
80        }
81
82        let rd = Reg::from_bits(rd_bits)?;
83        let rs1 = Reg::from_bits(rs1_bits)?;
84        let rs2 = Reg::from_bits(rs2_bits)?;
85        let bs = Rv32AesBs::from_bits(bs_bits)?;
86
87        match funct5 {
88            // aes32esi:  bs[31:30] | 0b10001[29:25]
89            0b10001 => Some(Self::Aes32Esi { rd, rs1, rs2, bs }),
90            // aes32esmi: bs[31:30] | 0b10011[29:25]
91            0b10011 => Some(Self::Aes32Esmi { rd, rs1, rs2, bs }),
92            _ => None,
93        }
94    }
95
96    #[inline(always)]
97    fn alignment() -> u8 {
98        align_of::<u32>() as u8
99    }
100
101    #[inline(always)]
102    fn size(&self) -> u8 {
103        size_of::<u32>() as u8
104    }
105}
106
107impl<Reg> fmt::Display for Rv32ZkneInstruction<Reg>
108where
109    Reg: fmt::Display,
110{
111    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
112        match self {
113            Self::Aes32Esi { rd, rs1, rs2, bs } => {
114                write!(f, "aes32esi {rd}, {rs1}, {rs2}, {bs}")
115            }
116            Self::Aes32Esmi { rd, rs1, rs2, bs } => {
117                write!(f, "aes32esmi {rd}, {rs1}, {rs2}, {bs}")
118            }
119        }
120    }
121}