From c54b62354151f7c6fa568269c6d892eeed195b82 Mon Sep 17 00:00:00 2001
From: Sachymetsu <sachymetsu@tutamail.com>
Date: Sun, 14 Jun 2026 09:43:37 +0200
Subject: [PATCH] Optimise advance_position & begin_op

Make use of better codegen techniques for advance_position to improve const-time behaviour
and compiler optimisations. Additional slice assertion for begin_op to do the same to
improve compiler optimisations.
---
 wharrgarbl-neko/src/lib.rs       | 21 +++++++++++++++------
 wharrgarbl-neko/src/operators.rs |  2 ++
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/wharrgarbl-neko/src/lib.rs b/wharrgarbl-neko/src/lib.rs
index 89c1a38..4c84802 100644
--- a/wharrgarbl-neko/src/lib.rs
+++ b/wharrgarbl-neko/src/lib.rs
@@ -80,12 +80,13 @@ impl<Sec: NekoSec> NekoState<Sec> {
         const {
             assert!(Sec::BLOCK_RATE < keccak::PLEN);
             assert!(Sec::BLOCK_RATE < u8::MAX as usize);
+            assert!(Sec::POS_RATE < u8::MAX as usize);
         };
 
         // OPS stack MUST be initialised with the INIT flag as the first op.
         let mut state: Array<u64, U25> = Array([0u64; keccak::PLEN]);
 
-        state[Sec::BLOCK_RATE].as_mut_bytes()[4] ^= ops::INIT.bits();
+        state[Sec::BLOCK_RATE].as_mut_bytes()[MAX_OPS] ^= ops::INIT.bits();
 
         let mut neko = Self {
             // The buffer state is the zeroed buffer layout for KeccakF1600: [u64; 25]
@@ -146,9 +147,15 @@ impl<Sec: NekoSec> NekoState<Sec> {
 
     #[inline(always)]
     fn advance_position(&mut self, advance: usize) {
-        let updated = self.position + advance;
-        assert!(updated <= Sec::POS_RATE);
-        self.position = updated;
+        let updated = (self.raw_position() as usize) + advance;
+        debug_assert!(updated <= Sec::POS_RATE);
+        if updated <= Sec::POS_RATE {
+            self.position = updated;
+        } else {
+            // SAFETY: the type enforces that `updated` is always smaller than
+            // `RATE * U64_CHUNK` & `u8::MAX`
+            unsafe { core::hint::unreachable_unchecked() };
+        }
     }
 
     #[inline(always)]
@@ -172,13 +179,14 @@ impl<Sec: NekoSec> NekoState<Sec> {
         self.ops_count += 1;
 
         // Encode the opflags to the available stack slot.
-        self.state[Sec::BLOCK_RATE].as_mut_bytes()[4..][op_index] ^= red_flags.bits();
+        self.state[Sec::BLOCK_RATE].as_mut_bytes()[4..][..MAX_OPS][op_index] ^= red_flags.bits();
     }
 
     fn permutation_p12(&mut self, continuation: OpFlags) {
         const {
             assert!(Sec::BLOCK_RATE < keccak::PLEN);
             assert!(Sec::BLOCK_RATE < u8::MAX as usize);
+            assert!(Sec::POS_RATE < u8::MAX as usize);
         };
 
         let position = self.raw_position();
@@ -192,7 +200,7 @@ impl<Sec: NekoSec> NekoState<Sec> {
 
         // XOR the permuter block into the BLOCK_RATE block, which concats the padding, ops count
         // & position, with the encoded ops stack.
-        self.state[Sec::BLOCK_RATE].as_mut_bytes()[0..4]
+        self.state[Sec::BLOCK_RATE].as_mut_bytes()[0..MAX_OPS]
             .iter_mut()
             .zip(permuter)
             .for_each(|(state, perm)| *state ^= perm);
@@ -212,6 +220,7 @@ impl<Sec: NekoSec> NekoState<Sec> {
         const {
             assert!(Sec::RATCHET < keccak::PLEN);
             assert!(Sec::RATCHET < u8::MAX as usize);
+            assert!(Sec::POS_RATE < u8::MAX as usize);
         };
         // Select the amount of bytes to zero, according to Security level
         // 128 bits = 16 bytes = 2 blocks to zero out to achieve forward secrecy
diff --git a/wharrgarbl-neko/src/operators.rs b/wharrgarbl-neko/src/operators.rs
index 738906c..8db3692 100644
--- a/wharrgarbl-neko/src/operators.rs
+++ b/wharrgarbl-neko/src/operators.rs
@@ -18,6 +18,7 @@ impl<'s, S: NekoSec> NekoOperateMut<'s, S> {
         const {
             assert!(S::BLOCK_RATE < keccak::PLEN);
             assert!(S::BLOCK_RATE < u8::MAX as usize);
+            assert!(S::POS_RATE < u8::MAX as usize);
         };
 
         while !self.data.is_empty() {
@@ -88,6 +89,7 @@ impl<'s, S: NekoSec> NekoOperate<'s, S> {
         const {
             assert!(S::BLOCK_RATE < keccak::PLEN);
             assert!(S::BLOCK_RATE < u8::MAX as usize);
+            assert!(S::POS_RATE < u8::MAX as usize);
         };
 
         while !self.data.is_empty() {
-- 
2.51.2