ManagedVec iterator refactor, ManagedVecPayloadIterator

multiversx · Dec 4, 2024 · 0df3c0a · 0df3c0a
1 parent 1db87bb
commit 0df3c0a
Show file tree

Hide file tree

Showing 4 changed files with 134 additions and 91 deletions.
diff --git a/framework/base/src/types/managed/wrapped.rs b/framework/base/src/types/managed/wrapped.rs
@@ -20,6 +20,7 @@ mod managed_vec_item;
 mod managed_vec_item_nested_tuple;
 mod managed_vec_item_payload;
 mod managed_vec_iter_owned;
+mod managed_vec_iter_payload;
 mod managed_vec_iter_ref;
 mod managed_vec_ref;
 mod managed_vec_ref_mut;
@@ -54,6 +55,7 @@ pub use managed_vec_item::{
 pub use managed_vec_item_nested_tuple::ManagedVecItemNestedTuple;
 pub use managed_vec_item_payload::*;
 pub use managed_vec_iter_owned::ManagedVecOwnedIterator;
+pub use managed_vec_iter_payload::ManagedVecPayloadIterator;
 pub use managed_vec_iter_ref::ManagedVecRefIterator;
 pub use managed_vec_ref::ManagedVecRef;
 pub use managed_vec_ref_mut::ManagedVecRefMut;

diff --git a/framework/base/src/types/managed/wrapped/managed_vec_iter_owned.rs b/framework/base/src/types/managed/wrapped/managed_vec_iter_owned.rs
@@ -1,6 +1,6 @@
-use crate::api::ManagedTypeApi;
+use crate::{api::ManagedTypeApi, types::ManagedType};
 
-use super::{ManagedVec, ManagedVecItem, ManagedVecItemPayload};
+use super::{ManagedVec, ManagedVecItem, ManagedVecPayloadIterator};
 
 impl<M, T> IntoIterator for ManagedVec<M, T>
 where
@@ -19,9 +19,7 @@ where
     M: ManagedTypeApi,
     T: ManagedVecItem,
 {
-    managed_vec: ManagedVec<M, T>,
-    byte_start: usize,
-    byte_end: usize,
+    payload_iter: ManagedVecPayloadIterator<M, T::PAYLOAD>,
 }
 
 impl<M, T> ManagedVecOwnedIterator<M, T>
@@ -30,11 +28,10 @@ where
     T: ManagedVecItem,
 {
     pub(crate) fn new(managed_vec: ManagedVec<M, T>) -> Self {
-        let byte_end = managed_vec.byte_len();
-        ManagedVecOwnedIterator {
-            managed_vec,
-            byte_start: 0,
-            byte_end,
+        unsafe {
+            ManagedVecOwnedIterator {
+                payload_iter: ManagedVecPayloadIterator::new(managed_vec.forget_into_handle()),
+            }
         }
     }
 }
@@ -47,26 +44,12 @@ where
     type Item = T;
 
     fn next(&mut self) -> Option<T> {
-        // managedrev<t>  / reference type
-        let next_byte_start = self.byte_start + T::payload_size();
-        if next_byte_start > self.byte_end {
-            return None;
-        }
-
-        let mut payload = T::PAYLOAD::new_buffer();
-        let _ = self
-            .managed_vec
-            .buffer
-            .load_slice(self.byte_start, payload.payload_slice_mut());
-
-        self.byte_start = next_byte_start;
+        let payload = self.payload_iter.next()?;
         Some(T::read_from_payload(&payload))
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
-        let size = T::payload_size();
-        let remaining = (self.byte_end - self.byte_start) / size;
-        (remaining, Some(remaining))
+        self.payload_iter.size_hint()
     }
 }
 
@@ -83,32 +66,7 @@ where
     T: ManagedVecItem,
 {
     fn next_back(&mut self) -> Option<Self::Item> {
-        if self.byte_start + T::payload_size() > self.byte_end {
-            return None;
-        }
-        self.byte_end -= T::payload_size();
-
-        let mut payload = T::PAYLOAD::new_buffer();
-        let _ = self
-            .managed_vec
-            .buffer
-            .load_slice(self.byte_end, payload.payload_slice_mut());
-
+        let payload = self.payload_iter.next_back()?;
         Some(T::read_from_payload(&payload))
     }
 }
-
-impl<M, T> Clone for ManagedVecOwnedIterator<M, T>
-where
-    M: ManagedTypeApi,
-    T: ManagedVecItem + Clone,
-{
-    fn clone(&self) -> Self {
-        let byte_end = self.byte_end;
-        Self {
-            managed_vec: self.managed_vec.clone(),
-            byte_start: self.byte_start,
-            byte_end,
-        }
-    }
-}
diff --git a/framework/base/src/types/managed/wrapped/managed_vec_iter_payload.rs b/framework/base/src/types/managed/wrapped/managed_vec_iter_payload.rs
@@ -0,0 +1,103 @@
+use core::marker::PhantomData;
+
+use crate::api::{ManagedBufferApiImpl, ManagedTypeApi};
+
+use super::ManagedVecItemPayload;
+
+pub struct ManagedVecPayloadIterator<M, P>
+where
+    M: ManagedTypeApi,
+    P: ManagedVecItemPayload,
+{
+    pub(super) vec_handle: M::ManagedBufferHandle,
+    byte_start: usize,
+    byte_end: usize,
+    _phantom: PhantomData<P>,
+}
+
+impl<M, P> ManagedVecPayloadIterator<M, P>
+where
+    M: ManagedTypeApi,
+    P: ManagedVecItemPayload,
+{
+    /// Unsafe because it works with the managed vec handle directly, so does not take ownership into account.
+    pub(crate) unsafe fn new(vec_handle: M::ManagedBufferHandle) -> Self {
+        let byte_end = M::managed_type_impl().mb_len(vec_handle.clone());
+        ManagedVecPayloadIterator {
+            vec_handle,
+            byte_start: 0,
+            byte_end,
+            _phantom: PhantomData,
+        }
+    }
+
+    /// Unsafe because it works with the managed vec handle directly, so does not take ownership into account.
+    pub(super) unsafe fn clone_iter(&self) -> Self {
+        ManagedVecPayloadIterator {
+            vec_handle: self.vec_handle.clone(),
+            byte_start: self.byte_start,
+            byte_end: self.byte_end,
+            _phantom: PhantomData,
+        }
+    }
+}
+
+impl<M, P> Iterator for ManagedVecPayloadIterator<M, P>
+where
+    M: ManagedTypeApi,
+    P: ManagedVecItemPayload,
+{
+    type Item = P;
+
+    fn next(&mut self) -> Option<P> {
+        let next_byte_start = self.byte_start + P::payload_size();
+        if next_byte_start > self.byte_end {
+            return None;
+        }
+
+        let mut payload = P::new_buffer();
+        let _ = M::managed_type_impl().mb_load_slice(
+            self.vec_handle.clone(),
+            self.byte_start,
+            payload.payload_slice_mut(),
+        );
+
+        self.byte_start = next_byte_start;
+        Some(payload)
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let size = P::payload_size();
+        let remaining = (self.byte_end - self.byte_start) / size;
+        (remaining, Some(remaining))
+    }
+}
+
+impl<M, P> ExactSizeIterator for ManagedVecPayloadIterator<M, P>
+where
+    M: ManagedTypeApi,
+    P: ManagedVecItemPayload,
+{
+}
+
+impl<M, P> DoubleEndedIterator for ManagedVecPayloadIterator<M, P>
+where
+    M: ManagedTypeApi,
+    P: ManagedVecItemPayload,
+{
+    fn next_back(&mut self) -> Option<Self::Item> {
+        if self.byte_start + P::payload_size() > self.byte_end {
+            return None;
+        }
+        self.byte_end -= P::payload_size();
+
+        let mut payload = P::new_buffer();
+        let _ = M::managed_type_impl().mb_load_slice(
+            self.vec_handle.clone(),
+            self.byte_end,
+            payload.payload_slice_mut(),
+        );
+
+        Some(payload)
+    }
+}
diff --git a/framework/base/src/types/managed/wrapped/managed_vec_iter_ref.rs b/framework/base/src/types/managed/wrapped/managed_vec_iter_ref.rs
@@ -1,6 +1,8 @@
-use crate::api::ManagedTypeApi;
+use core::marker::PhantomData;
 
-use super::{ManagedVec, ManagedVecItem, ManagedVecItemPayload};
+use crate::{api::ManagedTypeApi, types::ManagedType};
+
+use super::{ManagedVec, ManagedVecItem, ManagedVecPayloadIterator};
 
 impl<'a, M, T> IntoIterator for &'a ManagedVec<M, T>
 where
@@ -19,9 +21,8 @@ where
     M: ManagedTypeApi,
     T: ManagedVecItem,
 {
-    managed_vec: &'a ManagedVec<M, T>,
-    byte_start: usize,
-    byte_end: usize,
+    payload_iter: ManagedVecPayloadIterator<M, T::PAYLOAD>,
+    _phantom: PhantomData<&'a ManagedVec<M, T>>,
 }
 
 impl<'a, M, T> ManagedVecRefIterator<'a, M, T>
@@ -30,10 +31,11 @@ where
     T: ManagedVecItem,
 {
     pub(crate) fn new(managed_vec: &'a ManagedVec<M, T>) -> Self {
-        ManagedVecRefIterator {
-            managed_vec,
-            byte_start: 0,
-            byte_end: managed_vec.byte_len(),
+        unsafe {
+            ManagedVecRefIterator {
+                payload_iter: ManagedVecPayloadIterator::new(managed_vec.get_handle()),
+                _phantom: PhantomData,
+            }
         }
     }
 }
@@ -46,25 +48,12 @@ where
     type Item = T::Ref<'a>;
 
     fn next(&mut self) -> Option<Self::Item> {
-        let next_byte_start = self.byte_start + T::payload_size();
-        if next_byte_start > self.byte_end {
-            return None;
-        }
-
-        let mut payload = T::PAYLOAD::new_buffer();
-        let _ = self
-            .managed_vec
-            .buffer
-            .load_slice(self.byte_start, payload.payload_slice_mut());
-
-        self.byte_start = next_byte_start;
-
+        let payload = self.payload_iter.next()?;
         unsafe { Some(T::borrow_from_payload(&payload)) }
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
-        let remaining = (self.byte_end - self.byte_start) / T::payload_size();
-        (remaining, Some(remaining))
+        self.payload_iter.size_hint()
     }
 }
 
@@ -81,17 +70,7 @@ where
     T: ManagedVecItem,
 {
     fn next_back(&mut self) -> Option<Self::Item> {
-        if self.byte_start + T::payload_size() > self.byte_end {
-            return None;
-        }
-        self.byte_end -= T::payload_size();
-
-        let mut payload = T::PAYLOAD::new_buffer();
-        let _ = self
-            .managed_vec
-            .buffer
-            .load_slice(self.byte_end, payload.payload_slice_mut());
-
+        let payload = self.payload_iter.next_back()?;
         unsafe { Some(T::borrow_from_payload(&payload)) }
     }
 }
@@ -102,10 +81,11 @@ where
     T: ManagedVecItem,
 {
     fn clone(&self) -> Self {
-        Self {
-            managed_vec: self.managed_vec,
-            byte_start: self.byte_start,
-            byte_end: self.byte_end,
+        unsafe {
+            ManagedVecRefIterator {
+                payload_iter: self.payload_iter.clone_iter(),
+                _phantom: PhantomData,
+            }
         }
     }
 }