arrow_buffer/
bytes.rs

Help
1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! This module contains an implementation of a contiguous immutable memory region that knows
19//! how to de-allocate itself, [`Bytes`].
20//! Note that this is a low-level functionality of this crate.
21
22use core::slice;
23use std::ptr::NonNull;
24use std::{fmt::Debug, fmt::Formatter};
25
26use crate::alloc::Deallocation;
27
28/// A continuous, fixed-size, immutable memory region that knows how to de-allocate itself.
29///
30/// This structs' API is inspired by the `bytes::Bytes`, but it is not limited to using rust's
31/// global allocator nor u8 alignment.
32///
33/// In the most common case, this buffer is allocated using [`alloc`](std::alloc::alloc)
34/// with an alignment of [`ALIGNMENT`](crate::alloc::ALIGNMENT)
35///
36/// When the region is allocated by a different allocator, [Deallocation::Custom], this calls the
37/// custom deallocator to deallocate the region when it is no longer needed.
38pub struct Bytes {
39    /// The raw pointer to be beginning of the region
40    ptr: NonNull<u8>,
41
42    /// The number of bytes visible to this region. This is always smaller than its capacity (when available).
43    len: usize,
44
45    /// how to deallocate this region
46    deallocation: Deallocation,
47}
48
49impl Bytes {
50    /// Takes ownership of an allocated memory region,
51    ///
52    /// # Arguments
53    ///
54    /// * `ptr` - Pointer to raw parts
55    /// * `len` - Length of raw parts in **bytes**
56    /// * `deallocation` - Type of allocation
57    ///
58    /// # Safety
59    ///
60    /// This function is unsafe as there is no guarantee that the given pointer is valid for `len`
61    /// bytes. If the `ptr` and `capacity` come from a `Buffer`, then this is guaranteed.
62    #[inline]
63    pub(crate) unsafe fn new(ptr: NonNull<u8>, len: usize, deallocation: Deallocation) -> Bytes {
64        Bytes {
65            ptr,
66            len,
67            deallocation,
68        }
69    }
70
71    fn as_slice(&self) -> &[u8] {
72        self
73    }
74
75    #[inline]
76    pub fn len(&self) -> usize {
77        self.len
78    }
79
80    #[inline]
81    pub fn is_empty(&self) -> bool {
82        self.len == 0
83    }
84
85    #[inline]
86    pub fn ptr(&self) -> NonNull<u8> {
87        self.ptr
88    }
89
90    pub fn capacity(&self) -> usize {
91        match self.deallocation {
92            Deallocation::Standard(layout) => layout.size(),
93            // we only know the size of the custom allocation
94            // its underlying capacity might be larger
95            Deallocation::Custom(_, size) => size,
96        }
97    }
98
99    #[inline]
100    pub(crate) fn deallocation(&self) -> &Deallocation {
101        &self.deallocation
102    }
103}
104
105// Deallocation is Send + Sync, repeating the bound here makes that refactoring safe
106// The only field that is not automatically Send+Sync then is the NonNull ptr
107unsafe impl Send for Bytes where Deallocation: Send {}
108unsafe impl Sync for Bytes where Deallocation: Sync {}
109
110impl Drop for Bytes {
111    #[inline]
112    fn drop(&mut self) {
113        match &self.deallocation {
114            Deallocation::Standard(layout) => match layout.size() {
115                0 => {} // Nothing to do
116                _ => unsafe { std::alloc::dealloc(self.ptr.as_ptr(), *layout) },
117            },
118            // The automatic drop implementation will free the memory once the reference count reaches zero
119            Deallocation::Custom(_allocation, _size) => (),
120        }
121    }
122}
123
124impl std::ops::Deref for Bytes {
125    type Target = [u8];
126
127    fn deref(&self) -> &[u8] {
128        unsafe { slice::from_raw_parts(self.ptr.as_ptr(), self.len) }
129    }
130}
131
132impl PartialEq for Bytes {
133    fn eq(&self, other: &Bytes) -> bool {
134        self.as_slice() == other.as_slice()
135    }
136}
137
138impl Debug for Bytes {
139    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
140        write!(f, "Bytes {{ ptr: {:?}, len: {}, data: ", self.ptr, self.len,)?;
141
142        f.debug_list().entries(self.iter()).finish()?;
143
144        write!(f, " }}")
145    }
146}
147
148impl From<bytes::Bytes> for Bytes {
149    fn from(value: bytes::Bytes) -> Self {
150        let len = value.len();
151        Self {
152            len,
153            ptr: NonNull::new(value.as_ptr() as _).unwrap(),
154            deallocation: Deallocation::Custom(std::sync::Arc::new(value), len),
155        }
156    }
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162
163    #[test]
164    fn test_from_bytes() {
165        let bytes = bytes::Bytes::from(vec![1, 2, 3, 4]);
166        let arrow_bytes: Bytes = bytes.clone().into();
167
168        assert_eq!(bytes.as_ptr(), arrow_bytes.as_ptr());
169
170        drop(bytes);
171        drop(arrow_bytes);
172
173        let _ = Bytes::from(bytes::Bytes::new());
174    }
175}
arrow_buffer/bytes.rs

arrow_buffer/
bytes.rs
