ab_riscv_interpreter/rv32/zk/zkn/zkne/
rv32_zkne_helpers.rs

1//! Opaque helpers for RV32 Zkne extension
2
3use ab_riscv_primitives::prelude::*;
4use const_fn_specialization::const_fn_specialization;
5
6/// Software fallback for aes32esi and aes32esmi.
7///
8/// Both instructions share the same S-box and MixColumn machinery; the only difference is whether
9/// forward MixColumns is applied.
10#[expect(
11    clippy::inline_modules,
12    reason = "Small internal API, it is more readable this way"
13)]
14pub(in super::super) mod soft {
15    use crate::rv32::zk::zkn::zknd::rv32_zknd_helpers::{SBOX, gmul};
16
17    /// Compute the partial forward MixColumns contribution for a single substituted byte `b`.
18    ///
19    /// This is `aes_mixcolumn_byte_fwd` from the Sail reference:
20    /// the four output bytes of MixColumns when the input column has `b` in one position and zeros
21    /// elsewhere - packed into a little-endian `u32`.
22    ///
23    /// Column matrix multiply for MixColumns:
24    /// ```text
25    /// r0 = 0x02*b
26    /// r1 = 0x01*b
27    /// r2 = 0x01*b
28    /// r3 = 0x03*b
29    /// ```
30    #[inline(always)]
31    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
32    pub(super) const fn mix_col_byte(b: u8) -> u32 {
33        let r0 = u32::from(gmul(b, 0x02));
34        let r1 = u32::from(b);
35        let r2 = u32::from(b);
36        let r3 = u32::from(gmul(b, 0x03));
37        r0 | (r1 << 8) | (r2 << 16) | (r3 << 24)
38    }
39
40    /// `aes32esi rs1, rs2, bs`
41    ///
42    /// Pseudocode:
43    /// ```text
44    /// shamt = bs * 8
45    /// si    = (rs2 >> shamt) & 0xff
46    /// so    = SBOX[si] as u32
47    /// rd    = rs1 ^ rol32(so, shamt)
48    /// ```
49    #[inline(always)]
50    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
51    pub(super) const fn aes32esi(rs1: u32, rs2: u32, bs: u8) -> u32 {
52        let shamt = u32::from(bs) * 8;
53        let si = ((rs2 >> shamt) & 0xff) as u8;
54        let so = u32::from(SBOX[usize::from(si)]);
55        rs1 ^ so.rotate_left(shamt)
56    }
57
58    /// `aes32esmi rs1, rs2, bs`
59    ///
60    /// Pseudocode:
61    /// ```text
62    /// shamt = bs * 8
63    /// si    = (rs2 >> shamt) & 0xff
64    /// so    = SBOX[si]
65    /// mixed = mix_col_byte(so)
66    /// rd    = rs1 ^ rol32(mixed, shamt)
67    /// ```
68    #[inline(always)]
69    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
70    pub(super) const fn aes32esmi(rs1: u32, rs2: u32, bs: u8) -> u32 {
71        let shamt = u32::from(bs) * 8;
72        let si = ((rs2 >> shamt) & 0xff) as u8;
73        let so = SBOX[usize::from(si)];
74        let mixed = mix_col_byte(so);
75        rs1 ^ mixed.rotate_left(shamt)
76    }
77}
78
79#[const_fn_specialization]
80#[inline(always)]
81#[doc(hidden)]
82#[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
83pub fn aes32esi(rs1: u32, rs2: u32, bs: Rv32AesBs) -> u32 {
84    // TODO: Miri is excluded because corresponding intrinsic is not implemented there
85    cfg_select! {
86        all(not(miri), target_arch = "riscv32", target_feature = "zkne") => {
87            // SAFETY: Compile-time checked for supported feature
88            unsafe { core::arch::riscv32::aes32esi(rs1, rs2, u8::from(bs)) }
89        }
90        _ => soft::aes32esi(rs1, rs2, u8::from(bs)),
91    }
92}
93
94#[const_fn_specialization]
95#[inline(always)]
96#[doc(hidden)]
97#[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
98pub fn aes32esmi(rs1: u32, rs2: u32, bs: Rv32AesBs) -> u32 {
99    // TODO: Miri is excluded because corresponding intrinsic is not implemented there
100    cfg_select! {
101        all(not(miri), target_arch = "riscv32", target_feature = "zkne") => {
102            // SAFETY: Compile-time checked for supported feature
103            unsafe { core::arch::riscv32::aes32esmi(rs1, rs2, u8::from(bs)) }
104        }
105        _ => soft::aes32esmi(rs1, rs2, u8::from(bs)),
106    }
107}
108
109#[const_fn_specialization]
110#[inline(always)]
111#[doc(hidden)]
112#[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
113pub const fn aes32esi(rs1: u32, rs2: u32, bs: Rv32AesBs) -> u32 {
114    soft::aes32esi(rs1, rs2, bs as u8)
115}
116
117#[const_fn_specialization]
118#[inline(always)]
119#[doc(hidden)]
120#[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
121pub const fn aes32esmi(rs1: u32, rs2: u32, bs: Rv32AesBs) -> u32 {
122    soft::aes32esmi(rs1, rs2, bs as u8)
123}
ab_riscv_interpreter/rv32/zk/zkn/zkne/rv32_zkne_helpers.rs

ab_riscv_interpreter/rv32/zk/zkn/zkne/
rv32_zkne_helpers.rs
