ab_riscv_primitives/instructions/rv64/zk/zkn/
zknd.rs

1//! RV64 Zknd extension
2
3#[cfg(test)]
4mod tests;
5
6use crate::hart::HartConfig;
7use crate::instructions::{ImplementedExtension, Instruction, IsaExtension};
8use crate::registers::general_purpose::Register;
9use ab_riscv_macros::instruction;
10use core::fmt;
11
12/// AES key schedule round constant number
13#[derive(Debug, Clone, Copy)]
14#[derive_const(PartialEq, Eq)]
15#[repr(u8)]
16pub enum Rv64ZkndKsRnum {
17    R0 = 0x0,
18    R1 = 0x1,
19    R2 = 0x2,
20    R3 = 0x3,
21    R4 = 0x4,
22    R5 = 0x5,
23    R6 = 0x6,
24    R7 = 0x7,
25    R8 = 0x8,
26    R9 = 0x9,
27    Final = 0xA,
28}
29
30const impl From<Rv64ZkndKsRnum> for u8 {
31    #[inline(always)]
32    fn from(rnum: Rv64ZkndKsRnum) -> Self {
33        rnum as u8
34    }
35}
36
37const impl From<Rv64ZkndKsRnum> for usize {
38    #[inline(always)]
39    fn from(rnum: Rv64ZkndKsRnum) -> Self {
40        usize::from(rnum as u8)
41    }
42}
43
44impl fmt::Display for Rv64ZkndKsRnum {
45    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
46        fmt::Display::fmt(&(*self as u8), f)
47    }
48}
49
50impl Rv64ZkndKsRnum {
51    /// Round constants `RC[0..=9]`, indexed by rnum (0-based).
52    /// `RC[rnum]` corresponds to FIPS 197 `Rcon[rnum+1]`.
53    const RCON: [u8; 10] = [0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36];
54
55    /// Create from raw bits
56    #[inline(always)]
57    pub const fn from_bits(bits: u8) -> Option<Self> {
58        match bits {
59            0x0 => Some(Self::R0),
60            0x1 => Some(Self::R1),
61            0x2 => Some(Self::R2),
62            0x3 => Some(Self::R3),
63            0x4 => Some(Self::R4),
64            0x5 => Some(Self::R5),
65            0x6 => Some(Self::R6),
66            0x7 => Some(Self::R7),
67            0x8 => Some(Self::R8),
68            0x9 => Some(Self::R9),
69            0xA => Some(Self::Final),
70            _ => None,
71        }
72    }
73
74    /// Round constant (unless final)
75    #[inline(always)]
76    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
77    pub const fn constant(self) -> Option<u8> {
78        if matches!(self, Rv64ZkndKsRnum::Final) {
79            None
80        } else {
81            Some(Self::RCON[usize::from(self)])
82        }
83    }
84}
85
86/// RISC-V RV64 instructions shared between Zknd and Zkne extensions (AES key schedule).
87///
88/// Not an extension on its own, only exists such that both Zknd and Zkne can inherit it.
89#[instruction]
90#[derive(Debug, Clone, Copy)]
91#[derive_const(PartialEq, Eq)]
92#[rustfmt::skip]
93pub enum Rv64ZkndZkneSharedInstruction<Hart>
94where
95    Hart: HartConfig,
96{
97    /// AES key schedule step 1 (rnum in 0..=10)
98    Aes64Ks1i {
99        rd: Hart::Reg,
100        rs1: Hart::Reg,
101        rnum: Rv64ZkndKsRnum,
102    },
103    /// AES key schedule step 2
104    Aes64Ks2 { rd: Hart::Reg, rs1: Hart::Reg, rs2: Hart::Reg },
105}
106
107#[instruction]
108const impl<Reg, Hart> Instruction for Rv64ZkndZkneSharedInstruction<Hart>
109where
110    Reg: [const] Register<Type = u64>,
111    Hart: [const] HartConfig<Reg = Reg>,
112{
113    const OWN_ISA_EXTENSIONS: &'static [IsaExtension] = &[];
114
115    const ALIGNMENT: u8 = align_of::<u32>() as u8;
116
117    type Hart = Hart;
118
119    #[inline(always)]
120    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic(const))]
121    fn try_decode(instruction: u32) -> Option<Self> {
122        let opcode = (instruction & 0b111_1111) as u8;
123        let rd_bits = ((instruction >> 7) & 0x1f) as u8;
124        let funct3 = ((instruction >> 12) & 0b111) as u8;
125        let rs1_bits = ((instruction >> 15) & 0x1f) as u8;
126        let rs2_bits = ((instruction >> 20) & 0x1f) as u8;
127        let funct7 = ((instruction >> 25) & 0b111_1111) as u8;
128
129        match opcode {
130            // R-type: OP opcode (0x33)
131            //   aes64ks2: funct7=0b011_1111, funct3=0 -> MATCH=0x7e00_0033
132            0b011_0011 => {
133                if funct3 != 0b000 {
134                    None?;
135                }
136                let rd = Reg::from_bits(rd_bits)?;
137                let rs1 = Reg::from_bits(rs1_bits)?;
138                let rs2 = Reg::from_bits(rs2_bits)?;
139                match funct7 {
140                    0b011_1111 => Some(Self::Aes64Ks2 { rd, rs1, rs2 }),
141                    _ => None,
142                }
143            }
144            // I-type: OP-IMM opcode (0x13), funct3=0b001
145            //   aes64ks1i: imm[11:5]=0b001_1000, imm[4]=1, imm[3:0]=rnum     -> MATCH=0x3100_1013+
146            0b001_0011 => {
147                if funct3 != 0b001 {
148                    None?;
149                }
150                let rd = Reg::from_bits(rd_bits)?;
151                let rs1 = Reg::from_bits(rs1_bits)?;
152                let imm12 = instruction >> 20;
153                if (imm12 >> 5) == 0b001_1000 && (imm12 & 0b1_0000) != 0 {
154                    // bits[11:5]=0b0011000, bit[4]=1, bits[3:0]=rnum
155                    let rnum = Rv64ZkndKsRnum::from_bits((imm12 & 0xf) as u8)?;
156                    Some(Self::Aes64Ks1i { rd, rs1, rnum })
157                } else {
158                    None
159                }
160            }
161            _ => None,
162        }
163    }
164
165    #[inline(always)]
166    fn size(&self) -> u8 {
167        size_of::<u32>() as u8
168    }
169}
170
171#[instruction]
172impl<Reg, Hart> fmt::Display for Rv64ZkndZkneSharedInstruction<Hart>
173where
174    Reg: fmt::Display,
175    Hart: HartConfig<Reg = Reg>,
176{
177    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
178        match self {
179            Self::Aes64Ks1i { rd, rs1, rnum } => write!(f, "aes64ks1i {rd}, {rs1}, {rnum}"),
180            Self::Aes64Ks2 { rd, rs1, rs2 } => write!(f, "aes64ks2 {rd}, {rs1}, {rs2}"),
181        }
182    }
183}
184
185/// RISC-V RV64 Zknd instructions (AES decryption and key schedule)
186#[instruction(
187    reorder = [Aes64Ds, Aes64Dsm, Aes64Im],
188    inherit = [Rv64ZkndZkneSharedInstruction],
189)]
190#[derive(Debug, Clone, Copy)]
191#[derive_const(PartialEq, Eq)]
192#[rustfmt::skip]
193pub enum Rv64ZkndInstruction<Hart>
194where
195    Hart: HartConfig,
196{
197    /// AES final round decryption: InvShiftRows + InvSubBytes, no MixColumns
198    Aes64Ds { rd: Hart::Reg, rs1: Hart::Reg, rs2: Hart::Reg },
199    /// AES middle round decryption: InvShiftRows + InvSubBytes + InvMixColumns
200    Aes64Dsm { rd: Hart::Reg, rs1: Hart::Reg, rs2: Hart::Reg },
201    /// AES inverse MixColumns on each 32-bit word of rs1
202    Aes64Im { rd: Hart::Reg, rs1: Hart::Reg },
203}
204
205#[instruction]
206const impl<Reg, Hart> Instruction for Rv64ZkndInstruction<Hart>
207where
208    Reg: [const] Register<Type = u64>,
209    Hart: [const] HartConfig<Reg = Reg>,
210{
211    const OWN_ISA_EXTENSIONS: &'static [IsaExtension] = &[IsaExtension::new("zknd", 1, 0)];
212
213    const ALIGNMENT: u8 = align_of::<u32>() as u8;
214
215    type Hart = Hart;
216
217    #[inline(always)]
218    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic(const))]
219    fn try_decode(instruction: u32) -> Option<Self> {
220        let opcode = (instruction & 0b111_1111) as u8;
221        let rd_bits = ((instruction >> 7) & 0x1f) as u8;
222        let funct3 = ((instruction >> 12) & 0b111) as u8;
223        let rs1_bits = ((instruction >> 15) & 0x1f) as u8;
224        let rs2_bits = ((instruction >> 20) & 0x1f) as u8;
225        let funct7 = ((instruction >> 25) & 0b111_1111) as u8;
226
227        match opcode {
228            // R-type: OP opcode (0x33)
229            //   aes64ds:  funct7=0b001_1101, funct3=0 -> MATCH=0x3a00_0033
230            //   aes64dsm: funct7=0b001_1111, funct3=0 -> MATCH=0x3e00_0033
231            0b011_0011 => {
232                if funct3 != 0b000 {
233                    None?;
234                }
235                let rd = Reg::from_bits(rd_bits)?;
236                let rs1 = Reg::from_bits(rs1_bits)?;
237                let rs2 = Reg::from_bits(rs2_bits)?;
238                match funct7 {
239                    0b001_1101 => Some(Self::Aes64Ds { rd, rs1, rs2 }),
240                    0b001_1111 => Some(Self::Aes64Dsm { rd, rs1, rs2 }),
241                    _ => None,
242                }
243            }
244            // I-type: OP-IMM opcode (0x13), funct3=0b001
245            //   aes64im:   imm[11:0]=0x300  (funct7=0b001_1000, rs2=0b0_0000) -> MATCH=0x3000_1013
246            0b001_0011 => {
247                if funct3 != 0b001 {
248                    None?;
249                }
250                let rd = Reg::from_bits(rd_bits)?;
251                let rs1 = Reg::from_bits(rs1_bits)?;
252                let imm12 = instruction >> 20;
253                if imm12 == 0x300 {
254                    Some(Self::Aes64Im { rd, rs1 })
255                } else {
256                    None
257                }
258            }
259            _ => None,
260        }
261    }
262
263    #[inline(always)]
264    fn size(&self) -> u8 {
265        size_of::<u32>() as u8
266    }
267}
268
269#[instruction]
270impl<Reg, Hart> fmt::Display for Rv64ZkndInstruction<Hart>
271where
272    Reg: fmt::Display + Copy,
273    Hart: HartConfig<Reg = Reg>,
274{
275    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
276        match self {
277            Self::Aes64Ds { rd, rs1, rs2 } => write!(f, "aes64ds {rd}, {rs1}, {rs2}"),
278            Self::Aes64Dsm { rd, rs1, rs2 } => write!(f, "aes64dsm {rd}, {rs1}, {rs2}"),
279            Self::Aes64Im { rd, rs1 } => write!(f, "aes64im {rd}, {rs1}"),
280        }
281    }
282}
ab_riscv_primitives/instructions/rv64/zk/zkn/zknd.rs

ab_riscv_primitives/instructions/rv64/zk/zkn/
zknd.rs
