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

1//! Opaque helpers for RV32 Zkne extension
2
3use ab_riscv_primitives::prelude::*;
4
5/// Software fallback for aes32esi and aes32esmi.
6///
7/// Both instructions share the same S-box and MixColumn machinery; the only difference is whether
8/// forward MixColumns is applied.
9#[cfg(not(all(not(miri), target_arch = "riscv32", target_feature = "zkne")))]
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) 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) 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) 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#[inline(always)]
80#[doc(hidden)]
81#[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
82pub fn aes32esi(rs1: u32, rs2: u32, bs: Rv32AesBs) -> u32 {
83    // TODO: Miri is excluded because corresponding intrinsic is not implemented there
84    cfg_select! {
85        all(
86            not(miri),
87            target_arch = "riscv32",
88            target_feature = "zkne"
89        ) => {
90            // SAFETY: Compile-time checked for supported feature
91            unsafe {
92                core::arch::riscv32::aes32esi(rs1, rs2, u8::from(bs))
93            }
94        }
95        _ => { soft::aes32esi(rs1, rs2, u8::from(bs)) }
96    }
97}
98
99#[inline(always)]
100#[doc(hidden)]
101#[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
102pub fn aes32esmi(rs1: u32, rs2: u32, bs: Rv32AesBs) -> u32 {
103    // TODO: Miri is excluded because corresponding intrinsic is not implemented there
104    cfg_select! {
105        all(
106            not(miri),
107            target_arch = "riscv32",
108            target_feature = "zkne"
109        ) => {
110            // SAFETY: Compile-time checked for supported feature
111            unsafe {
112                core::arch::riscv32::aes32esmi(rs1, rs2, u8::from(bs))
113            }
114        }
115        _ => { soft::aes32esmi(rs1, rs2, u8::from(bs)) }
116    }
117}
ab_riscv_interpreter/rv32/zk/zkn/zkne/rv32_zkne_helpers.rs

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