taos_query/util/
inline_read.rs

1use tokio::io::*;
2
3use super::AsyncInlinable;
4
5#[async_trait::async_trait]
6pub trait AsyncInlinableRead: AsyncRead + Unpin + Send {
7    #[inline]
8    /// Read `N` bytes as `usize`.
9    ///
10    /// Only 1/2/4/8 is valid as `N`.
11    async fn read_len_with_width<const N: usize>(&mut self) -> std::io::Result<usize> {
12        let mut bytes: [u8; N] = [0; N];
13        self.read_exact(&mut bytes).await?;
14        let len = match N {
15            1 => bytes[0] as usize,
16            2 => unsafe { *std::mem::transmute::<*const u8, *const u16>(bytes.as_ptr()) as usize },
17            4 => unsafe { *std::mem::transmute::<*const u8, *const u32>(bytes.as_ptr()) as usize },
18            8 => unsafe { *std::mem::transmute::<*const u8, *const u64>(bytes.as_ptr()) as usize },
19            _ => unreachable!(),
20        };
21        Ok(len)
22    }
23
24    #[inline]
25    /// Read a bytes `len` with width `N` and the next `len - N` bytes into a `Vec`。
26    ///
27    /// The bytes contains:
28    ///
29    /// ```text
30    /// +--------------+-----------------+
31    /// | len: N bytes | data: len - N bytes |
32    /// +--------------+-----------------+
33    /// ```
34    ///
35    async fn read_len_with_data<const N: usize>(&mut self) -> std::io::Result<Vec<u8>> {
36        let len = self.read_len_with_width::<N>().await?;
37        let mut buf = Vec::with_capacity(len);
38        buf.extend(&(len as u64).to_le_bytes()[0..N]);
39        buf.resize(len, 0);
40        self.read_exact(&mut buf[N..]).await?;
41        Ok(buf)
42    }
43    #[inline]
44    /// Read inlined bytes with specific length width `N`.
45    ///
46    /// The inlined bytes are constructed as:
47    ///
48    /// ```text
49    /// +--------------+-----------------+
50    /// | len: N bytes | data: len bytes |
51    /// +--------------+-----------------+
52    /// ```
53    ///
54    async fn read_inlined_bytes<const N: usize>(&mut self) -> std::io::Result<Vec<u8>> {
55        let len = self.read_len_with_width::<N>().await?;
56        let mut buf = vec![0; len];
57        self.read_exact(&mut buf).await?;
58        Ok(buf)
59    }
60
61    #[inline]
62    /// Read inlined string with specific length width `N`.
63    ///
64    /// The inlined string are constructed as:
65    ///
66    /// ```text
67    /// +--------------+-----------------+
68    /// | len: N bytes | data: len bytes |
69    /// +--------------+-----------------+
70    /// ```
71    ///
72    async fn read_inlined_str<const N: usize>(&mut self) -> std::io::Result<String> {
73        self.read_inlined_bytes::<N>().await.and_then(|vec| {
74            String::from_utf8(vec)
75                .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
76        })
77    }
78
79    #[inline]
80    /// Read some bytes into inlinable object.
81    async fn read_inlinable<T: AsyncInlinable>(&mut self) -> std::io::Result<T>
82    where
83        T: Sized,
84    {
85        self.read_inlinable().await
86    }
87}
88
89impl<T> AsyncInlinableRead for T where T: AsyncRead + Send + Unpin {}
taos_query/util/inline_read.rs

taos_query/util/
inline_read.rs
