use alloc::vec::Vec;
use core::fmt;
use core::ops::{Deref, DerefMut};
pub struct InlineVec<T>(Repr<T>);
#[derive(Clone)]
enum Repr<T> {
Empty,
One(T),
Spilled(Vec<T>),
}
impl<T> InlineVec<T> {
pub const fn new() -> Self {
Self(Repr::Empty)
}
pub const fn one(value: T) -> Self {
Self(Repr::One(value))
}
#[cfg(test)]
pub(crate) fn spilled(&self) -> bool {
matches!(self.0, Repr::Spilled(_))
}
pub fn as_slice(&self) -> &[T] {
match &self.0 {
Repr::Empty => &[],
Repr::One(value) => core::slice::from_ref(value),
Repr::Spilled(values) => values,
}
}
pub fn as_mut_slice(&mut self) -> &mut [T] {
match &mut self.0 {
Repr::Empty => &mut [],
Repr::One(value) => core::slice::from_mut(value),
Repr::Spilled(values) => values,
}
}
pub fn push(&mut self, value: T) {
match &mut self.0 {
Repr::Empty => self.0 = Repr::One(value),
Repr::Spilled(values) => values.push(value),
Repr::One(_) => {
let Repr::One(first) = core::mem::replace(&mut self.0, Repr::Empty) else {
unreachable!("the arm above matched One")
};
self.0 = Repr::Spilled(alloc::vec![first, value]);
}
}
}
pub fn pop(&mut self) -> Option<T> {
match &mut self.0 {
Repr::Empty => None,
Repr::Spilled(values) => values.pop(),
Repr::One(_) => match core::mem::replace(&mut self.0, Repr::Empty) {
Repr::One(value) => Some(value),
_ => unreachable!("the arm above matched One"),
},
}
}
pub fn retain(&mut self, mut keep: impl FnMut(&T) -> bool) {
match &mut self.0 {
Repr::Empty => {}
Repr::Spilled(values) => values.retain(|value| keep(value)),
Repr::One(value) => {
if !keep(value) {
self.0 = Repr::Empty;
}
}
}
}
pub fn clear(&mut self) {
self.0 = Repr::Empty;
}
pub(crate) fn into_vec(self) -> Vec<T> {
match self.0 {
Repr::Empty => Vec::new(),
Repr::One(value) => alloc::vec![value],
Repr::Spilled(values) => values,
}
}
}
impl<T> Default for InlineVec<T> {
fn default() -> Self {
Self::new()
}
}
impl<T: Clone> Clone for InlineVec<T> {
fn clone(&self) -> Self {
Self(self.0.clone())
}
}
impl<T> Deref for InlineVec<T> {
type Target = [T];
fn deref(&self) -> &[T] {
self.as_slice()
}
}
impl<T> DerefMut for InlineVec<T> {
fn deref_mut(&mut self) -> &mut [T] {
self.as_mut_slice()
}
}
impl<T: PartialEq> PartialEq for InlineVec<T> {
fn eq(&self, other: &Self) -> bool {
self.as_slice() == other.as_slice()
}
}
impl<T: Eq> Eq for InlineVec<T> {}
impl<T: PartialEq<U>, U, const N: usize> PartialEq<[U; N]> for InlineVec<T> {
fn eq(&self, other: &[U; N]) -> bool {
self.as_slice() == other.as_slice()
}
}
impl<T: PartialEq<U>, U> PartialEq<[U]> for InlineVec<T> {
fn eq(&self, other: &[U]) -> bool {
self.as_slice() == other
}
}
impl<T: PartialEq<U>, U> PartialEq<Vec<U>> for InlineVec<T> {
fn eq(&self, other: &Vec<U>) -> bool {
self.as_slice() == other.as_slice()
}
}
impl<T: fmt::Debug> fmt::Debug for InlineVec<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.as_slice().fmt(f)
}
}
impl<T> From<Vec<T>> for InlineVec<T> {
fn from(values: Vec<T>) -> Self {
let mut values = values;
match values.len() {
0 => Self::new(),
1 => Self(Repr::One(values.pop().expect("length is one"))),
_ => Self(Repr::Spilled(values)),
}
}
}
impl<T, const N: usize> From<[T; N]> for InlineVec<T> {
fn from(values: [T; N]) -> Self {
values.into_iter().collect()
}
}
impl<T> From<InlineVec<T>> for Vec<T> {
fn from(inline: InlineVec<T>) -> Self {
inline.into_vec()
}
}
impl<T> FromIterator<T> for InlineVec<T> {
fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
let mut iter = iter.into_iter();
let Some(first) = iter.next() else {
return Self::new();
};
let Some(second) = iter.next() else {
return Self(Repr::One(first));
};
let mut values = Vec::with_capacity(iter.size_hint().0 + 2);
values.push(first);
values.push(second);
values.extend(iter);
Self(Repr::Spilled(values))
}
}
impl<T> Extend<T> for InlineVec<T> {
fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
for value in iter {
self.push(value);
}
}
}
pub enum IntoIter<T> {
#[doc(hidden)]
Inline(core::option::IntoIter<T>),
#[doc(hidden)]
Spilled(alloc::vec::IntoIter<T>),
}
impl<T> Iterator for IntoIter<T> {
type Item = T;
fn next(&mut self) -> Option<T> {
match self {
Self::Inline(iter) => iter.next(),
Self::Spilled(iter) => iter.next(),
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
match self {
Self::Inline(iter) => iter.size_hint(),
Self::Spilled(iter) => iter.size_hint(),
}
}
}
impl<T> ExactSizeIterator for IntoIter<T> {}
impl<T> IntoIterator for InlineVec<T> {
type Item = T;
type IntoIter = IntoIter<T>;
fn into_iter(self) -> IntoIter<T> {
match self.0 {
Repr::Empty => IntoIter::Inline(None.into_iter()),
Repr::One(value) => IntoIter::Inline(Some(value).into_iter()),
Repr::Spilled(values) => IntoIter::Spilled(values.into_iter()),
}
}
}
impl<'a, T> IntoIterator for &'a InlineVec<T> {
type Item = &'a T;
type IntoIter = core::slice::Iter<'a, T>;
fn into_iter(self) -> core::slice::Iter<'a, T> {
self.as_slice().iter()
}
}
impl<'a, T> IntoIterator for &'a mut InlineVec<T> {
type Item = &'a mut T;
type IntoIter = core::slice::IterMut<'a, T>;
fn into_iter(self) -> core::slice::IterMut<'a, T> {
self.as_mut_slice().iter_mut()
}
}
#[cfg(test)]
mod tests {
use super::*;
use alloc::vec;
#[test]
fn an_empty_value_holds_nothing_inline() {
let empty = InlineVec::<u32>::new();
assert!(empty.is_empty());
assert_eq!(empty.len(), 0);
assert!(!empty.spilled());
assert_eq!(empty.as_slice(), &[] as &[u32]);
}
#[test]
fn a_single_element_stays_inline() {
let mut one = InlineVec::new();
one.push(7u32);
assert!(!one.spilled(), "one element must not allocate");
assert_eq!(one.as_slice(), [7]);
assert_eq!(InlineVec::one(7u32).as_slice(), [7]);
}
#[test]
fn a_second_element_spills_to_the_heap_in_order() {
let mut many = InlineVec::new();
many.push(1u32);
many.push(2);
many.push(3);
assert!(many.spilled());
assert_eq!(many.as_slice(), [1, 2, 3]);
}
#[test]
fn collecting_picks_the_representation_from_the_length() {
let none: InlineVec<u32> = core::iter::empty().collect();
assert!(none.is_empty() && !none.spilled());
let one: InlineVec<u32> = core::iter::once(5).collect();
assert!(!one.spilled());
assert_eq!(one.as_slice(), [5]);
let many: InlineVec<u32> = (0..4).collect();
assert!(many.spilled());
assert_eq!(many.as_slice(), [0, 1, 2, 3]);
}
#[test]
fn a_short_vec_gives_up_its_allocation() {
assert!(!InlineVec::from(vec![9u32]).spilled());
assert!(InlineVec::from(vec![9u32, 10]).spilled());
assert!(!InlineVec::<u32>::from(vec![]).spilled());
assert_eq!(InlineVec::from(vec![9u32, 10]).as_slice(), [9, 10]);
}
#[test]
fn the_two_representations_compare_equal_on_equal_contents() {
let inline = InlineVec::one(4u32);
let mut spilled = InlineVec::new();
spilled.push(4u32);
spilled.push(5);
spilled.pop();
assert!(spilled.spilled() && !inline.spilled());
assert_eq!(inline, spilled);
assert_eq!(
alloc::format!("{inline:?}"),
alloc::format!("{spilled:?}"),
"the representation must not show in Debug"
);
}
#[test]
fn retain_drops_the_inline_element_and_filters_the_spilled_one() {
let mut one = InlineVec::one(1u32);
one.retain(|&v| v != 1);
assert!(one.is_empty());
let mut many: InlineVec<u32> = (0..6).collect();
many.retain(|v| v % 2 == 0);
assert_eq!(many.as_slice(), [0, 2, 4]);
}
#[test]
fn a_shrunk_value_keeps_its_heap_storage_until_cleared() {
let mut many: InlineVec<u32> = (0..3).collect();
many.retain(|&v| v == 0);
assert!(many.spilled());
assert_eq!(many.as_slice(), [0]);
many.clear();
assert!(!many.spilled());
assert!(many.is_empty());
}
#[test]
fn pop_returns_elements_from_the_end_of_both_representations() {
let mut one = InlineVec::one(1u32);
assert_eq!(one.pop(), Some(1));
assert_eq!(one.pop(), None);
let mut many: InlineVec<u32> = (0..3).collect();
assert_eq!(many.pop(), Some(2));
assert_eq!(many.pop(), Some(1));
assert_eq!(many.pop(), Some(0));
assert_eq!(many.pop(), None);
}
#[test]
fn iteration_covers_every_representation() {
let cases = [
InlineVec::<u32>::new(),
InlineVec::one(1),
(1..4).collect::<InlineVec<u32>>(),
];
let expected: [&[u32]; 3] = [&[], &[1], &[1, 2, 3]];
for (case, want) in cases.into_iter().zip(expected) {
assert_eq!(case.iter().copied().collect::<Vec<_>>(), want);
assert_eq!((&case).into_iter().copied().collect::<Vec<_>>(), want);
assert_eq!(case.clone().into_vec(), want);
assert_eq!(case.into_iter().collect::<Vec<_>>(), want);
}
}
#[test]
fn mutable_access_reaches_both_representations() {
let mut one = InlineVec::one(1u32);
for value in &mut one {
*value += 1;
}
assert_eq!(one.as_slice(), [2]);
let mut many: InlineVec<u32> = (0..3).collect();
many.as_mut_slice()[1] = 9;
assert_eq!(many.as_slice(), [0, 9, 2]);
}
#[test]
fn extend_grows_through_the_same_spill_rule() {
let mut values = InlineVec::new();
values.extend([1u32]);
assert!(!values.spilled());
values.extend([2u32, 3]);
assert_eq!(values.as_slice(), [1, 2, 3]);
}
#[test]
fn the_inline_element_costs_no_more_than_the_vec_it_replaces() {
assert_eq!(
size_of::<InlineVec<u32>>(),
size_of::<Vec<u32>>(),
"the inline case must ride in the Vec's niche"
);
assert_eq!(size_of::<InlineVec<u64>>(), size_of::<Vec<u64>>());
assert_eq!(
size_of::<InlineVec<(u32, u32)>>(),
size_of::<Vec<(u32, u32)>>()
);
}
#[test]
fn arrays_and_vecs_convert_and_compare_across_both_representations() {
let one: InlineVec<u32> = [1].into();
let many: InlineVec<u32> = [1, 2].into();
assert!(!one.spilled() && many.spilled());
assert_eq!(one, [1]);
assert_eq!(many, [1, 2]);
assert_eq!(many, vec![1u32, 2]);
assert_eq!(many, *[1u32, 2].as_slice());
assert_ne!(one, [2]);
assert_ne!(one, [1, 2]);
}
#[test]
fn slice_methods_are_available_through_deref() {
let values: InlineVec<u32> = (10..13).collect();
assert!(values.contains(&11));
assert_eq!(values[2], 12);
assert_eq!(values.first(), Some(&10));
assert_eq!(values.to_vec(), vec![10, 11, 12]);
}
#[test]
fn dropping_runs_element_destructors_once() {
use alloc::rc::Rc;
let witness = Rc::new(());
let mut values = InlineVec::new();
values.push(Rc::clone(&witness));
assert_eq!(Rc::strong_count(&witness), 2);
values.push(Rc::clone(&witness));
assert_eq!(Rc::strong_count(&witness), 3, "the spill must not leak");
drop(values);
assert_eq!(Rc::strong_count(&witness), 1);
}
}