This is an automated email from the ASF dual-hosted git repository. xuanwo pushed a commit to branch exact-buf-write in repository https://gitbox.apache.org/repos/asf/incubator-opendal.git
commit cd8663503a5a6bd7c9d9763a03ca3e60e0572840 Author: Xuanwo <[email protected]> AuthorDate: Wed Aug 23 16:27:19 2023 +0800 Add split support Signed-off-by: Xuanwo <[email protected]> --- core/src/raw/oio/cursor.rs | 162 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 162 insertions(+) diff --git a/core/src/raw/oio/cursor.rs b/core/src/raw/oio/cursor.rs index ad1ac4661..ec4589d18 100644 --- a/core/src/raw/oio/cursor.rs +++ b/core/src/raw/oio/cursor.rs @@ -220,6 +220,86 @@ impl ChunkedCursor { pub fn push(&mut self, bs: Bytes) { self.inner.push_back(bs); } + + /// split_off will split the cursor into two cursors at given size. + /// + /// After split, `self` will contains the `0..at` part and the returned cursor contains + /// `at..` parts. + /// + /// # Panics + /// + /// - Panics if `at > len` + /// - Panics if `idx != 0`, the cursor must be reset before split. + pub fn split_off(&mut self, mut at: usize) -> Self { + assert!( + at <= self.len(), + "split_off at must smaller than current size" + ); + assert_eq!(self.idx, 0, "split_off must reset cursor first"); + + let mut chunks = VecDeque::new(); + let mut size = self.len() - at; + + while let Some(mut bs) = self.inner.pop_back() { + if size > bs.len() { + size -= bs.len(); + chunks.push_front(bs); + } else if size == bs.len() { + chunks.push_front(bs); + break; + } else { + let remaining = bs.split_off(bs.len() - size); + chunks.push_front(remaining); + self.inner.push_back(bs); + break; + } + } + + Self { + inner: chunks, + idx: 0, + } + } + + /// split_to will split the cursor into two cursors at given size. + /// + /// After split, `self` will contains the `at..` part and the returned cursor contains + /// `0..at` parts. + /// + /// # Panics + /// + /// - Panics if `at > len` + /// - Panics if `idx != 0`, the cursor must be reset before split. + pub fn split_to(&mut self, at: usize) -> Self { + assert!( + at <= self.len(), + "split_to at must smaller than current size" + ); + assert_eq!(self.idx, 0, "split_to must reset cursor first"); + + let mut chunks = VecDeque::new(); + let mut size = at; + + while let Some(mut bs) = self.inner.pop_front() { + if size > bs.len() { + size -= bs.len(); + chunks.push_back(bs); + } else if size == bs.len() { + chunks.push_back(bs); + break; + } else { + let remaining = bs.split_off(size); + chunks.push_back(bs); + self.inner.push_front(remaining); + break; + } + } + + Self { + inner: chunks, + idx: 0, + } + } } impl oio::Stream for ChunkedCursor { @@ -457,4 +537,86 @@ mod tests { Ok(()) } + + #[test] + fn test_chunked_cursor_split_to() { + let mut base = ChunkedCursor::new(); + base.push(Bytes::from("Hello")); + base.push(Bytes::from("Wor")); + base.push(Bytes::from("ld")); + + // Case 1: split less than first chunk + let mut c1 = base.clone(); + let c2 = c1.split_to(3); + + assert_eq!(c1.len(), 7); + assert_eq!( + &c1.inner, + &[Bytes::from("lo"), Bytes::from("Wor"), Bytes::from("ld")] + ); + + assert_eq!(c2.len(), 3); + assert_eq!(&c2.inner, &[Bytes::from("Hel")]); + + // Case 2: split larger than first chunk + let mut c1 = base.clone(); + let c2 = c1.split_to(6); + + assert_eq!(c1.len(), 4); + assert_eq!(&c1.inner, &[Bytes::from("or"), Bytes::from("ld")]); + + assert_eq!(c2.len(), 6); + assert_eq!(&c2.inner, &[Bytes::from("Hello"), Bytes::from("W")]); + + // Case 3: split at chunk edge + let mut c1 = base.clone(); + let c2 = c1.split_to(8); + + assert_eq!(c1.len(), 2); + assert_eq!(&c1.inner, &[Bytes::from("ld")]); + + assert_eq!(c2.len(), 8); + assert_eq!(&c2.inner, &[Bytes::from("Hello"), Bytes::from("Wor")]); + } + + #[test] + fn test_chunked_cursor_split_off() { + let mut base = ChunkedCursor::new(); + base.push(Bytes::from("Hello")); + base.push(Bytes::from("Wor")); + base.push(Bytes::from("ld")); + + // Case 1: split less than first chunk + let mut c1 = base.clone(); + let c2 = c1.split_off(3); + + assert_eq!(c1.len(), 3); + assert_eq!(&c1.inner, &[Bytes::from("Hel")]); + + assert_eq!(c2.len(), 7); + assert_eq!( + &c2.inner, + &[Bytes::from("lo"), Bytes::from("Wor"), Bytes::from("ld")] + ); + + // Case 2: split larger than first chunk + let mut c1 = base.clone(); + let c2 = c1.split_off(6); + + assert_eq!(c1.len(), 6); + assert_eq!(&c1.inner, &[Bytes::from("Hello"), Bytes::from("W")]); + + assert_eq!(c2.len(), 4); + assert_eq!(&c2.inner, &[Bytes::from("or"), Bytes::from("ld")]); + + // Case 3: split at chunk edge + let mut c1 = base.clone(); + let c2 = c1.split_off(8); + + assert_eq!(c1.len(), 8); + assert_eq!(&c1.inner, &[Bytes::from("Hello"), Bytes::from("Wor")]); + + assert_eq!(c2.len(), 2); + assert_eq!(&c2.inner, &[Bytes::from("ld")]); + } }
