Rust手动管理内存

为了确保安全性，Rust在零抽象的基础上，限制了很多易导致潜在bug的操作，比如直接指针操作，以及长生命周期等等，但是在某些情况下，如果涉及到操作比较底层的数据，往往难免自行维护内存。

本文主要以layout库来进行手动内存维护演示。

假设要实现一个内存分配器，但是需要满足以下条件：

• 手动从堆中申请内存
• 手动执行内存回收
• 内存空间能够被复用

由于要手动回收内存，可行的方案是实现一个智能指针包裹裸指针，为了将裸指针指向的内存区视为待分配的数据结构，可以采用强制类型转换。

use std::alloc::{alloc, dealloc, Layout};
use std::ops::{Deref, DerefMut};
use std::fmt;
use fmt::Debug;
use std::collections::HashMap;

#[derive(Debug)]
struct Header {
    a: i32,
}

struct Cap<T>(*mut u8, Layout, std::marker::PhantomData<T>);

impl<T> Cap<T> {
    fn new() -> Self {
        let layout = Layout::new::<T>();
        unsafe {
            let ptr = alloc(layout);
            Cap(ptr, layout, std::marker::PhantomData)
        }
    }
}

impl<T: Debug> fmt::Display for Cap<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        unsafe {
            write!(f, "({:?})", & *(self.0 as *mut T))
        }
    }
}

impl<T> Drop for Cap<T> {
    fn drop(&mut self) {}
}

impl<T> Deref for Cap<T> {
    type Target = T;
    fn deref(&self) -> &T {
        unsafe {
            & *(self.0 as *mut T)
        }
    }
}

impl<T> DerefMut for Cap<T> {
    fn deref_mut(&mut self) -> &mut T {
        unsafe {
            &mut *(self.0 as *mut T)
        }
    }
}

struct MemCache<T> {
    cache_map: HashMap<usize, Vec<Cap<T>>>,
}

impl<T> MemCache<T> {
    fn new() ->Self {
        MemCache {
            cache_map: HashMap::new(),
        }
    }

    fn alloc(&mut self) -> Cap<T> {
        let obj_size = std::mem::size_of::<T>();
        match self.cache_map.get_mut(&obj_size) {
            Some(bucket) => {                
                if bucket.is_empty() {
                    return Cap::<T>::new()
                } else {
                    let ele = bucket.remove(0);
                    return ele
                }
            },
            _ => {
                return Cap::<T>::new();
            },
        }
    }

    fn drop(&mut self, d: Cap<T>) {
        let obj_size = std::mem::size_of::<T>(); 
        match self.cache_map.get_mut(&obj_size) {
            Some(bucket) => {
                bucket.push(d);
            }
            _ => {
                let mut b = Vec::new();
                b.push(d);
                self.cache_map.insert(obj_size, b);
            }
        };
    }
}

fn main() {
    let mut cache = MemCache::<Header>::new();

    let mut t1 = cache.alloc();
    t1.a = 38;
    println!("Cap {:#}", t1);
    cache.drop(t1);

    let mut t2 = cache.alloc();
    t2.a = 36;
    println!("Cap {:#}", t2);
    cache.drop(t2);
}

主要思路如下：

• 用智能指针封装裸指针，在deref时，将裸指针转换为具体的类型
• 释放智能指针时，用一个HashMap缓存住，在分配阶段如果有缓存则用缓存

上述的实现不能自动执行drop，需要手动执行，所以在Cap的Drop实现是一个空函数。

为了实现自动释放，需要在Cap中记录下MemCache，参考实现如下：

use std::fmt;
use std::mem;
use std::ops::{Deref, DerefMut};
use std::ptr;

use crossbeam_queue::SegQueue;
use stable_deref_trait::StableDeref;

pub trait Allocate {
    fn alloc() -> Self;
}

#[derive(Debug)]
struct Header {
    a: i32,
}

impl Allocate for Header {
    fn alloc() -> Self {
        Header {
            a: 89,
        }
    }
}

pub struct BytePool<T>
where
    T: Allocate,
{
    small_bask: SegQueue<T>,
}

pub struct ByteBuffer<'a, T: Allocate> {
    data: mem::ManuallyDrop<T>,
    pool: &'a BytePool<T>,
}

impl<T: Allocate + fmt::Debug> fmt::Debug for ByteBuffer<'_, T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}", self.data.deref())
    }
}

impl<T: Allocate + fmt::Debug> fmt::Display for ByteBuffer<'_, T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}", self.data.deref())
    }
}

impl<T: Allocate> Default for BytePool<T> {
    fn default() -> Self {
        BytePool::<T> {
            small_bask: SegQueue::new(),
        }
    }
}

impl<T: Allocate> BytePool<T> {
    pub fn new() -> Self {
        BytePool::default()
    }

    pub fn alloc(&self) -> ByteBuffer<'_, T> {
        let list = &self.small_bask;
        if let Some(el) = list.pop() {
            return ByteBuffer::new(el, self);
        }

        let data = T::alloc();
        ByteBuffer::new(data, self)
    }

    fn push_raw_block(&self, buffer: T) {
        self.small_bask.push(buffer);
    }
}

impl<'a, T: Allocate> Drop for ByteBuffer<'a, T> {
    fn drop(&mut self) {
        let data = mem::ManuallyDrop::into_inner(unsafe { ptr::read(&self.data) });
        self.pool.push_raw_block(data);
    }
}

impl<'a, T: Allocate> ByteBuffer<'a, T> {
    fn new(data: T, pool: &'a BytePool<T>) -> Self {
        ByteBuffer {
            data: mem::ManuallyDrop::new(data),
            pool,
        }
    }

    pub fn size(&self) -> usize {
        std::mem::size_of::<T>()
    }
}

impl<'a, T: Allocate> Deref for ByteBuffer<'a, T> {
    type Target = T;

    #[inline]
    fn deref(&self) -> &Self::Target {
        self.data.deref()
    }
}

impl<'a, T: Allocate> DerefMut for ByteBuffer<'a, T> {
    #[inline]
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.data.deref_mut()
    }
}

unsafe impl<'a, T: StableDeref + Allocate> StableDeref for ByteBuffer<'a, T> {}

fn main() {
    let p = BytePool::<Header>::new();
    {
        let mut t = p.alloc();
        t.a = 1;
        println!("{:#}", t);
    }

    let mut t = p.alloc();
    t.a = 1;
    println!("{:#}", t);
}