/* This Source Code Form is subject to the terms of the Mozilla Public * License, v. 2.0. If a copy of the MPL was not distributed with this * file, You can obtain one at https://mozilla.org/MPL/2.0/. * * (c) ZeroTier, Inc. * https://www.zerotier.com/ */ use std::fmt::Debug; use std::io::Write; use std::mem::{size_of, MaybeUninit}; use std::ptr::{slice_from_raw_parts, slice_from_raw_parts_mut}; use serde::ser::SerializeSeq; use serde::{Deserialize, Deserializer, Serialize, Serializer}; #[derive(Clone, Copy, Debug)] pub struct OutOfCapacityError(pub T); impl std::fmt::Display for OutOfCapacityError { fn fmt(self: &Self, stream: &mut std::fmt::Formatter) -> std::fmt::Result { std::fmt::Display::fmt("ArrayVec out of space", stream) } } impl ::std::error::Error for OutOfCapacityError { fn description(self: &Self) -> &str { "ArrayVec out of space" } } /// A simple vector backed by a static sized array with no memory allocations and no overhead construction. pub struct ArrayVec { pub(crate) s: usize, pub(crate) a: [MaybeUninit; C], } impl Default for ArrayVec { #[inline(always)] fn default() -> Self { Self::new() } } impl PartialEq for ArrayVec { #[inline(always)] fn eq(&self, other: &Self) -> bool { let tmp: &[T] = self.as_ref(); tmp.eq(other.as_ref()) } } impl Eq for ArrayVec {} impl Clone for ArrayVec { #[inline] fn clone(&self) -> Self { debug_assert!(self.s <= C); Self { s: self.s, a: unsafe { let mut tmp: [MaybeUninit; C] = MaybeUninit::uninit().assume_init(); for i in 0..self.s { tmp.get_unchecked_mut(i).write(self.a[i].assume_init_ref().clone()); } tmp }, } } } impl From<[T; S]> for ArrayVec { #[inline] fn from(v: [T; S]) -> Self { if S <= C { let mut tmp = Self::new(); for i in 0..S { tmp.push(v[i].clone()); } tmp } else { panic!(); } } } impl ToString for ArrayVec { #[inline] fn to_string(&self) -> String { crate::hex::to_string(self.as_bytes()) } } impl Debug for ArrayVec { #[inline] fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.write_str(self.to_string().as_str()) } } impl Write for ArrayVec { #[inline] fn write(&mut self, buf: &[u8]) -> std::io::Result { for i in buf.iter() { if self.try_push(*i).is_err() { return Err(std::io::Error::new(std::io::ErrorKind::Other, "ArrayVec out of space")); } } Ok(buf.len()) } #[inline(always)] fn flush(&mut self) -> std::io::Result<()> { Ok(()) } } impl TryFrom> for ArrayVec { type Error = OutOfCapacityError; #[inline(always)] fn try_from(mut value: Vec) -> Result { let mut tmp = Self::new(); for x in value.drain(..) { tmp.try_push(x)?; } Ok(tmp) } } impl TryFrom<&Vec> for ArrayVec { type Error = OutOfCapacityError; #[inline(always)] fn try_from(value: &Vec) -> Result { let mut tmp = Self::new(); for x in value.iter() { tmp.try_push(x.clone())?; } Ok(tmp) } } impl TryFrom<&[T]> for ArrayVec { type Error = OutOfCapacityError; #[inline(always)] fn try_from(value: &[T]) -> Result { let mut tmp = Self::new(); for x in value.iter() { tmp.try_push(x.clone())?; } Ok(tmp) } } impl ArrayVec { #[inline(always)] pub fn new() -> Self { assert_eq!(size_of::<[T; C]>(), size_of::<[MaybeUninit; C]>()); Self { s: 0, a: unsafe { MaybeUninit::uninit().assume_init() } } } #[inline] pub fn push(&mut self, v: T) { let i = self.s; if i < C { unsafe { self.a.get_unchecked_mut(i).write(v) }; self.s = i + 1; } else { panic!(); } } #[inline] pub fn try_push(&mut self, v: T) -> Result<(), OutOfCapacityError> { if self.s < C { let i = self.s; unsafe { self.a.get_unchecked_mut(i).write(v) }; self.s = i + 1; Ok(()) } else { Err(OutOfCapacityError(v)) } } #[inline(always)] pub fn as_bytes(&self) -> &[T] { unsafe { &*slice_from_raw_parts(self.a.as_ptr().cast(), self.s) } } #[inline(always)] pub fn is_empty(&self) -> bool { self.s == 0 } #[inline(always)] pub fn len(&self) -> usize { self.s } #[inline(always)] pub fn capacity_remaining(&self) -> usize { C - self.s } #[inline] pub fn pop(&mut self) -> Option { if self.s > 0 { let i = self.s - 1; debug_assert!(i < C); self.s = i; Some(unsafe { self.a.get_unchecked(i).assume_init_read() }) } else { None } } } impl ArrayVec where T: Copy, { /// Push a slice of copyable objects, panic if capacity exceeded. #[inline] pub fn push_slice(&mut self, v: &[T]) { let start = self.s; let end = self.s + v.len(); if end <= C { for i in start..end { unsafe { self.a.get_unchecked_mut(i).write(*v.get_unchecked(i - start)) }; } self.s = end; } else { panic!(); } } } impl Drop for ArrayVec { #[inline(always)] fn drop(&mut self) { for i in 0..self.s { unsafe { self.a.get_unchecked_mut(i).assume_init_drop() }; } } } impl AsRef<[T]> for ArrayVec { #[inline(always)] fn as_ref(&self) -> &[T] { unsafe { &*slice_from_raw_parts(self.a.as_ptr().cast(), self.s) } } } impl AsMut<[T]> for ArrayVec { #[inline(always)] fn as_mut(&mut self) -> &mut [T] { unsafe { &mut *slice_from_raw_parts_mut(self.a.as_mut_ptr().cast(), self.s) } } } impl Serialize for ArrayVec { #[inline] fn serialize(&self, serializer: S) -> Result where S: Serializer, { let mut seq = serializer.serialize_seq(Some(self.len()))?; let sl: &[T] = self.as_ref(); for i in 0..self.s { seq.serialize_element(&sl[i])?; } seq.end() } } struct ArrayVecVisitor<'de, T: Deserialize<'de>, const L: usize>(std::marker::PhantomData<&'de T>); impl<'de, T: Deserialize<'de>, const L: usize> serde::de::Visitor<'de> for ArrayVecVisitor<'de, T, L> { type Value = ArrayVec; #[inline] fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { formatter.write_str(format!("array of up to {} elements", L).as_str()) } #[inline] fn visit_seq(self, mut seq: A) -> Result where A: serde::de::SeqAccess<'de>, { let mut a = ArrayVec::::new(); while let Some(x) = seq.next_element()? { a.push(x); } Ok(a) } } impl<'de, T: Deserialize<'de> + 'de, const L: usize> Deserialize<'de> for ArrayVec { #[inline] fn deserialize(deserializer: D) -> Result, D::Error> where D: Deserializer<'de>, { deserializer.deserialize_seq(ArrayVecVisitor(std::marker::PhantomData::default())) } } #[cfg(test)] mod tests { use super::ArrayVec; #[test] fn array_vec() { let mut v = ArrayVec::::new(); for i in 0..128 { v.push(i); } assert_eq!(v.len(), 128); assert!(!v.try_push(1000).is_ok()); assert_eq!(v.len(), 128); for _ in 0..128 { assert!(v.pop().is_some()); } assert!(v.pop().is_none()); } }