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)]
15#[derive_const(PartialEq, Eq)]
16pub enum Rv32ZkneInstruction<Reg> {
17    /// AES final round encryption step: SubBytes on one byte of rs2, rotated to the byte lane
18    /// selected by bs, XOR'd into rs1.
19    ///
20    /// `rd = rs1 ^ rol32(SBOX[(rs2 >> (bs*8)) & 0xff] as u32, bs*8)`
21    Aes32Esi {
22        rd: Reg,
23        rs1: Reg,
24        rs2: Reg,
25        bs: Rv32AesBs,
26    },
27    /// AES middle round encryption step: SubBytes + partial MixColumns on one byte of rs2, rotated
28    /// to the byte lane selected by bs, XOR'd into rs1.
29    ///
30    /// `rd = rs1 ^ rol32(MixColByte(SBOX[(rs2 >> (bs*8)) & 0xff]), bs*8)`
31    Aes32Esmi {
32        rd: Reg,
33        rs1: Reg,
34        rs2: Reg,
35        bs: Rv32AesBs,
36    },
37}
38
39/// Encoding layout (R-type, opcode 0x33, funct3 0x0):
40///
41/// ```text
42/// [31:30] bs       - 2-bit byte select
43/// [29:25] funct5   - 0b1_0001 (aes32esi) / 0b1_0011 (aes32esmi)
44/// [24:20] rs2
45/// [19:15] rs1
46/// [14:12] funct3   - 0b000
47/// [11:7]  rd
48/// [6:0]   opcode   - 0b011_0011 (OP)
49/// ```
50///
51/// Ratified match/mask values (from riscv-opcodes):
52///   MATCH_AES32ESI  = 0x2200_0033, MASK_AES32ESI  = 0x3e00_707f
53///   MATCH_AES32ESMI = 0x2600_0033, MASK_AES32ESMI = 0x3e00_707f
54///
55/// `rd` and `rs1` are independent fields. The assembler convention places
56/// the accumulator in both rd and rs1 (the `rt` pattern), but the hardware
57/// does not require rd == rs1 and the decoder must not enforce it.
58#[instruction]
59const impl<Reg> Instruction for Rv32ZkneInstruction<Reg>
60where
61    Reg: [const] Register<Type = u32>,
62{
63    const ALIGNMENT: u8 = align_of::<u32>() as u8;
64
65    type Reg = Reg;
66
67    #[inline(always)]
68    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic(const))]
69    fn try_decode(instruction: u32) -> Option<Self> {
70        let opcode = (instruction & 0b111_1111) as u8;
71        let rd_bits = ((instruction >> 7) & 0x1f) as u8;
72        let funct3 = ((instruction >> 12) & 0b111) as u8;
73        let rs1_bits = ((instruction >> 15) & 0x1f) as u8;
74        let rs2_bits = ((instruction >> 20) & 0x1f) as u8;
75        let funct5 = ((instruction >> 25) & 0b1_1111) as u8;
76        let bs_bits = ((instruction >> 30) & 0b11) as u8;
77
78        // R-type OP opcode only
79        if opcode != 0b011_0011 {
80            None?;
81        }
82        if funct3 != 0b000 {
83            None?;
84        }
85
86        let rd = Reg::from_bits(rd_bits)?;
87        let rs1 = Reg::from_bits(rs1_bits)?;
88        let rs2 = Reg::from_bits(rs2_bits)?;
89        let bs = Rv32AesBs::from_bits(bs_bits)?;
90
91        match funct5 {
92            // aes32esi:  bs[31:30] | 0b1_0001[29:25]
93            0b1_0001 => Some(Self::Aes32Esi { rd, rs1, rs2, bs }),
94            // aes32esmi: bs[31:30] | 0b1_0011[29:25]
95            0b1_0011 => Some(Self::Aes32Esmi { rd, rs1, rs2, bs }),
96            _ => None,
97        }
98    }
99
100    #[inline(always)]
101    fn size(&self) -> u8 {
102        size_of::<u32>() as u8
103    }
104}
105
106#[instruction]
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}