Rust Vec 源码分析

1 Vec 数据结构

pub struct Vec<T, #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator = Global> {
    buf: RawVec<T, A>,
    len: usize,
}

pub(crate) struct RawVec<T, A: Allocator = Global> {
    ptr: Unique<T>,
    cap: usize,
    alloc: A,
}

Vec 内存大小 24 字节，包括一个指向堆上的 ptr，一个 cap，一个 len，每个都是 8 字节大小的 u64。

fn main() {
    let mut v = Vec::new();
    v.push(1);
    // size: 24
    println!("size: {}", std::mem::size_of::<Vec<i32>>());

    let (ptr, cap, len): (u64, u64, u64) = unsafe { std::mem::transmute(v) };
    // ptr: 0x600001194010, cap: 4, len: 1
    println!("ptr: 0x{:0x}, cap: {}, len: {}", ptr, cap, len);
}

2 扩容

以 2 的 n 次幂的方式进行扩容，在第一次扩容的时候会有些特殊 (一个小优化)。

pub fn push(&mut self, value: T) {
 	// This will panic or abort if we would allocate > isize::MAX bytes
    // or if the length increment would overflow for zero-sized types.
    if self.len == self.buf.capacity() {
    	// 在这里面进行扩容
        self.reserve(1);
    }
    unsafe {
        let end = self.as_mut_ptr().add(self.len);
        ptr::write(end, value);
        self.len += 1;
    }
}

pub fn reserve(&mut self, additional: usize) {
	// 在这里面进行扩容
	self.buf.reserve(self.len, additional);
}

pub fn reserve(&mut self, len: usize, additional: usize) {
    // Callers expect this function to be very cheap when there is already sufficient capacity.
    // Therefore, we move all the resizing and error-handling logic from grow_amortized and
    // handle_reserve behind a call, while making sure that this function is likely to be
    // inlined as just a comparison and a call if the comparison fails.
    #[cold]
    fn do_reserve_and_handle<T, A: Allocator>(
        slf: &mut RawVec<T, A>,
        len: usize,
        additional: usize,
    ) {
    	// 在这里面进行扩容
        handle_reserve(slf.grow_amortized(len, additional));
    }

    if self.needs_to_grow(len, additional) {
    	// 在这里面进行扩容，然后跳转到 do_reserve_and_handle
        do_reserve_and_handle(self, len, additional);
    }
}

fn grow_amortized(&mut self, len: usize, additional: usize) -> Result<(), TryReserveError> {
    // This is ensured by the calling contexts.
    debug_assert!(additional > 0);

    if mem::size_of::<T>() == 0 {
        // Since we return a capacity of `usize::MAX` when `elem_size` is
        // 0, getting to here necessarily means the `RawVec` is overfull.
        return Err(CapacityOverflow.into());
    }

    // Nothing we can really do about these checks, sadly.
    let required_cap = len.checked_add(additional).ok_or(CapacityOverflow)?;

    // This guarantees exponential growth. The doubling cannot overflow
    // because `cap <= isize::MAX` and the type of `cap` is `usize`.
    // 扩容后的容量变为原来容量的 2 倍
    let cap = cmp::max(self.cap * 2, required_cap);
    // 也有特殊情况，比如你在第一次 push 元素的时候，如果这个元素的大小小于等于 1024 字节，那么 cap 就等于 4
    let cap = cmp::max(Self::MIN_NON_ZERO_CAP, cap);

	// 开辟了一块新内存
    let new_layout = Layout::array::<T>(cap);

    // `finish_grow` is non-generic over `T`.
    let ptr = finish_grow(new_layout, self.current_memory(), &mut self.alloc)?;
    self.set_ptr(ptr);
    Ok(())
}

const MIN_NON_ZERO_CAP: usize = if mem::size_of::<T>() == 1 {
	8
} else if mem::size_of::<T>() <= 1024 {
    4
} else {
    1
};

fn finish_grow<A>(
    new_layout: Result<Layout, LayoutError>,
    current_memory: Option<(NonNull<u8>, Layout)>,
    alloc: &mut A,
) -> Result<NonNull<[u8]>, TryReserveError>
where
    A: Allocator,
{
    // Check for the error here to minimize the size of `RawVec::grow_*`.
    let new_layout = new_layout.map_err(|_| CapacityOverflow)?;

    alloc_guard(new_layout.size())?;

    let memory = if let Some((ptr, old_layout)) = current_memory {
        debug_assert_eq!(old_layout.align(), new_layout.align());
        unsafe {
            // The allocator checks for alignment equality
            intrinsics::assume(old_layout.align() == new_layout.align());
            // 复制数据到新内存中
            alloc.grow(ptr, old_layout, new_layout)
        }
    } else {
        alloc.allocate(new_layout)
    };

    memory.map_err(|_| AllocError { layout: new_layout, non_exhaustive: () }.into())
}

unsafe fn grow(
    &self,
    ptr: NonNull<u8>,
    old_layout: Layout,
    new_layout: Layout,
) -> Result<NonNull<[u8]>, AllocError> {
    debug_assert!(
        new_layout.size() >= old_layout.size(),
        "`new_layout.size()` must be greater than or equal to `old_layout.size()`"
    );

    let new_ptr = self.allocate(new_layout)?;

    // SAFETY: because `new_layout.size()` must be greater than or equal to
    // `old_layout.size()`, both the old and new memory allocation are valid for reads and
    // writes for `old_layout.size()` bytes. Also, because the old allocation wasn't yet
    // deallocated, it cannot overlap `new_ptr`. Thus, the call to `copy_nonoverlapping` is
    // safe. The safety contract for `dealloc` must be upheld by the caller.
    unsafe {
    	// 复制数据到新内存中
        ptr::copy_nonoverlapping(ptr.as_ptr(), new_ptr.as_mut_ptr(), old_layout.size());
        self.deallocate(ptr, old_layout);
    }

    Ok(new_ptr)
}