smux_rust 0.2.1

use std::sync::Arc;
use tokio::sync::Mutex;

/// De Bruijn 序列查找表，用于快速计算最高位位置
const DEBRUIJN_POS: [u8; 32] = [
    0, 9, 1, 10, 13, 21, 2, 29, 11, 14, 16, 18, 22, 25, 3, 30, 8, 12, 20, 28, 15, 17, 24, 7,
    19, 27, 23, 6, 26, 5, 4, 31,
];

/// 内存分配器，用于优化帧的内存分配
/// 使用对象池来减少内存分配和碎片化
pub struct Allocator {
    /// 缓冲区池数组，每个池对应 2^n 大小的缓冲区
    buffers: Vec<Arc<Mutex<Vec<Vec<u8>>>>>,
}

impl Allocator {
    /// 创建新的分配器
    /// 为 1B 到 64KB 的缓冲区提供池化分配
    /// 内存碎片化保证不超过 50%
    pub fn new() -> Self {
        let mut buffers = Vec::with_capacity(17); // 1B -> 64K (2^0 到 2^16)
        for _i in 0..17 {
            buffers.push(Arc::new(Mutex::new(Vec::new())));
        }
        Self { buffers }
    }

    /// 从池中获取一个合适大小的缓冲区
    /// 返回的缓冲区容量至少为 size
    pub async fn get(&self, size: usize) -> Option<Vec<u8>> {
        if size == 0 || size > 65536 {
            return None;
        }

        let bits = msb(size);
        // 如果 size 正好是 2^bits，使用 bits，否则使用 bits+1
        let actual_bits = if size == 1 << bits { bits } else { bits + 1 };
        
        let pool = &self.buffers[actual_bits as usize];
        let mut pool_guard = pool.lock().await;

        if let Some(mut buf) = pool_guard.pop() {
            buf.resize(size, 0);
            Some(buf)
        } else {
            // 如果池为空，分配新缓冲区
            let mut buf = vec![0; 1 << actual_bits];
            buf.resize(size, 0);
            Some(buf)
        }
    }

    /// 将缓冲区返回到池中
    /// 缓冲区的容量必须是 2^n
    pub async fn put(&self, mut buf: Vec<u8>) -> Result<(), String> {
        let cap = buf.capacity();
        if cap == 0 || cap > 65536 {
            return Err("allocator Put() incorrect buffer size".to_string());
        }

        let bits = msb(cap);
        if cap != 1 << bits {
            return Err("allocator Put() incorrect buffer size".to_string());
        }

        // 清空缓冲区
        buf.clear();
        buf.resize(cap, 0);

        let pool = &self.buffers[bits as usize];
        let mut pool_guard = pool.lock().await;
        pool_guard.push(buf);
        Ok(())
    }
}

impl Default for Allocator {
    fn default() -> Self {
        Self::new()
    }
}

/// 计算最高有效位的位置
/// 使用 De Bruijn 序列算法
fn msb(size: usize) -> u8 {
    let v = size as u32;
    let v = v | (v >> 1);
    let v = v | (v >> 2);
    let v = v | (v >> 4);
    let v = v | (v >> 8);
    let v = v | (v >> 16);
    DEBRUIJN_POS[((v.wrapping_mul(0x07C4ACDD)) >> 27) as usize]
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_allocator() {
        let alloc = Allocator::new();

        // 测试分配和回收
        let buf1 = alloc.get(100).await.unwrap();
        assert_eq!(buf1.len(), 100); // 分配后长度应该是请求的大小
        assert!(buf1.capacity() >= 100); // 容量应该至少是请求的大小
        let cap1 = buf1.capacity();
        alloc.put(buf1).await.unwrap();

        // 再次获取应该从池中获取
        let buf2 = alloc.get(100).await.unwrap();
        assert_eq!(buf2.capacity(), cap1);
        assert_eq!(buf2.len(), 100);
        
        // 测试不同大小的分配
        let buf3 = alloc.get(1000).await.unwrap();
        assert_eq!(buf3.len(), 1000);
        assert!(buf3.capacity() >= 1000);
        alloc.put(buf3).await.unwrap();
        
        // 测试无效大小
        assert!(alloc.get(0).await.is_none());
        assert!(alloc.get(70000).await.is_none());
    }

    #[test]
    fn test_msb() {
        // msb 返回最高有效位的索引（从 0 开始）
        // 1 = 0b1, 最高位是第 0 位
        assert_eq!(msb(1), 0);
        // 2 = 0b10, 最高位是第 1 位
        assert_eq!(msb(2), 1);
        // 3 = 0b11, 最高位是第 1 位
        assert_eq!(msb(3), 1);
        // 4 = 0b100, 最高位是第 2 位
        assert_eq!(msb(4), 2);
        // 5 = 0b101, 最高位是第 2 位
        assert_eq!(msb(5), 2);
        // 8 = 0b1000, 最高位是第 3 位
        assert_eq!(msb(8), 3);
        // 16 = 0b10000, 最高位是第 4 位
        assert_eq!(msb(16), 4);
    }
}