MemoryTracker.hpp

Go to the documentation of this file.

#pragma once
#include <vector>
#include <mutex>
#include <cstdint>
#include <iostream>
#include <cstdint>
 
 
namespace memviz {
    /*
    * @brief Metadata stored for each tracked allocation.
    */
    struct AllocationInfo {
        std::uintptr_t ptr;
        std::size_t size;
    };
    namespace detail {
        inline thread_local bool in_hook = false;
 
        struct HookGuard {
            bool armed = false;
            HookGuard() {
                if (!in_hook) { in_hook = true; armed = true; }
            }
            ~HookGuard() {
                if (armed) in_hook = false;
            }
 
        };
 
        inline auto lower_bound_addr(std::vector<AllocationInfo>& t, std::uintptr_t addr) {
            return std::lower_bound(t.begin(), t.end(), addr,
                [](const AllocationInfo& a, std::uintptr_t k){ return a.ptr < k; });
        }
    } // namespace detail 
 
 
static inline std::uintptr_t key_of(const void* p) noexcept {
    return reinterpret_cast<std::uintptr_t>(p);
}
 
inline auto& table() {
    static auto* m = new std::vector<AllocationInfo>();
    return *m;
}
 
 
inline std::mutex& lock() {
    static auto* mu = new std::mutex();
    return *mu;
}
 
inline void on_alloc(const void* ptr, std::size_t sz) {
    if (!ptr) 
        return;
 
    detail::HookGuard g;    
 
    if (!g.armed) 
        return;          // already inside: skip
 
    std::scoped_lock _(lock());
 
    const auto addr = key_of(ptr);
 
    auto it = detail::lower_bound_addr(table(), addr);
    if (it != table().end() && it->ptr == addr) {
        it->size = sz;
    } else {
        table().insert(it, AllocationInfo{addr, sz});   // keeps vector sorted
    }
}
 
 
inline void on_free(const void* ptr) {
    if (!ptr) return;
    detail::HookGuard g;
    if (!g.armed) return;
 
    std::scoped_lock lk(lock());
    auto& t = table();
    const auto addr = key_of(ptr);
 
    auto it = detail::lower_bound_addr(t, addr);
    if (it != t.end() && it->ptr == addr) {
        t.erase(it);               
    }
}
 
[[nodiscard]] inline bool isLive(const void* ptr) {
    std::scoped_lock lk(lock());
    const auto& t = table();
    const auto addr = key_of(ptr);
 
    auto it = detail::lower_bound_addr(const_cast<std::vector<AllocationInfo>&>(t), addr);
    return it != t.end() && it->ptr == addr;
}
 
inline void reportLeaks(std::ostream& os = std::cout) {
    std::scoped_lock g(lock());
    std::size_t total_bytes = 0;
    int i = 0;
 
    for (const auto& info : table()) {
 
        os << "[" << i++ << "] LEAK " << reinterpret_cast<void*>(info.ptr) << " (" << info.size << " bytes)\n";
        total_bytes += info.size;
    }
    os << "[SUMMARY] " << i << " potential leak(s), total " << total_bytes << " bytes.\n";
}
 
inline void printAllocationLog(std::ostream& os = std::cout) {
    std::scoped_lock g(lock());
    std::size_t i = 0;
    for (const auto& info : table()) {
        os << "[" << i++ << "] LIVE " << reinterpret_cast<void*>(info.ptr) << " (" << info.size << " bytes)\n";
    }
}
 
} // namespace memviz