aws_smithy_eventstream/buf/
count.rs

Help
1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6//! A [`Buf`] implementation that counts bytes read.
7
8use bytes::Buf;
9
10/// A [`Buf`] implementation that counts bytes read.
11pub(crate) struct CountBuf<'a, B>
12where
13    B: Buf,
14{
15    buffer: &'a mut B,
16    count: usize,
17}
18
19impl<'a, B> CountBuf<'a, B>
20where
21    B: Buf,
22{
23    /// Creates a new `CountBuf` by wrapping the given `buffer`.
24    pub(crate) fn new(buffer: &'a mut B) -> Self {
25        CountBuf { buffer, count: 0 }
26    }
27
28    /// Consumes the `CountBuf` and returns the number of bytes read.
29    pub(crate) fn into_count(self) -> usize {
30        self.count
31    }
32}
33
34impl<'a, B> Buf for CountBuf<'a, B>
35where
36    B: Buf,
37{
38    fn remaining(&self) -> usize {
39        self.buffer.remaining()
40    }
41
42    fn chunk(&self) -> &[u8] {
43        self.buffer.chunk()
44    }
45
46    fn advance(&mut self, cnt: usize) {
47        self.count += cnt;
48        self.buffer.advance(cnt);
49    }
50}
51
52#[cfg(test)]
53mod tests {
54    use super::CountBuf;
55    use bytes::Buf;
56
57    #[test]
58    fn count_no_data_read() {
59        let mut data: &[u8] = &[];
60        let buf = CountBuf::new(&mut data);
61        assert_eq!(0, buf.into_count());
62    }
63
64    #[test]
65    fn count_data_read() {
66        let mut data: &[u8] = &[0, 0, 0, 5, 0, 10u8];
67        let mut buf = CountBuf::new(&mut data);
68        assert_eq!(5, buf.get_i32());
69        assert_eq!(4, buf.into_count());
70
71        let mut data: &[u8] = &[0, 0, 0, 5, 0, 10u8];
72        let mut buf = CountBuf::new(&mut data);
73        assert_eq!(5, buf.get_i32());
74        assert_eq!(10, buf.get_i16());
75        assert_eq!(6, buf.into_count());
76    }
77
78    #[test]
79    fn chunk_called_multiple_times_before_advance() {
80        let mut data: &[u8] = &[0, 0, 0, 5, 0, 10u8];
81        let mut buf = CountBuf::new(&mut data);
82        for _ in 0..3 {
83            buf.chunk();
84        }
85        buf.advance(4);
86        assert_eq!(10, buf.get_i16());
87        assert_eq!(6, buf.into_count());
88    }
89}
aws_smithy_eventstream/buf/count.rs

aws_smithy_eventstream/buf/
count.rs
