ab_riscv_interpreter/rv64/zk/zkn/zkne/
rv64_zkne_helpers.rs

1//! Opaque helpers for RV64 Zkne extension
2
3use const_fn_specialization::const_fn_specialization;
4
5cfg_select! {
6    all(not(miri), target_arch = "riscv64", target_feature = "zkne") => {
7        // Nothing, calling native intrinsics
8    }
9    all(
10        target_arch = "x86_64",
11        target_feature = "aes",
12        target_feature = "sse4.1"
13    ) => {
14        /// x86-64 AES-NI implementation
15        #[expect(
16            clippy::inline_modules,
17            reason = "Small internal API, it is more readable this way"
18        )]
19        mod x86_64 {
20            use core::arch::x86_64::{
21                _mm_aesenc_si128, _mm_aesenclast_si128, _mm_extract_epi64, _mm_set_epi64x,
22                _mm_setzero_si128,
23            };
24
25            /// `_mm_aesenclast_si128(state, zero)` computes ShiftRows + SubBytes then XORs with
26            /// the round key. Zero key -> no-op XOR, matching `aes64es`.
27            #[inline]
28            #[target_feature(enable = "aes,sse4.1")]
29            #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
30            pub(super) fn aes64es(rs1: u64, rs2: u64) -> u64 {
31                let state = _mm_set_epi64x(rs2.cast_signed(), rs1.cast_signed());
32                let zero = _mm_setzero_si128();
33                let result = _mm_aesenclast_si128(state, zero);
34                _mm_extract_epi64::<0>(result).cast_unsigned()
35            }
36
37            /// `_mm_aesenc_si128(state, zero)` computes ShiftRows + SubBytes + MixColumns then
38            /// XORs with the round key. Zero key -> no-op XOR, matching `aes64esm`.
39            #[inline]
40            #[target_feature(enable = "aes,sse4.1")]
41            #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
42            pub(super) fn aes64esm(rs1: u64, rs2: u64) -> u64 {
43                let state = _mm_set_epi64x(rs2.cast_signed(), rs1.cast_signed());
44                let zero = _mm_setzero_si128();
45                let result = _mm_aesenc_si128(state, zero);
46                _mm_extract_epi64::<0>(result).cast_unsigned()
47            }
48        }
49    }
50    all(target_arch = "aarch64", target_feature = "aes") => {
51        /// AArch64 AES implementation
52        ///
53        /// AESE XORs the round key first, then applies SubBytes + ShiftRows (note: ARM ShiftRows
54        /// direction matches the forward cipher). With a zero round key the XOR is a no-op,
55        /// leaving pure SubBytes + ShiftRows - identical to what `aes64es` requires.
56        #[expect(
57            clippy::inline_modules,
58            reason = "Small internal API, it is more readable this way"
59        )]
60        mod aarch64 {
61            use core::arch::aarch64::{
62                vaeseq_u8, vaesmcq_u8, vcombine_u64, vcreate_u64, vdupq_n_u8, vgetq_lane_u64,
63                vreinterpretq_u8_u64, vreinterpretq_u64_u8,
64            };
65
66            #[inline]
67            #[target_feature(enable = "aes")]
68            #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
69            pub(super) fn aes64es(rs1: u64, rs2: u64) -> u64 {
70                let state = vreinterpretq_u8_u64(vcombine_u64(vcreate_u64(rs1), vcreate_u64(rs2)));
71                let zero = vdupq_n_u8(0);
72                let result = vaeseq_u8(state, zero);
73                vgetq_lane_u64::<0>(vreinterpretq_u64_u8(result))
74            }
75
76            /// `vaesmcq_u8(vaeseq_u8(state, zero))` maps exactly to `aes64esm`
77            #[inline]
78            #[target_feature(enable = "aes")]
79            #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
80            pub(super) fn aes64esm(rs1: u64, rs2: u64) -> u64 {
81                let state = vreinterpretq_u8_u64(vcombine_u64(vcreate_u64(rs1), vcreate_u64(rs2)));
82                let zero = vdupq_n_u8(0);
83                let after_sub_shift = vaeseq_u8(state, zero);
84                let result = vaesmcq_u8(after_sub_shift);
85                vgetq_lane_u64::<0>(vreinterpretq_u64_u8(result))
86            }
87        }
88    }
89    _ => {
90        // Nothing, the software fallback below is always available
91    }
92}
93
94/// Software fallback for aes64es, aes64esm
95#[expect(
96    clippy::inline_modules,
97    reason = "Small internal API, it is more readable this way"
98)]
99mod soft {
100    use crate::const_utils::ConstRange;
101    use crate::rv32::zk::zkn::zknd::rv32_zknd_helpers::{SBOX, gmul};
102
103    #[inline(always)]
104    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
105    const fn mix_col(col: u32) -> u32 {
106        let s0 = col as u8;
107        let s1 = (col >> 8) as u8;
108        let s2 = (col >> 16) as u8;
109        let s3 = (col >> 24) as u8;
110        let r0 = gmul(s0, 0x02) ^ gmul(s1, 0x03) ^ s2 ^ s3;
111        let r1 = s0 ^ gmul(s1, 0x02) ^ gmul(s2, 0x03) ^ s3;
112        let r2 = s0 ^ s1 ^ gmul(s2, 0x02) ^ gmul(s3, 0x03);
113        let r3 = gmul(s0, 0x03) ^ s1 ^ s2 ^ gmul(s3, 0x02);
114        u32::from(r0) | (u32::from(r1) << 8) | (u32::from(r2) << 16) | (u32::from(r3) << 24)
115    }
116
117    /// Byte in row `row` of column `col` of the 128-bit state `(rs1, rs2)`
118    #[inline(always)]
119    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
120    const fn state_byte(rs1: u64, rs2: u64, col: usize, row: usize) -> u8 {
121        let word = if col < 2 { rs1 } else { rs2 };
122        (word >> ((col % 2) * 32 + row * 8)) as u8
123    }
124
125    /// Apply ShiftRows + SubBytes to the full 128-bit state `(rs1, rs2)` and return the
126    /// low 64-bit half of the result.
127    ///
128    /// State layout: column-major, little-endian 64-bit halves.
129    /// ShiftRows shifts row `r` left by `r` columns (cyclically over 4).
130    /// Output low half contains post-transform columns 0 and 1.
131    #[inline(always)]
132    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
133    pub(super) const fn aes64es(rs1: u64, rs2: u64) -> u64 {
134        let mut out = 0;
135        for c in ConstRange::new(0usize, 2) {
136            for r in ConstRange::new(0usize, 4) {
137                let src_col = (c + r) & 3;
138                let b = SBOX[state_byte(rs1, rs2, src_col, r) as usize];
139                out |= u64::from(b) << (c * 32 + r * 8);
140            }
141        }
142        out
143    }
144
145    #[inline(always)]
146    #[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
147    pub(super) const fn aes64esm(rs1: u64, rs2: u64) -> u64 {
148        let lo = aes64es(rs1, rs2);
149        let col0 = mix_col(lo as u32);
150        let col1 = mix_col((lo >> 32) as u32);
151        u64::from(col0) | (u64::from(col1) << 32)
152    }
153}
154
155#[const_fn_specialization]
156#[inline(always)]
157#[doc(hidden)]
158#[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
159pub fn aes64es(rs1: u64, rs2: u64) -> u64 {
160    // TODO: Miri is excluded because corresponding intrinsic is not implemented there
161    cfg_select! {
162        all(not(miri), target_arch = "riscv64", target_feature = "zkne") => {
163            // SAFETY: Compile-time checked for supported feature
164            unsafe { core::arch::riscv64::aes64es(rs1, rs2) }
165        }
166        all(
167            target_arch = "x86_64",
168            target_feature = "aes",
169            target_feature = "sse4.1"
170        ) => {
171            // SAFETY: Compile-time checked for supported feature
172            unsafe { x86_64::aes64es(rs1, rs2) }
173        }
174        all(target_arch = "aarch64", target_feature = "aes") => {
175            // SAFETY: Compile-time checked for supported feature
176            unsafe { aarch64::aes64es(rs1, rs2) }
177        }
178        _ => soft::aes64es(rs1, rs2),
179    }
180}
181
182#[const_fn_specialization]
183#[inline(always)]
184#[doc(hidden)]
185#[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
186pub fn aes64esm(rs1: u64, rs2: u64) -> u64 {
187    // TODO: Miri is excluded because corresponding intrinsic is not implemented there
188    cfg_select! {
189        all(not(miri), target_arch = "riscv64", target_feature = "zkne") => {
190            // SAFETY: Compile-time checked for supported feature
191            unsafe { core::arch::riscv64::aes64esm(rs1, rs2) }
192        }
193        all(
194            target_arch = "x86_64",
195            target_feature = "aes",
196            target_feature = "sse4.1"
197        ) => {
198            // SAFETY: Compile-time checked for supported feature
199            unsafe { x86_64::aes64esm(rs1, rs2) }
200        }
201        all(target_arch = "aarch64", target_feature = "aes") => {
202            // SAFETY: Compile-time checked for supported feature
203            unsafe { aarch64::aes64esm(rs1, rs2) }
204        }
205        _ => soft::aes64esm(rs1, rs2),
206    }
207}
208
209#[const_fn_specialization]
210#[inline(always)]
211#[doc(hidden)]
212#[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
213pub const fn aes64es(rs1: u64, rs2: u64) -> u64 {
214    soft::aes64es(rs1, rs2)
215}
216
217#[const_fn_specialization]
218#[inline(always)]
219#[doc(hidden)]
220#[cfg_attr(feature = "no-panic", no_panic_const::no_panic)]
221pub const fn aes64esm(rs1: u64, rs2: u64) -> u64 {
222    soft::aes64esm(rs1, rs2)
223}
ab_riscv_interpreter/rv64/zk/zkn/zkne/rv64_zkne_helpers.rs

ab_riscv_interpreter/rv64/zk/zkn/zkne/
rv64_zkne_helpers.rs
