通过自定义Vec＜T＞理解Rust中的内存分配方式

Rust中一般通过syscall分配内存， 可以自定义分配内存算法，实现GlobalAlloc trait, 并通过#[global_allocator]注解告诉编译器自定义的内存分配算法。

alloc和Layout实战

在我们深入了解之前，先通过以下简单代码段对 alloc, dealloc 以及Layout对Rust内存分配有一个基本概念。

use std::alloc::{alloc, dealloc, Layout};
fn main(){
	let layout = Layout::new::<i32>();
    let ptr = unsafe { alloc(layout) } as *mut i32;

    unsafe {
        *ptr = 42;
    }

    let value = unsafe { *ptr };
    println!("value: {}", value);

    unsafe {
        dealloc(ptr as *mut u8, layout);
    }
}

下面我们对这段代码做一个解释， 看看每行代码都做了那些事情。

1. let layout = Layout::new::<i32>();
   这行代码就像是我们预先请求了一段新的数据存储单元， 通过Layout告诉Rust我们需要一段足够大的内存来存储i32， 你也可以换成其他数据类型比如：String，Structs等等，Rust会恰到好处的处理所需的存储单元大小。
2. let ptr = unsafe { alloc(layout) } as *mut i32 ;
   在这里我们通过调用std::alloc::alloc这个函数, 传入我们请求的内存布局, Rust会分配一块内存单元并返回这块内存的可变的i32指针 (*mut i32), 后面我们可以通过这个可变指针来修改或读取指针指向的数据.

这里使用了unsafe关键字来告诉Rust编译器在这段unsafe代码块中我们不在需要编译器来保证内存安全. 因此,内存安全完全由开发者自行控制,

3. let value = unsafe { *ptr };
   这里我们通过unsafe代码块来访问指针指向的内存空间. 但是需要注意的是指针指向的内存空间是没有安全保障的, 可能为null, 可以通过is_null()判断. 之后打印了指针指向的value.

4. unsafe { dealloc(ptr as *mut u8, layout) };
   通过dealloc方法释放之前申请的分配好的内存. 这里使用的mut u8指针类型来释放内存 而不是mut i32, 因为i32是申请的数据类型,也可以是String, 自定义的结构体等.

自定义Vec

Rust中的动态数组是通过NonNull智能指针, 并通过Layout::array来实现指针数组.

#![feature(allocator_api)]

use std::{alloc::Layout, ptr::NonNull};

struct Arr<T: Sized> {
    // 数据指针
    ptr: NonNull<T>,
    // 数组实际长度
    len: usize,
    // 数组容量
    cap: usize,
}

impl<T> Arr<T> {
    fn len(&self) -> usize {
        self.len
    }

    fn cap(&self) -> usize {
        self.cap
    }

    fn new() -> Self {
        Arr {
            // 初始化的指针, 注意: 此时指针数据并未初始化, 如果在数据未初始化之前访问会出现内存访问错误
            ptr: NonNull::dangling(),
            len: 0,
            cap: 0,
        }
    }

    /// 向数组中添加元素
    fn insert(&mut self, value: T) {
        // 当数组长度=数组容量时需要扩容
        if self.len == self.cap {
            // 简单扩容逻辑, 当前容量 * 2
            let new_cap = self.cap.max(1) * 2;
            // 申请一块新的连续的内存布局(内存块, 用来容纳数据指针)
            let new_layout = Layout::array::<T>(new_cap).unwrap();
            // 分配内存单元, 并返回内存块的可变指针, 用于指向数据 -> value:T
            let new_ptr = unsafe { std::alloc::alloc(new_layout) as *mut T };

            // 当数组中长度不为0时, 说明原来有数据.
            if self.len != 0 {
                unsafe {
                    // dealloc old layout
                    // 将原来数组中的指针copy到当前的layout中
                    std::ptr::copy_nonoverlapping(self.ptr.as_ptr(), new_ptr, self.len);
                    // 获取旧的layout, 准备释放这块内存
                    let old_layout = Layout::array::<T>(self.cap).unwrap();
                    // 释放旧的内存
                    std::alloc::dealloc(self.ptr.as_ptr() as *mut u8, old_layout);
                }
            }
            unsafe {
                // 重新设置数组指针
                self.ptr = NonNull::new_unchecked(new_ptr);
            }
            self.cap = new_cap;
        }

        unsafe {
            // 通过std::ptr::write将指针指向指定数据
            // 这里的add方法用于返回数组中 n * align_of::<T>()的offset. 这里传入self.len表示访问数组中最后一个元素.
            // 将最后一个元素指针指向value 完成push操作
            std::ptr::write(self.ptr.as_ptr().add(self.len), value);
        }
        self.len += 1;
    }

    fn pop(&mut self) -> Option<T> {
        if self.len == 0 {
            return None;
        }
        self.len -= 1;
        unsafe { Some(std::ptr::read(self.ptr.as_ptr().add(self.len))) }
    }

    fn get(&self, index: usize) -> Option<&T> {
        if index < self.len {
            return unsafe { Some(&*self.ptr.as_ptr().add(index)) };
        }
        None
    }
}

impl<T> Drop for Arr<T> {
    // 当数组出了作用域之后, 会被释放, 为了同时释放数组中各个指针指向的数据(堆内存中的T)
    fn drop(&mut self) {
        if self.len != 0 {
            // 这里通过for循环将数组中指针指向的所有元素在内存中释放.
            for i in 0..self.len {
                unsafe {
                    std::ptr::drop_in_place(self.ptr.as_ptr().add(i));
                }
            }
            // 释放所有数据后, 释放内存布局
            let layout = Layout::array::<T>(self.cap).unwrap();
            unsafe {
                std::alloc::dealloc(self.ptr.as_ptr() as *mut u8, layout);
            }
        }
    }
}

#[test]
fn test_unsafe() {
    let mut a = Arr::<i32>::new();
    a.insert(1);
    a.insert(3);
    a.insert(2);
    for i in 0..20 {
        a.insert(i);
        println!("a.cap: {:?}", a.cap());
    }

    println!("get {:?}", a.get(0));
    println!("get {:?}", a.get(2));
    println!("get {:?}", a.get(1));
    println!("{}", "=".repeat(20));
    println!("pop {:?}", a.pop());
    println!("pop {:?}", a.pop());
    println!("pop {:?}", a.pop());
}

pop get 方法比较简单这里就不在代码中一一解释了.
这里最重要的几行代码, 再做一个重点说明:

• self.ptr.as_ptr().add(self.len)): 用于获取内存中最后一个元素下一个的内存块(这个是指T类型)内存指针
• std::ptr::write(self.ptr.as_ptr().add(self.len, item): 这里把新的value写入到数组的尾部, 即新申请的那块内存区域.
• std::ptr::drop_in_place(self.ptr.as_ptr().add(i)): drop数组中的所有元素, 确保没有资源泄露的问题.