TieredMap.hpp

// <TieredMap> -*- C++ -*-
 
#pragma once
 
#include <iostream>
#include <utility>
#include <memory>
 
#include "sparta/utils/Utils.hpp"
#include "sparta/utils/SpartaException.hpp"
#include "sparta/utils/SpartaAssert.hpp"
 
 
namespace sparta
{
    template <typename KeyT=uint64_t, typename ValT=void*>
    class TieredMap
    {
    public:
 
 
        typedef std::pair<const KeyT, ValT> pair_t;
 
        typedef std::vector<void*> node_t;
 
        // Forward declaration for friendship with iterator
        class const_iterator;
 
        class iterator {
        public:
            typedef typename std::vector<std::unique_ptr<pair_t>>::iterator internal_iter;
 
            iterator(internal_iter itr) :
                itr_(itr)
            { }
 
            iterator& operator++() { ++itr_; return *this; }
            iterator operator++(int) { iterator ret(itr_); ++itr_; return ret; }
 
            iterator& operator=(const iterator& rhp) {
                itr_ = rhp.itr_;
                return *this;
            }
 
            bool operator==(const iterator& rhp) const {
                return rhp.itr_ == itr_;
            }
 
            bool operator!=(const iterator& rhp) const {
                return rhp.itr_ != itr_;
            }
 
            bool operator==(const const_iterator& rhp) const {
                return rhp.itr_ == itr_;
            }
 
            bool operator!=(const const_iterator& rhp) const {
                return rhp.itr_ != itr_;
            }
 
            ValT& operator*() { return (*itr_)->second; }
            const ValT& operator*() const { return (*itr_)->second; }
 
            ValT* operator->() { return &(*itr_)->second; }
            const ValT* operator->() const { return &(*itr_)->second; }
 
            friend class const_iterator;
 
        private:
            internal_iter itr_;
        };
 
        class const_iterator {
        public:
            typedef typename std::vector<std::unique_ptr<pair_t>>::const_iterator internal_iter;
 
            const_iterator(internal_iter itr) :
                itr_(itr)
            { }
 
            const_iterator& operator++() { ++itr_; return *this; }
            const_iterator operator++(int) { const_iterator ret(itr_); ++itr_; return ret; }
 
            const_iterator& operator=(const const_iterator& rhp) {
                itr_ = rhp.itr_;
                return *this;
            }
 
            bool operator==(const const_iterator& rhp) const {
                return rhp.itr_ == itr_;
            }
 
            bool operator!=(const const_iterator& rhp) const {
                return rhp.itr_ != itr_;
            }
 
            bool operator==(const iterator& rhp) const {
                return rhp.itr_ == itr_;
            }
 
            bool operator!=(const iterator& rhp) const {
                return rhp.itr_ != itr_;
            }
 
            const ValT& operator*() const { return (*itr_)->second; }
 
            const ValT* operator->() const { return &(*itr_)->second; }
 
            friend class iterator;
 
        private:
            internal_iter itr_;
        };
 
 
 
        TieredMap(uint64_t node_size=512) :
            total_size_(0),
            num_nodes_(0), // First node is allocated in ctor
            num_tiers_(0), // First tier is allocated in ctor
            node_size_(node_size),
            tier_shift_(0),
            tier_idx_mask_(~computeMask<uint64_t>(node_size, tier_shift_)) // Computes tier_shift_
        {
            if(false == isPowerOf2(node_size_)){
                throw SpartaException("node_size must be a power of 2, is ")
                    << node_size;
            }
 
            allocateRoot_();
        }
 
        ~TieredMap() {
            // Free all nodes
            freeNode_(t0_);
        }
 
 
 
        uint64_t size() const { return pairs_.size(); }
 
        uint64_t getNumTiers() const { return num_tiers_; }
 
        uint64_t getNumNodes() const { return num_nodes_; }
 
        uint64_t getEstimatedMemory() const {
            return sizeof(TieredMap) // This map
                + (num_nodes_ * sizeof(node_t)) // Each node (container)
                + (num_nodes_ * sizeof(void*) * 0.7) // Content of each node. 0.7 for average utilization
                + (sizeof(pairs_))
                + (pairs_.size() * sizeof(typename decltype(tier_shifts_)::value_type)) // Pair container and content
                + (sizeof(tier_shifts_))
                + (tier_shifts_.size() * sizeof(typename decltype(tier_shifts_)::value_type)); // tier_shift container and content
        }
 
 
 
        const_iterator begin() const { return pairs_.begin(); }
        const_iterator end() const { return pairs_.end(); }
        iterator begin() { return pairs_.begin(); }
        iterator end() { return pairs_.end(); }
 
 
 
        void clear() {
            freeNode_(t0_);
            pairs_.clear();
            tier_shifts_.clear();
            num_nodes_ = 0;
            num_tiers_ = 0;
 
            allocateRoot_(); // Adds first node and tier
        }
 
        const pair_t* find(const KeyT& k) const {
            const pair_t* result = tryGet(k);
            if(!result){
                return nullptr;
            }
            return result;
        }
 
        pair_t* find(const KeyT& k) {
            pair_t* result = tryGet(k);
            if(!result){
                return nullptr;
            }
            return result;
        }
 
        ValT& operator[](const KeyT& k) {
            pair_t* result = tryGet(k);
            if(result){
                return (result->second);
            }
            return set_(k, 0);
        }
 
        const ValT& operator[](const KeyT& k) const {
            pair_t* result = tryGet(k);
            if(result){
                return (result->second);
            }
            return set_(k, 0);
        }
 
        pair_t* tryGet(const KeyT& k) {
            void* v = t0_;
            size_t tidx = 0;
 
            uint64_t vidx = k >> tier_shifts_[tidx];
            if(vidx >= node_size_){
                return nullptr;
            }
 
            do{
                vidx = k >> tier_shifts_[tidx];
                vidx &= tier_idx_mask_;
 
                node_t& node = *static_cast<node_t*>(v);
                if(vidx >= node.size()){
                    return nullptr;
                }
                v = node[vidx]; // Points to node in next tier or final pair result
                if(!v){
                    return nullptr;
                }
                ++tidx;
            }while(tidx < num_tiers_);
 
            return static_cast<pair_t*>(v);
        }
 
        const pair_t* tryGet(const KeyT& k) const {
            void* v = t0_;
            size_t tidx = 0;
 
            uint64_t vidx = k >> tier_shifts_[tidx];
            if(vidx >= node_size_){
                return nullptr;
            }
 
            do{
                vidx = k >> tier_shifts_[tidx];
                vidx &= tier_idx_mask_;
 
                const node_t& node = *static_cast<const node_t*>(v);
                if(vidx >= node.size()){
                    return nullptr;
                }
                v = node[vidx]; // Points to next tier or final pair result
                if(!v){
                    return nullptr;
                }
                ++tidx;
            }while(tidx < num_tiers_);
 
            return static_cast<const pair_t*>(v);
        }
 
 
    private:
 
        void allocateRoot_() {
            sparta_assert(num_nodes_ == 0
                        && num_tiers_ == 0
                        && tier_shifts_.size() == 0,
                        "num_nodes_, num_tiers_, and tier_shifts_.size() must be 0 during allocateRoot_");
            tier_shifts_.push_back(0);
            sparta_assert(tier_shifts_.size() == 1); // Expects exactly 1 shift to start
            t0_ = new node_t; // Allocate a node
            num_nodes_++;
            num_tiers_++;
        }
 
        ValT& set_(const KeyT& k, const ValT n) {
            void* v = t0_;
            size_t tidx = 0;
 
            uint64_t vidx = k >> tier_shifts_[tidx]; // Value index (within tier)
            if(vidx >= node_size_){
                // Add a new tier because current tier does not represent a
                // large enough space to contain this key
                addTier_();
                return set_(k, n);
            }
 
            do{
                // Compute index for this tier
                vidx = k >> tier_shifts_[tidx];
                vidx &= tier_idx_mask_;
 
                node_t& node = *static_cast<node_t*>(v);
                if(vidx >= node.size()){
                    node.resize(vidx+1, nullptr);
                }
                void*& vt = node[vidx]; // Points to next tier or final pair result
                if(!vt){
                    if(tidx < num_tiers_ - 1){
                        //std::cout << "Allocating new vector at tier[0x" << std::hex << vidx << std::dec << "] tidx " << tidx << ", key=0x" << std::hex << k << std::dec << std::endl;
                        // Create a new node at node[vidx]
                        vt = new node_t();
                        ++num_nodes_;
                    }else{
                        //std::cout << "Allocating new pair at tier[0x" << std::hex << vidx << std::dec << "] tidx " << tidx << ", key=0x" << std::hex << k << std::dec << std::endl;
                        // Create a new result pair at node[vidx]
                        pair_t* npair = new pair_t(k, n);
                        pairs_.emplace_back(npair);
                        v = vt = npair;
                        break;
                    }
                }else{
                    if(tidx == num_tiers_ - 1){
                        pair_t* p = static_cast<pair_t*>(vt);
                        p->second = n;
                        v = p;
                        break;
                    }
                }
                v = vt;
                ++tidx;
            }while(tidx < num_tiers_);
 
            return static_cast<pair_t*>(v)->second;
        }
 
        void addTier_() {
            uint64_t new_shift = tier_shifts_[0] + tier_shift_;
            tier_shifts_.insert(tier_shifts_.begin(), new_shift);
            void* old = t0_;
            t0_ = new node_t;
            t0_->push_back(old); // Old tier at idx=0
            ++num_nodes_;
            ++num_tiers_;
        }
 
        void freeNode_(node_t* n, uint64_t tier=0) {
            sparta_assert(n);
            if(tier < num_tiers_ - 1){
                for(node_t::iterator i = n->begin(); i != n->end(); ++i){
                    node_t* child = static_cast<node_t*>(*i);
                    if(child){
                        freeNode_(child, tier+1);
                    }
                }
            }else{
                for(node_t::iterator i = n->begin(); i != n->end(); ++i){
                    pair_t* p = static_cast<pair_t*>(*i);
                    if(p){
                        // Do not delete p, it is managed by the pairs_ vector.
                    }
                }
            }
 
            delete n;
        }
 
        const uint64_t total_size_;
 
        uint64_t num_nodes_;
 
        uint64_t num_tiers_;
 
        const uint64_t node_size_;
 
        uint64_t tier_shift_ = 0;
 
        const uint64_t tier_idx_mask_;
 
        std::vector<std::unique_ptr<pair_t>> pairs_;
 
        std::vector<uint64_t> tier_shifts_;
 
        node_t* t0_;
 
    }; // class TieredMap
 
} // namespace sparta