From 98470d6373db9c8fb05dd91d6acce6c78a2e4fb7 Mon Sep 17 00:00:00 2001
From: monica <mamoniot@protonmail.com>
Date: Wed, 22 Mar 2023 11:51:10 -0400
Subject: [PATCH] saved

---
 zssp/src/fragged.rs | 55 +++++++++++++++++++++------------------------
 1 file changed, 25 insertions(+), 30 deletions(-)

diff --git a/zssp/src/fragged.rs b/zssp/src/fragged.rs
index 73eca9682..8c4783e34 100644
--- a/zssp/src/fragged.rs
+++ b/zssp/src/fragged.rs
@@ -38,17 +38,13 @@ impl<Fragment, const MAX_FRAGMENTS: usize> Fragged<Fragment, MAX_FRAGMENTS> {
         // that the array of MaybeUninit<Fragment> can be freely cast into an array of
         // Fragment. They also check that the maximum number of fragments is not too large
         // for the fact that we use bits in a u64 to track which fragments are received.
-        debug_assert!(MAX_FRAGMENTS <= 64);
-        debug_assert_eq!(size_of::<MaybeUninit<Fragment>>(), size_of::<Fragment>());
-        debug_assert_eq!(
+        assert!(MAX_FRAGMENTS <= 64);
+        assert_eq!(size_of::<MaybeUninit<Fragment>>(), size_of::<Fragment>());
+        assert_eq!(
             size_of::<[MaybeUninit<Fragment>; MAX_FRAGMENTS]>(),
             size_of::<[Fragment; MAX_FRAGMENTS]>()
         );
-        Self {
-            have: AtomicU64::new(0),
-            counter_want: RwLock::new((0, 0)),
-            frags: unsafe { zeroed() },
-        }
+        unsafe { zeroed() }
     }
 
     /// Add a fragment and return an assembled packet container if all fragments have been received.
@@ -56,21 +52,22 @@ impl<Fragment, const MAX_FRAGMENTS: usize> Fragged<Fragment, MAX_FRAGMENTS> {
     /// When a fully assembled packet is returned the internal state is reset and this object can
     /// be reused to assemble another packet.
     #[inline(always)]
-    pub fn assemble(&self, counter: u64, fragment: Fragment, fragment_no: u8, fragment_count: u8) -> Option<Assembled<Fragment, MAX_FRAGMENTS>> {
+    pub fn assemble(&mut self, counter: u64, fragment: Fragment, fragment_no: u8, fragment_count: u8) -> Option<Assembled<Fragment, MAX_FRAGMENTS>> {
         if fragment_no < fragment_count && (fragment_count as usize) <= MAX_FRAGMENTS {
             let r = self.counter_want.read().unwrap();
-            let (cur_counter, want) = *r;
-            let mut r_guard = Some(r);
+            let (mut r_counter, mut r_want) = *r;
+            let mut _r_guard = Some(r);
             let mut _w_guard = None;
+
             // If the counter has changed, reset the structure to receive a new packet.
-            if counter != cur_counter {
-                r_guard.take();
+            if counter != r_counter {
+                _r_guard.take();
                 let mut w = self.counter_want.write().unwrap();
-                let (cur_counter, _) = *w;
-                if counter != cur_counter {
+                let (w_counter, w_want) = *w;
+                if counter != w_counter {
                     if needs_drop::<Fragment>() {
                         let mut have = self.have.load(Ordering::Relaxed);
-                        if have != want {
+                        if have != w_want {
                             let mut i = 0;
                             while have != 0 {
                                 if (have & 1) != 0 {
@@ -82,28 +79,26 @@ impl<Fragment, const MAX_FRAGMENTS: usize> Fragged<Fragment, MAX_FRAGMENTS> {
                             }
                         }
                     }
-
+                    r_counter = counter;
+                    r_want = 1u64.wrapping_shl(fragment_count as u32) - 1;
+                    *w = (r_counter, r_want);
                     self.have.store(0, Ordering::Relaxed);
-                    let want = 1u64.wrapping_shl(fragment_count as u32) - 1;
-                    *w = (counter, want);
                 }
                 _w_guard = Some(w);
             }
 
-            // Validate inputs, mainly that the fragment_count matches what we expect
-            // and that fragment_no is actually a wanted fragment
             let got = 1u64.wrapping_shl(fragment_no as u32);
-            if want & got > 0 && 1u64.wrapping_shl(fragment_count as u32) == want + 1 {
-
-                // Now we check whether fragment_no is missing
+            if got & r_want > 0 && 1u64.wrapping_shl(fragment_count as u32) - 1 == r_want {
                 let have = self.have.fetch_or(got, Ordering::Relaxed);
                 if have & got == 0 {
+
                     unsafe {
                         (*self.frags.get()).get_unchecked_mut(fragment_no as usize).write(fragment);
                     }
-                    if (have | got) == want {
-                        // The fragments are effectively moved into the Assembled<> container and returned.
-                        // That container will drop them when it is dropped.
+                    if have | got == r_want {
+                        // Setting 'have' to 0 resets the state of this object, and the fragments
+                        // are effectively moved into the Assembled<> container and returned. That
+                        // container will drop them when it is dropped.
                         return Some(Assembled(unsafe { std::mem::transmute_copy(&self.frags) }, fragment_count as usize));
                     }
                 }
@@ -117,10 +112,10 @@ impl<Fragment, const MAX_FRAGMENTS: usize> Drop for Fragged<Fragment, MAX_FRAGME
     #[inline(always)]
     fn drop(&mut self) {
         if needs_drop::<Fragment>() {
-            let r = self.counter_want.read().unwrap();
-            let (_, want) = *r;
             let mut have = self.have.load(Ordering::Relaxed);
-            if have != want {
+            let r = self.counter_want.read().unwrap();
+            let (_, r_want) = *r;
+            if have != r_want {
                 let mut i = 0;
                 while have != 0 {
                     if (have & 1) != 0 {