RXSmallVector - 灵活运用堆栈内存的数组

152 阅读3分钟

问题

  • 如果某个函数(如加密 encode):

那么怎么实现这个数组,比较适合这种情况?

void encode(const char* str) {
    int len = strlen(str);
    std::vector<char> buf;
    buf.reserve(len);
    return exec_encode(str, buf.data());
}

方案一:std::vector

void encode(const char* str) {
    int len = strlen(str);
    std::vector<char> buf;
    // 预分配足够的堆内存,大小等于 len
    buf.reserve(len);
    return exec_encode(str, buf.data(), len);
}
  • 优点:

  • 缺点:

方案二:std::array

const int kEncodeBufMax = 1000 * 1;
void encode(const char* str) {
    int len = strlen(str);
    // 只能分配固定大小的栈内存
    std::array<char, kEncodeBufMax> buf;
    return exec_encode(str, buf.data(), len);
}
  • 优点:

  • 缺点:

方案三:std::vector + std::array

const int kEncodeBufMax = 32;
void encode(const char* str) {
    int len = strlen(str);
    // 预分配固定大小的栈内存
    std::array<char, kEncodeBufMax> buf_stack;
    std::vector<char> buf_heap;
    char* buf = buf_stack.data(); 
    if (len > kEncodeBufMax) {
        // 如果预分配的栈内存不足以处理数据,则转为使用堆内存
        buf_heap.resize(len);
        buf = buf_heap.data();
    }
    return exec_encode(str, buf, len);
}
  • 优点:

  • 缺点:

总结

  • 综上所述,方案三是比较合适的解决方案:

  • 但是方案三这种需要开发者自行维护栈内存/堆内存的切换,显然是不友好且风险较大的。

  • 这时候,我们就可以基于方案三,封装一个“智能”的数组 small vector,可以支持栈内存和堆内存切换。

RXSmallVector

  • small_vector 用于小数据量时使用栈内存,大数据量时使用堆内存。兼顾栈内存不需要内存分配效率高,而堆内存量大管够+随便扩容的特点。

内存布局

其核心的数据结构是通过 union 两个 Stack 成员和 Heap 成员:

  • Stack 成员可以在函数体内申请定长(kSize)的栈内存。

  • Heap 成员则在超出栈内存范围时,申请堆内存,通过 ptr 指针访问。

    struct Stack {
        uint32_t size;
        alignas(T) uint8_t buffer[sizeof(T) * kSize];
    };

    struct Heap {
        uint32_t size;
        uint32_t capacity;
        T* ptr;
    };

    union {
        Stack stack;
        Heap heap;
    } data_;

代码实现

优点:

  • 默认使用栈内存,在栈内存充足时,不会触发内存申请。

  • 当栈内存不足时,会(扩容)申请堆内存,并移动栈内存数据到堆内存中。

  • 如果使用了堆内存(已经申请了堆内存),clear/pop_back 操作都不会销毁原有的堆内存(只清空数据),避免大数据量插入重新申请内存。

  • 支持部分STL方法(提供了迭代器)

#pragma once

#include <cstddef>
#include <cstdlib>
#include <utility>

namespace realx {

template <typename T, uint32_t kSize, typename VT = typename std::decay<T>::type,
          typename = typename std::enable_if<!std::is_array<VT>::value>::type>
class RXSmallVector {
public:
    static_assert(kSize * sizeof(T) <= 128, "Stack memory is too large. It is recommended to use std::vector instead");
    using iterator = T*;
    using const_iterator = const T*;
    using reference = T&;
    using const_reference = const T&;

public:
    // 默认构造
    RXSmallVector() {
        set_use_heap(false);
        set_size(0);
    }

    // 列表初始化
    RXSmallVector(std::initializer_list<T> ilist) : RXSmallVector() {
        assign(ilist);
    }

    ~RXSmallVector() {
        release();
    }

    // 拷贝构造
    RXSmallVector(const RXSmallVector& o) : RXSmallVector() {
        assign(o.begin(), o.end());
    }

    // 移动构造
    RXSmallVector(RXSmallVector&& o) : RXSmallVector() {
        o.swap(*this);
    }

    // 拷贝赋值
    RXSmallVector& operator=(const RXSmallVector& o) {
        if (this != &o) {
            if (o.empty()) {
                release();
            } else {
                assign(o.begin(), o.end());
            }
        }
        return *this;
    }

    // 移动赋值
    RXSmallVector& operator=(RXSmallVector&& o) {
        if (this != &o) {
            o.swap(*this);
        }
        return *this;
    }

    reference operator[](uint32_t index) {
        RX_CHECK(index < size()) << "index(" << index << ") > size(" << size() << ")";
        return begin()[index];
    }

    const_reference operator[](uint32_t index) const {
        RX_CHECK(index < size()) << "index(" << index << ") > size(" << size() << ")";
        return begin()[index];
    }

    inline bool empty() const {
        return !size();
    }

    // 当前已有的数据量
    inline uint32_t size() const {
        return data_.heap.size >> 1;
    }

    // 已申请的数据容量
    inline uint32_t capacity() const {
        return is_use_heap() ? data_.heap.capacity : kSize;
    }

    iterator begin() {
        return is_use_heap() ? heap_ptr() : stack_ptr();
    }

    iterator end() {
        return begin() + size();
    }

    const_iterator begin() const {
        return is_use_heap() ? heap_ptr() : stack_ptr();
    }

    const_iterator end() const {
        return begin() + size();
    }

    template <typename... Args>
    bool emplace_back(Args... args) {
        if (size() < capacity() || reserve(calculate_next_capacity(size()))) {
            new (end()) T(std::forward<Args>(args)...);
            set_size(size() + 1);
            return true;
        }
        return false;
    }

    bool push_back(T&& val) {
        if (size() < capacity() || reserve(calculate_next_capacity(size()))) {
            new (end()) T(std::forward<T>(val));
            set_size(size() + 1);
            return true;
        }
        return false;
    }

    bool push_back(const T& val) {
        if (size() < capacity() || reserve(calculate_next_capacity(size()))) {
            new (end()) T(val);
            set_size(size() + 1);
            return true;
        }
        return false;
    }

    void pop_back() {
        uint32_t current_size = size();
        if (current_size > 0) {
            set_size(current_size - 1);
            end()->~T();
        }
    }

    template <typename Iter>
    void assign(Iter first, Iter last) {
        RX_CHECK(first <= last) << "error: first >= last";
        clear();
        uint32_t count = static_cast<uint32_t>(last - first);
        if (capacity() < count) {
            reserve(calculate_next_capacity(count));
        }
        iterator p = end();
        for (auto it = first; it != last; ++it, ++p) {
            new (p) T(*it);
        }
        set_size(count);
    }

    void assign(std::initializer_list<T> ilist) {
        assign(ilist.begin(), ilist.end());
    }

    // 清除数据
    void clear() {
        for (auto it = begin(), it_end = end(); it != it_end; ++it) {
            it->~T();
        }
        set_size(0);
    }

    // 清除数据。并释放堆内存
    void release() {
        clear();
        if (is_use_heap() && heap_ptr()) {
            free(heap_ptr());
            set_heap_ptr(nullptr);
        }
    }

    // 分配内存,不添加新数据
    bool reserve(uint32_t new_capacity) {
        if (new_capacity <= kSize || (is_use_heap() && new_capacity <= capacity())) {
            return true;
        }
        do {
            bool use_heap = is_use_heap();
            T* new_memory;
            if (use_heap) {
                // move elements from heap(old) memory to heap(new) memory
                new_memory = static_cast<T*>(realloc(heap_ptr(), new_capacity * sizeof(T)));
                if (!new_memory) {
                    RX_LOG(kRXLSError) << "bad_alloc: realloc failed!";
                    return false;
                }
            } else {
                // malloc heap memory
                new_memory = static_cast<T*>(malloc(new_capacity * sizeof(T)));
                if (!new_memory) {
                    RX_LOG(kRXLSError) << "bad_alloc: malloc failed!";
                    return false;
                }
                // move elements from stack memory to heap memory
                uint32_t last_size = size();
                for (uint32_t i = 0; i < last_size; ++i) {
                    new (new_memory + i) T(std::move(stack_ptr()[i]));
                }
                clear();
                set_size(last_size);
                set_use_heap(true);
            }
            set_heap_ptr(new_memory);
        } while (false);
        set_capacity(new_capacity);
        return true;
    }

private:
    inline T* stack_ptr() const {
        return (T*)(data_.stack.buffer);
    }

    inline T* heap_ptr() const {
        return data_.heap.ptr;
    }

    inline void set_heap_ptr(T* ptr) {
        data_.heap.ptr = ptr;
    }

    inline bool is_use_heap() const {
        return (data_.heap.size & kSizeFlagUseHeap);
    }

    inline void set_use_heap(bool is_use_heap) {
        data_.heap.size = (size() << 1) | (is_use_heap ? kSizeFlagUseHeap : 0);
    }

    inline void set_size(uint32_t size) {
        data_.heap.size = (size << 1) | (is_use_heap() ? kSizeFlagUseHeap : 0);
    }

    inline void set_capacity(uint32_t capacity) {
        data_.heap.capacity = capacity;
    }

    inline constexpr uint32_t calculate_next_capacity(uint32_t capacity) const {
        return (capacity * sizeof(T) < 256 || capacity <= 2) ? (capacity << 1) : capacity + (capacity >> 1);
    }

    void swap(RXSmallVector<T, kSize>& o) {
        if (is_use_heap()) {
            if (o.is_use_heap()) {
                // this: heap, o: heap
                std::swap(data_.heap, o.data_.heap);
            } else {
                // this: heap, o: stack
                Heap heap = data_.heap;
                std::swap(data_.stack, o.data_.stack);
                o.data_.heap = heap;
            }
        } else {
            if (o.is_use_heap()) {
                // this: stack, o: heap
                Heap heap = o.data_.heap;
                std::swap(data_.stack, o.data_.stack);
                data_.heap = heap;
            } else {
                // this: stack, o: stack
                std::swap(data_.stack, o.data_.stack);
            }
        }
    }

private:
    struct Stack {
        uint32_t size;
        alignas(T) uint8_t buffer[sizeof(T) * kSize];
    };

    struct Heap {
        uint32_t size;
        uint32_t capacity;
        T* ptr;
    };

    union {
        Stack stack;
        Heap heap;
    } data_;

    static constexpr uint32_t kSizeFlagUseHeap = 1 << 0;

    static_assert(sizeof(Stack) >= sizeof(Heap), "error: sizeof(Stack) < sizeof(Heap)");
};

}  // namespace realx

使用示例

  • 以上述问题为例,使用 RXSmallVector 对比方案三的代码
void encode(const char* str) {
    int len = strlen(str);
    realx::RXSmallVector<char, 32> buf;
    // 如果栈内存不足,可以扩容
    if (buf.capacity() < len) {
        buf.reserve(len);
    }
    return exec_encode(str, &buf[0], len);
}

待改进

  • 这个数据结构查看堆栈信息并不方便

  • 接口不完备,vector 的很多接口在这里都没有实现(当然也是觉得暂时没有使用的必要)。