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 {
89                // The intrinsic takes the byte select as a const generic
90                match bs {
91                    Rv32AesBs::B0 => core::arch::riscv32::aes32esi::<0>(rs1, rs2),
92                    Rv32AesBs::B1 => core::arch::riscv32::aes32esi::<1>(rs1, rs2),
93                    Rv32AesBs::B2 => core::arch::riscv32::aes32esi::<2>(rs1, rs2),
94                    Rv32AesBs::B3 => core::arch::riscv32::aes32esi::<3>(rs1, rs2),
95                }
96            }
97        }
98        _ => soft::aes32esi(rs1, rs2, u8::from(bs)),
99    }
100}
101
102#[const_fn_specialization]
103#[inline(always)]
104#[doc(hidden)]
105#[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
106pub fn aes32esmi(rs1: u32, rs2: u32, bs: Rv32AesBs) -> u32 {
107    // TODO: Miri is excluded because corresponding intrinsic is not implemented there
108    cfg_select! {
109        all(not(miri), target_arch = "riscv32", target_feature = "zkne") => {
110            // SAFETY: Compile-time checked for supported feature
111            unsafe {
112                // The intrinsic takes the byte select as a const generic
113                match bs {
114                    Rv32AesBs::B0 => core::arch::riscv32::aes32esmi::<0>(rs1, rs2),
115                    Rv32AesBs::B1 => core::arch::riscv32::aes32esmi::<1>(rs1, rs2),
116                    Rv32AesBs::B2 => core::arch::riscv32::aes32esmi::<2>(rs1, rs2),
117                    Rv32AesBs::B3 => core::arch::riscv32::aes32esmi::<3>(rs1, rs2),
118                }
119            }
120        }
121        _ => soft::aes32esmi(rs1, rs2, u8::from(bs)),
122    }
123}
124
125#[const_fn_specialization]
126#[inline(always)]
127#[doc(hidden)]
128#[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
129pub const fn aes32esi(rs1: u32, rs2: u32, bs: Rv32AesBs) -> u32 {
130    soft::aes32esi(rs1, rs2, bs as u8)
131}
132
133#[const_fn_specialization]
134#[inline(always)]
135#[doc(hidden)]
136#[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
137pub const fn aes32esmi(rs1: u32, rs2: u32, bs: Rv32AesBs) -> u32 {
138    soft::aes32esmi(rs1, rs2, bs as u8)
139}
ab_riscv_interpreter/rv32/zk/zkn/zkne/rv32_zkne_helpers.rs

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