Compiler Explorer

Source code

#include <memory_resource>
#include <unordered_map>
#include <string>
#include <vector>
#include <iostream>

// Transparent memory resource that logs allocations
struct TracingResource final : std::pmr::memory_resource {
    TracingResource(std::string_view label, std::pmr::memory_resource* upstream) : label_(label), upstream_(upstream) {}
    std::pmr::memory_resource* upstream() const { return upstream_; }
private:
    // Allocation
    void* do_allocate(std::size_t bytes, std::size_t alignment) override {
        std::cerr << label_ << " allocating " << bytes << " bytes of memory, at " << alignment << " bytes alignment.\n";
        return upstream_->allocate(bytes,alignment);
    }
    // Deallocation
    void do_deallocate(void* p, std::size_t bytes, std::size_t alignment) override {
        std::cerr << label_ << " deallocating " << bytes << " bytes of memory, at " << alignment << " bytes alignment.\n";
        return upstream_->deallocate(p,bytes,alignment);
    }
    bool do_is_equal(const std::pmr::memory_resource& other) const noexcept override {
        auto ptr = dynamic_cast<const TracingResource*>(&other);
        if (ptr == this) return true;
        if (&other == upstream_) return true;
        return false;
    }

std::string label_;
    std::pmr::memory_resource* upstream_;
};

// Side-note, this custom hash allows us to take a std::string_view
// as argument and use std::unordered_map methods, without
// repeated copies and conversions between std::string,
// std::pmr::string, or even pmr strings that do not use 
// the same backing resource.
// For more context see Heterogenous lookup.
struct StringHash {
  using is_transparent = void;
  [[nodiscard]] size_t operator()(const char *txt) const {
    return std::hash<std::string_view>{}(txt);
  }
  [[nodiscard]] size_t operator()(std::string_view txt) const {
    return std::hash<std::string_view>{}(txt);
  }
  [[nodiscard]] size_t operator()(const std::string &txt) const {
    return std::hash<std::string>{}(txt);
  }
  [[nodiscard]] size_t operator()(const std::pmr::string &txt) const {
    return std::hash<std::pmr::string>{}(txt);
  }
};

// An allocator aware type
struct WordCounter {
    // The demonstration uses std::pmr for brevity, however, 
    // the same approach can be applied to static allocators, 
    // with the caveat that you would need to use
    // std::scoped_allocator_adaptor.
    using allocator_type = std::pmr::polymorphic_allocator<>;

// Default constructor, with an optional allocator argument
    WordCounter(const allocator_type& alloc = {}) : counter_(alloc) {}
    // Copy constructor
    WordCounter(const WordCounter& other, 
                const allocator_type& alloc = {}) 
      : counter_(other.counter_, alloc) {}
    // Move constructor, note that the move is conditional
    // if (alloc != other.alloc), we have to default to a copy.
    WordCounter(WordCounter&& other, 
                const allocator_type& alloc = {})
      : counter_(std::move(other.counter_), alloc) {}

// Assignment operators remain without change
    WordCounter& operator=(const WordCounter&) = default;
    WordCounter& operator=(WordCounter&&) = default;

// For demonstration
    void increment(std::string_view word) {
        if (auto it = counter_.find(word); it != counter_.end())
            ++(it->second);
        else
            counter_.emplace(word, 1);
    }
    void decrement(std::string_view word) {
        if (auto it = counter_.find(word);
            it != counter_.end() && it->second != 0)
            --(it->second);
    }
private:
    // Our storage that gets customized with the allocator
    std::pmr::unordered_map<
        std::pmr::string, uint64_t,
        // custom hash and std::equal_to<>
        // to allow for heterogenous lookup
        StringHash, std::equal_to<>> counter_; 
};

// If we cannot put the allocator as the last argument,
// we can use the std::allocator_arg_t tag.
template <typename... Types>
struct VariadicType {
    using allocator_type = std::pmr::polymorphic_allocator<>;

// Tag first, followed by the allocator, other arguments follow.
    VariadicType(std::allocator_arg_t, const allocator_type& alloc, 
        Types&&... args) {}
    // And we need to keep the non-allocator version around.
    VariadicType(Types&&... args) {}

// For copy/move, we can use the normal style
    VariadicType(const VariadicType&, const allocator_type = {}) {}
    VariadicType(VariadicType&&, const allocator_type = {});
};

int main() {
    std::pmr::monotonic_buffer_resource buffer;
    
    // We can wrap allocator aware types in containers and the outer
    // allocator will be correctly applied to the elements.
    // Note: for non-pmr, you would need std::scoped_allocator_adaptor
    TracingResource t1("WordCounter:", &buffer);
    std::pmr::vector<WordCounter> counters(&t1);

// Construct WordCounter, using the allocator constructor.
    // The vector will allocate memory from the allocator.
    counters.emplace_back();
    // The map will allocate the bucket array and the node using the 
    // allocator. The string fits into small string optimization.
    counters[0].increment("hello");
    // Same as above, except the bucket array is already allocated
    // and the string also needs to allocate.
    counters[0].increment("this string is long enough");

TracingResource t2("Variadic:", &buffer);
    std::pmr::vector<VariadicType<int,int,int>> variadic(&t2);
    variadic.emplace_back(1,2,3); // Calls the allocator constructor
}