Skip to main content

servers/
repeated_field.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15// The Clear trait is copied from https://github.com/stepancheg/rust-protobuf/blob/v2.28.0/protobuf/src/clear.rs
16// The RepeatedField struct is copied from https://github.com/stepancheg/rust-protobuf/blob/v2.28.0/protobuf/src/repeated.rs
17// This code is to leverage the pooling mechanism to avoid frequent heap allocation/de-allocation when decoding deeply nested structs.
18
19use std::borrow::Borrow;
20use std::cmp::Ordering;
21use std::default::Default;
22use std::hash::{Hash, Hasher};
23use std::iter::{FromIterator, IntoIterator};
24use std::ops::{Deref, DerefMut, Index, IndexMut};
25use std::{fmt, slice, vec};
26
27use bytes::Bytes;
28
29const NULL_BYTES: &[u8] = &[];
30
31/// anything that can be cleared
32pub trait Clear {
33    /// Clear this make, make it equivalent to newly created object.
34    fn clear(&mut self);
35}
36
37impl Clear for &[u8] {
38    fn clear(&mut self) {
39        *self = NULL_BYTES;
40    }
41}
42
43impl<T> Clear for Option<T> {
44    fn clear(&mut self) {
45        self.take();
46    }
47}
48
49impl Clear for String {
50    fn clear(&mut self) {
51        String::clear(self);
52    }
53}
54
55impl<T> Clear for Vec<T> {
56    fn clear(&mut self) {
57        Vec::clear(self);
58    }
59}
60
61impl Clear for Bytes {
62    fn clear(&mut self) {
63        Bytes::clear(self);
64    }
65}
66
67/// Wrapper around vector to avoid deallocations on clear.
68pub struct RepeatedField<T> {
69    vec: Vec<T>,
70    len: usize,
71}
72
73impl<T> RepeatedField<T> {
74    /// Return number of elements in this container.
75    #[inline]
76    pub fn len(&self) -> usize {
77        self.len
78    }
79
80    /// Returns true if this container is empty.
81    #[inline]
82    pub fn is_empty(&self) -> bool {
83        self.len == 0
84    }
85
86    /// Clear.
87    #[inline]
88    pub fn clear(&mut self) {
89        self.len = 0;
90    }
91}
92
93impl<T> Default for RepeatedField<T> {
94    #[inline]
95    fn default() -> RepeatedField<T> {
96        RepeatedField {
97            vec: Vec::new(),
98            len: 0,
99        }
100    }
101}
102
103impl<T> RepeatedField<T> {
104    /// Create new empty container.
105    #[inline]
106    pub fn new() -> RepeatedField<T> {
107        Default::default()
108    }
109
110    /// Create a contained with data from given vec.
111    #[inline]
112    pub fn from_vec(vec: Vec<T>) -> RepeatedField<T> {
113        let len = vec.len();
114        RepeatedField { vec, len }
115    }
116
117    /// Convert data into vec.
118    #[inline]
119    pub fn into_vec(self) -> Vec<T> {
120        let mut vec = self.vec;
121        vec.truncate(self.len);
122        vec
123    }
124
125    /// Return current capacity.
126    #[inline]
127    pub fn capacity(&self) -> usize {
128        self.vec.capacity()
129    }
130
131    /// View data as slice.
132    #[inline]
133    pub fn as_slice(&self) -> &[T] {
134        &self.vec[..self.len]
135    }
136
137    /// View data as mutable slice.
138    #[inline]
139    pub fn as_mut_slice(&mut self) -> &mut [T] {
140        &mut self.vec[..self.len]
141    }
142
143    /// Get subslice of this container.
144    #[inline]
145    pub fn slice(&self, start: usize, end: usize) -> &[T] {
146        &self.as_ref()[start..end]
147    }
148
149    /// Get slice from given index.
150    #[inline]
151    pub fn slice_from(&self, start: usize) -> &[T] {
152        &self.as_ref()[start..]
153    }
154
155    /// Get slice to given index.
156    #[inline]
157    pub fn slice_to(&self, end: usize) -> &[T] {
158        &self.as_ref()[..end]
159    }
160
161    /// View this container as two slices split at given index.
162    #[inline]
163    pub fn split_at(&self, mid: usize) -> (&[T], &[T]) {
164        self.as_ref().split_at(mid)
165    }
166
167    /// View this container as two mutable slices split at given index.
168    #[inline]
169    pub fn split_at_mut(&mut self, mid: usize) -> (&mut [T], &mut [T]) {
170        self.as_mut_slice().split_at_mut(mid)
171    }
172
173    /// View all but first elements of this container.
174    #[inline]
175    pub fn tail(&self) -> &[T] {
176        &self.as_ref()[1..]
177    }
178
179    /// Last element of this container.
180    #[inline]
181    pub fn last(&self) -> Option<&T> {
182        self.as_ref().last()
183    }
184
185    /// Mutable last element of this container.
186    #[inline]
187    pub fn last_mut(&mut self) -> Option<&mut T> {
188        self.as_mut_slice().last_mut()
189    }
190
191    /// View all but last elements of this container.
192    #[inline]
193    pub fn init(&self) -> &[T] {
194        let s = self.as_ref();
195        &s[0..s.len() - 1]
196    }
197
198    /// Push an element to the end.
199    #[inline]
200    pub fn push(&mut self, value: T) {
201        if self.len == self.vec.len() {
202            self.vec.push(value);
203        } else {
204            self.vec[self.len] = value;
205        }
206        self.len += 1;
207    }
208
209    /// Pop last element.
210    #[inline]
211    pub fn pop(&mut self) -> Option<T> {
212        if self.len == 0 {
213            None
214        } else {
215            self.vec.truncate(self.len);
216            self.len -= 1;
217            self.vec.pop()
218        }
219    }
220
221    /// Insert an element at specified position.
222    #[inline]
223    pub fn insert(&mut self, index: usize, value: T) {
224        assert!(index <= self.len);
225        self.vec.insert(index, value);
226        self.len += 1;
227    }
228
229    /// Remove an element from specified position.
230    #[inline]
231    pub fn remove(&mut self, index: usize) -> T {
232        assert!(index < self.len);
233        self.len -= 1;
234        self.vec.remove(index)
235    }
236
237    /// Retains only the elements specified by the predicate.
238    ///
239    /// In other words, remove all elements `e` such that `f(&e)` returns `false`.
240    /// This method operates in place, visiting each element exactly once in the
241    /// original order, and preserves the order of the retained elements.
242    ///
243    /// # Examples
244    ///
245    /// ```
246    /// use servers::repeated_field::RepeatedField;
247    ///
248    /// let mut vec = RepeatedField::from(vec![1, 2, 3, 4]);
249    /// vec.retain(|&x| x % 2 == 0);
250    /// assert_eq!(vec, RepeatedField::from(vec![2, 4]));
251    /// ```
252    pub fn retain<F>(&mut self, f: F)
253    where
254        F: FnMut(&T) -> bool,
255    {
256        // suboptimal
257        self.vec.truncate(self.len);
258        self.vec.retain(f);
259        self.len = self.vec.len();
260    }
261
262    /// Truncate at specified length.
263    #[inline]
264    pub fn truncate(&mut self, len: usize) {
265        if self.len > len {
266            self.len = len;
267        }
268    }
269
270    /// Reverse in place.
271    #[inline]
272    pub fn reverse(&mut self) {
273        self.as_mut_slice().reverse()
274    }
275
276    /// Immutable data iterator.
277    #[inline]
278    pub fn iter(&self) -> slice::Iter<'_, T> {
279        self.as_ref().iter()
280    }
281
282    /// Mutable data iterator.
283    #[inline]
284    pub fn iter_mut(&mut self) -> slice::IterMut<'_, T> {
285        self.as_mut_slice().iter_mut()
286    }
287
288    /// Sort elements with given comparator.
289    #[inline]
290    pub fn sort_by<F>(&mut self, compare: F)
291    where
292        F: Fn(&T, &T) -> Ordering,
293    {
294        self.as_mut_slice().sort_by(compare)
295    }
296
297    /// Get data as raw pointer.
298    #[inline]
299    pub fn as_ptr(&self) -> *const T {
300        self.vec.as_ptr()
301    }
302
303    /// Get data a mutable raw pointer.
304    #[inline]
305    pub fn as_mut_ptr(&mut self) -> *mut T {
306        self.vec.as_mut_ptr()
307    }
308}
309
310impl<T: Default + Clear> RepeatedField<T> {
311    /// Push default value.
312    /// This operation could be faster than `rf.push(Default::default())`,
313    /// because it may reuse previously allocated and cleared element.
314    pub fn push_default(&mut self) -> &mut T {
315        if self.len == self.vec.len() {
316            self.vec.push(Default::default());
317        } else {
318            self.vec[self.len].clear();
319        }
320        self.len += 1;
321        self.last_mut().unwrap()
322    }
323}
324
325impl<T> From<Vec<T>> for RepeatedField<T> {
326    #[inline]
327    fn from(values: Vec<T>) -> RepeatedField<T> {
328        RepeatedField::from_vec(values)
329    }
330}
331
332impl<'a, T: Clone> From<&'a [T]> for RepeatedField<T> {
333    #[inline]
334    fn from(values: &'a [T]) -> RepeatedField<T> {
335        RepeatedField::from_slice(values)
336    }
337}
338
339impl<T> From<RepeatedField<T>> for Vec<T> {
340    #[inline]
341    fn from(val: RepeatedField<T>) -> Self {
342        val.into_vec()
343    }
344}
345
346impl<T: Clone> RepeatedField<T> {
347    /// Copy slice data to `RepeatedField`
348    #[inline]
349    pub fn from_slice(values: &[T]) -> RepeatedField<T> {
350        RepeatedField::from_vec(values.to_vec())
351    }
352
353    /// Copy slice data to `RepeatedField`
354    #[inline]
355    pub fn from_ref<X: AsRef<[T]>>(values: X) -> RepeatedField<T> {
356        RepeatedField::from_slice(values.as_ref())
357    }
358
359    /// Copy this data into new vec.
360    #[inline]
361    pub fn to_vec(&self) -> Vec<T> {
362        self.as_ref().to_vec()
363    }
364}
365
366impl<T: Clone> Clone for RepeatedField<T> {
367    #[inline]
368    fn clone(&self) -> RepeatedField<T> {
369        RepeatedField {
370            vec: self.to_vec(),
371            len: self.len(),
372        }
373    }
374}
375
376impl<T> FromIterator<T> for RepeatedField<T> {
377    #[inline]
378    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> RepeatedField<T> {
379        RepeatedField::from_vec(FromIterator::from_iter(iter))
380    }
381}
382
383impl<'a, T> IntoIterator for &'a RepeatedField<T> {
384    type Item = &'a T;
385    type IntoIter = slice::Iter<'a, T>;
386
387    fn into_iter(self) -> slice::Iter<'a, T> {
388        self.iter()
389    }
390}
391
392impl<'a, T> IntoIterator for &'a mut RepeatedField<T> {
393    type Item = &'a mut T;
394    type IntoIter = slice::IterMut<'a, T>;
395
396    fn into_iter(self) -> slice::IterMut<'a, T> {
397        self.iter_mut()
398    }
399}
400
401impl<T> IntoIterator for RepeatedField<T> {
402    type Item = T;
403    type IntoIter = vec::IntoIter<T>;
404
405    fn into_iter(mut self) -> vec::IntoIter<T> {
406        self.vec.truncate(self.len);
407        self.vec.into_iter()
408    }
409}
410
411impl<T: PartialEq> PartialEq for RepeatedField<T> {
412    #[inline]
413    fn eq(&self, other: &RepeatedField<T>) -> bool {
414        self.as_ref() == other.as_ref()
415    }
416}
417
418impl<T: Eq> Eq for RepeatedField<T> {}
419
420impl<T: PartialEq> PartialEq<[T]> for RepeatedField<T> {
421    fn eq(&self, other: &[T]) -> bool {
422        self.as_slice() == other
423    }
424}
425
426impl<T: PartialEq> PartialEq<RepeatedField<T>> for [T] {
427    fn eq(&self, other: &RepeatedField<T>) -> bool {
428        self == other.as_slice()
429    }
430}
431
432impl<T: PartialEq> RepeatedField<T> {
433    /// True iff this container contains given element.
434    #[inline]
435    pub fn contains(&self, value: &T) -> bool {
436        self.as_ref().contains(value)
437    }
438}
439
440impl<T: Hash> Hash for RepeatedField<T> {
441    fn hash<H: Hasher>(&self, state: &mut H) {
442        self.as_ref().hash(state);
443    }
444}
445
446impl<T> AsRef<[T]> for RepeatedField<T> {
447    #[inline]
448    fn as_ref(&self) -> &[T] {
449        &self.vec[..self.len]
450    }
451}
452
453impl<T> Borrow<[T]> for RepeatedField<T> {
454    #[inline]
455    fn borrow(&self) -> &[T] {
456        &self.vec[..self.len]
457    }
458}
459
460impl<T> Deref for RepeatedField<T> {
461    type Target = [T];
462    #[inline]
463    fn deref(&self) -> &[T] {
464        &self.vec[..self.len]
465    }
466}
467
468impl<T> DerefMut for RepeatedField<T> {
469    #[inline]
470    fn deref_mut(&mut self) -> &mut [T] {
471        &mut self.vec[..self.len]
472    }
473}
474
475impl<T> Index<usize> for RepeatedField<T> {
476    type Output = T;
477
478    #[inline]
479    fn index(&self, index: usize) -> &T {
480        &self.as_ref()[index]
481    }
482}
483
484impl<T> IndexMut<usize> for RepeatedField<T> {
485    #[inline]
486    fn index_mut(&mut self, index: usize) -> &mut T {
487        &mut self.as_mut_slice()[index]
488    }
489}
490
491impl<T> Extend<T> for RepeatedField<T> {
492    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
493        self.vec.truncate(self.len);
494        self.vec.extend(iter);
495        self.len = self.vec.len();
496    }
497}
498
499impl<'a, T: Copy + 'a> Extend<&'a T> for RepeatedField<T> {
500    fn extend<I: IntoIterator<Item = &'a T>>(&mut self, iter: I) {
501        self.vec.truncate(self.len);
502        self.vec.extend(iter);
503        self.len = self.vec.len();
504    }
505}
506
507impl<T: fmt::Debug> fmt::Debug for RepeatedField<T> {
508    #[inline]
509    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
510        self.as_ref().fmt(f)
511    }
512}
513
514#[cfg(test)]
515mod tests {
516    use crate::repeated_field::RepeatedField;
517
518    #[test]
519    fn test_null_ptr() {
520        let mut vec: RepeatedField<&'static [u8]> = RepeatedField::new();
521        let borrowed_value = vec.push_default();
522        *borrowed_value = b"hello";
523        vec.clear();
524        let new_value = vec.push_default();
525        assert!((*new_value).is_empty());
526    }
527}