LRUCache.hpp

Go to the documentation of this file.

#pragma once
 
#include <cstddef>
#include <list>
#include <memory>
#include <mutex>
#include <shared_mutex>
#include <unordered_map>
 
#include "vc/core/types/Cache.hpp"
 
namespace volcart
{
 
struct NoOpMutex {
};
 
template <typename TKey, typename TValue, class TMutex = std::shared_mutex>
class LRUCache final : public Cache<TKey, TValue>
{
public:
    using BaseClass = Cache<TKey, TValue>;
    using BaseClass::capacity_;
 
    using TPair = std::pair<TKey, TValue>;
 
    using TListIterator = typename std::list<TPair>::iterator;
 
    using Pointer = std::shared_ptr<LRUCache<TKey, TValue>>;
 
    LRUCache() : BaseClass() {}
 
    explicit LRUCache(std::size_t capacity) : BaseClass(capacity) {}
 
    static auto New() -> Pointer
    {
        return std::make_shared<LRUCache<TKey, TValue>>();
    }
 
    static auto New(std::size_t capacity) -> Pointer
    {
        return std::make_shared<LRUCache<TKey, TValue>>(capacity);
    }
    void setCapacity(std::size_t capacity) override
    {
        std::unique_lock lock(cache_mutex_);
        if (capacity <= 0) {
            throw std::invalid_argument(
                "Cannot create cache with capacity <= 0");
        }
        capacity_ = capacity;
        evict();
    }
 
    auto capacity() const -> std::size_t override { return capacity_; }
 
    auto size() const -> std::size_t override { return lookup_.size(); }
    auto get(const TKey& k) -> TValue override
    {
        std::unique_lock lock(cache_mutex_);
        auto lookupIter = lookup_.find(k);
        if (lookupIter == std::end(lookup_)) {
            throw std::invalid_argument("Key not in cache");
        }
 
        items_.splice(std::begin(items_), items_, lookupIter->second);
        return lookupIter->second->second;
    }
 
    void put(const TKey& k, const TValue& v) override
    {
        std::unique_lock lock(cache_mutex_);
        // If already in cache, need to refresh it
        auto lookupIter = lookup_.find(k);
        if (lookupIter != std::end(lookup_)) {
            auto& [key, val] = *lookupIter->second;
            if (BaseClass::on_eject_) {
                BaseClass::on_eject_(key, val);
            }
            items_.erase(lookupIter->second);
            lookup_.erase(lookupIter);
        }
 
        items_.push_front(TPair(k, v));
        lookup_[k] = std::begin(items_);
        evict();
    }
 
    auto contains(const TKey& k) -> bool override
    {
        std::shared_lock lock(cache_mutex_);
        return lookup_.find(k) != std::end(lookup_);
    }
 
    void onEvict(std::function<bool(TKey&, TValue&)> fn) override
    {
        BaseClass::on_eject_ = fn;
    }
 
    void resetOnEvict() override { BaseClass::on_eject_ = {}; }
 
    void evict() override
    {  // Already below capacity
        if (lookup_.size() <= capacity_) {
            return;
        }
 
        // If we don't have an onEject callback, fast remove the last element
        if (not BaseClass::on_eject_) {
            while (lookup_.size() > capacity_) {
                auto& [key, _] = items_.back();
                lookup_.erase(key);
                items_.pop_back();
            }
            return;
        }
 
        // Else iterate the items in reverse until we find one we can remove
        for (auto it = items_.rbegin(); it != items_.rend();) {
            // Get the key and value
            auto& [key, value] = *it;
 
            // Eject this item
            if (BaseClass::on_eject_(key, value)) {
                lookup_.erase(key);
                it = std::next(it);
                it = std::reverse_iterator(items_.erase(it.base()));
            } else {
                ++it;
            }
 
            // Stop when we've reset to capacity
            if (lookup_.size() <= capacity_) {
                break;
            }
        }
    }
 
    void purge() override
    {
        std::unique_lock lock(cache_mutex_);
        // If we have an on_eject_ function, check every item
        if (BaseClass::on_eject_) {
            for (auto it = items_.begin(); it != items_.end();) {
                auto& [key, value] = *it;
                if (BaseClass::on_eject_(key, value)) {
                    lookup_.erase(key);
                    it = items_.erase(it);
                } else {
                    ++it;
                }
            }
        }
 
        // Otherwise, remove everything
        else {
            lookup_.clear();
            items_.clear();
        }
    }
private:
    std::list<TPair> items_;
    std::unordered_map<TKey, TListIterator> lookup_;
    mutable TMutex cache_mutex_;
};
}  // namespace volcart