memory.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: AGPL-3.0-or-later
// Copyright 2026 Johnson Ogundeji
#pragma once
 
 
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <memory>
#include <new>
#include <vector>
 
namespace signet::forge {
 
class Arena {
public:
    explicit Arena(size_t block_size = 64 * 1024)
        : block_size_(block_size) {}
 
    ~Arena() = default;
 
    void* allocate(size_t size, size_t alignment = 8) {
        if (size == 0) return nullptr;
 
        // Try to fit in the current (last) block
        if (!blocks_.empty()) {
            Block& current = blocks_.back();
            void* ptr = try_allocate_from(current, size, alignment);
            if (ptr) {
                total_used_ += size;
                return ptr;
            }
        }
 
        // Current block cannot satisfy — allocate a new block
        if (size > SIZE_MAX - alignment) return nullptr; // overflow guard
        size_t new_block_size = (size + alignment > block_size_)
                                    ? (size + alignment)
                                    : block_size_;
        allocate_block(new_block_size);
 
        Block& fresh = blocks_.back();
        void* ptr = try_allocate_from(fresh, size, alignment);
        // A fresh block is guaranteed to fit (we sized it above)
        if (ptr) {
            total_used_ += size;
        }
        return ptr;
    }
 
    void* allocate_zeroed(size_t size, size_t alignment = alignof(std::max_align_t)) {
        void* ptr = allocate(size, alignment);
        if (ptr) std::memset(ptr, 0, size);
        return ptr;
    }
 
    template <typename T>
    T* allocate_array(size_t count) {
        if (count == 0) return nullptr;
        if (count > SIZE_MAX / sizeof(T)) return nullptr; // overflow guard
        void* raw = allocate(count * sizeof(T), alignof(T));
        return static_cast<T*>(raw);
    }
 
    const uint8_t* copy(const uint8_t* data, size_t size) {
        if (size == 0 || data == nullptr) return nullptr;
        void* dst = allocate(size, 1);
        if (!dst) return nullptr; // allocation failed — propagate null
        std::memcpy(dst, data, size);
        return static_cast<const uint8_t*>(dst);
    }
 
    const char* copy_string(const std::string& str) {
        size_t len = str.size() + 1; // include null terminator
        void* dst = allocate(len, 1);
        if (!dst) return nullptr; // allocation failed — propagate null
        std::memcpy(dst, str.c_str(), len);
        return static_cast<const char*>(dst);
    }
 
    [[nodiscard]] size_t bytes_used() const { return total_used_; }
 
    void reset() {
        for (auto& block : blocks_) {
            block.used = 0;
        }
        total_used_ = 0;
    }
 
    Arena(const Arena&) = delete;            
    Arena& operator=(const Arena&) = delete; 
 
    Arena(Arena&&) noexcept = default;            
    Arena& operator=(Arena&&) noexcept = default; 
 
private:
    struct Block {
        std::unique_ptr<uint8_t[]> data; 
        size_t size;                     
        size_t used;                     
    };
 
    std::vector<Block> blocks_;
    size_t block_size_;
    size_t total_used_ = 0;
 
    static void* try_allocate_from(Block& block, size_t size, size_t alignment) {
        // Round up 'used' to the next multiple of alignment
        size_t aligned_offset = align_up(block.used, alignment);
        if (aligned_offset + size > block.size) {
            return nullptr;
        }
        void* ptr = block.data.get() + aligned_offset;
        block.used = aligned_offset + size;
        return ptr;
    }
 
    void allocate_block(size_t size) {
        try {
            Block block;
            block.data = std::make_unique<uint8_t[]>(size);
            block.size = size;
            block.used = 0;
            blocks_.push_back(std::move(block));
        } catch (const std::bad_alloc&) {
            // Push an empty block so callers see a block with size 0
            Block empty;
            empty.size = 0;
            empty.used = 0;
            blocks_.push_back(std::move(empty));
        }
    }
 
    static size_t align_up(size_t offset, size_t alignment) {
        size_t mask = alignment - 1;
        // Overflow guard: if offset + mask would wrap, return SIZE_MAX (CWE-190)
        if (offset > SIZE_MAX - mask) return SIZE_MAX;
        return (offset + mask) & ~mask;
    }
};
 
} // namespace signet::forge