From a4d178c95e0f7058260e0561de1cdabf0e92a585 Mon Sep 17 00:00:00 2001
From: EAimTY <ea.imty@gmail.com>
Date: Wed, 25 Jan 2023 20:08:05 +0900
Subject: [PATCH] assembling packet to existed buffer

---
 tuic/src/model/mod.rs    | 55 ++++++++++++++++++++++++++--------------
 tuic/src/model/packet.rs |  7 ++---
 2 files changed, 38 insertions(+), 24 deletions(-)

diff --git a/tuic/src/model/mod.rs b/tuic/src/model/mod.rs
index 4cd8cb6..f526e41 100644
--- a/tuic/src/model/mod.rs
+++ b/tuic/src/model/mod.rs
@@ -5,6 +5,7 @@ use crate::protocol::{
 use parking_lot::Mutex;
 use std::{
     collections::HashMap,
+    mem,
     sync::{
         atomic::{AtomicU16, Ordering},
         Arc, Weak,
@@ -202,7 +203,7 @@ where
         Dissociate::<side::Rx>::new(assoc_id)
     }
 
-    fn insert<A>(
+    fn insert(
         &mut self,
         assoc_id: u16,
         pkt_id: u16,
@@ -211,10 +212,7 @@ where
         size: u16,
         addr: Address,
         data: B,
-    ) -> Result<Option<(A, Address)>, AssembleError>
-    where
-        A: Assembled<B>,
-    {
+    ) -> Result<Option<Assemblable<B>>, AssembleError> {
         self.sessions
             .entry(assoc_id)
             .or_insert_with(|| UdpSession::new(self.task_associate_count.register()))
@@ -273,7 +271,7 @@ where
         Packet::<side::Rx, B>::new(sessions, assoc_id, pkt_id, frag_total, frag_id, size, addr)
     }
 
-    fn insert<A>(
+    fn insert(
         &mut self,
         pkt_id: u16,
         frag_total: u8,
@@ -281,10 +279,7 @@ where
         size: u16,
         addr: Address,
         data: B,
-    ) -> Result<Option<(A, Address)>, AssembleError>
-    where
-        A: Assembled<B>,
-    {
+    ) -> Result<Option<Assemblable<B>>, AssembleError> {
         let res = self
             .pkt_buf
             .entry(pkt_id)
@@ -328,17 +323,14 @@ where
         }
     }
 
-    fn insert<A>(
+    fn insert(
         &mut self,
         frag_total: u8,
         frag_id: u8,
         size: u16,
         addr: Address,
         data: B,
-    ) -> Result<Option<(A, Address)>, AssembleError>
-    where
-        A: Assembled<B>,
-    {
+    ) -> Result<Option<Assemblable<B>>, AssembleError> {
         if data.as_ref().len() != size as usize {
             return Err(AssembleError::InvalidFragmentSize);
         }
@@ -363,20 +355,45 @@ where
         }
 
         if self.frag_received == self.frag_total {
-            let iter = self.buf.iter_mut().map(|x| x.take().unwrap());
-            Ok(Some((A::assemble(iter)?, self.addr.take())))
+            Ok(Some(Assemblable::new(
+                mem::take(&mut self.buf),
+                self.addr.take(),
+            )))
         } else {
             Ok(None)
         }
     }
 }
 
-pub trait Assembled<B>
+pub struct Assemblable<B> {
+    buf: Vec<Option<B>>,
+    addr: Address,
+}
+
+impl<B> Assemblable<B>
+where
+    B: AsRef<[u8]>,
+{
+    fn new(buf: Vec<Option<B>>, addr: Address) -> Self {
+        Self { buf, addr }
+    }
+
+    pub fn assemble<A>(self, buf: &mut A) -> Address
+    where
+        A: Assembler<B>,
+    {
+        let data = self.buf.into_iter().map(|b| b.unwrap());
+        buf.assemble(data);
+        self.addr
+    }
+}
+
+pub trait Assembler<B>
 where
     Self: Sized,
     B: AsRef<[u8]>,
 {
-    fn assemble(buf: impl IntoIterator<Item = B>) -> Result<Self, AssembleError>;
+    fn assemble(&mut self, data: impl IntoIterator<Item = B>);
 }
 
 #[derive(Debug, Error)]
diff --git a/tuic/src/model/packet.rs b/tuic/src/model/packet.rs
index d4c4483..181fc7c 100644
--- a/tuic/src/model/packet.rs
+++ b/tuic/src/model/packet.rs
@@ -1,6 +1,6 @@
 use super::{
     side::{self, Side},
-    AssembleError, Assembled, UdpSessions,
+    Assemblable, AssembleError, UdpSessions,
 };
 use crate::protocol::{Address, Header, Packet as PacketHeader};
 use parking_lot::Mutex;
@@ -77,10 +77,7 @@ where
         }
     }
 
-    pub fn assemble<A>(self, data: B) -> Result<Option<(A, Address)>, AssembleError>
-    where
-        A: Assembled<B>,
-    {
+    pub fn assemble(self, data: B) -> Result<Option<Assemblable<B>>, AssembleError> {
         let Side::Rx(rx) = self.inner else { unreachable!() };
         let mut sessions = rx.sessions.lock();