pht.h

// Copyright (c) 2014-2026 Savoir-faire Linux Inc.
// SPDX-License-Identifier: MIT
#pragma once
 
#include <string>
#include <vector>
#include <memory>
#include <map>
#include <functional>
#include <stdexcept>
#include <bitset>
#include <iostream>
#include <sstream>
 
#include "../value.h"
#include "../dhtrunner.h"
#include "../infohash.h"
 
namespace dht {
namespace indexation {

struct OPENDHT_PUBLIC Prefix
{
    Prefix() {}
    Prefix(InfoHash h)
        : size_(h.size() * 8)
        , content_(h.begin(), h.end())
    {}
    Prefix(const Blob& d, const Blob& f = {})
        : size_(d.size() * 8)
        , flags_(f)
        , content_(d)
    {}
 
    Prefix(const Prefix& p, size_t first)
        : size_(std::min(first, p.content_.size() * 8))
        , content_(Blob(p.content_.begin(), p.content_.begin() + size_ / 8))
    {
        auto rem = size_ % 8;
        if (not p.flags_.empty()) {
            flags_ = Blob(p.flags_.begin(), p.flags_.begin() + size_ / 8);
            if (rem)
                flags_.push_back(p.flags_[size_ / 8] & (0xFF << (8 - rem)));
        }
 
        if (rem)
            content_.push_back(p.content_[size_ / 8] & (0xFF << (8 - rem)));
    }

    Prefix getPrefix(ssize_t len) const
    {
        if ((size_t) std::abs(len) >= content_.size() * 8)
            throw std::out_of_range("len larger than prefix size.");
        if (len < 0)
            len += size_;
 
        return Prefix(*this, len);
    }

    bool isFlagActive(size_t pos) const { return flags_.empty() or isActiveBit(flags_, pos); }

    bool isContentBitActive(size_t pos) const { return isActiveBit(content_, pos); }
 
    Prefix getFullSize() { return Prefix(*this, content_.size() * 8); }

    Prefix getSibling() const
    {
        Prefix copy = *this;
        if (size_)
            copy.swapContentBit(size_ - 1);
 
        return copy;
    }
 
    InfoHash hash() const
    {
        Blob copy(content_);
        copy.push_back(size_);
        return InfoHash::get(copy);
    }

    static inline unsigned commonBits(const Prefix& p1, const Prefix& p2)
    {
        unsigned i, j;
        uint8_t x;
        auto longest_prefix_size = std::min(p1.size_, p2.size_);
 
        for (i = 0; i < longest_prefix_size; i++) {
            if (p1.content_.data()[i] != p2.content_.data()[i] or not p1.isFlagActive(i) or not p2.isFlagActive(i)) {
                break;
            }
        }
 
        if (i == longest_prefix_size)
            return 8 * longest_prefix_size;
 
        x = p1.content_.data()[i] ^ p2.content_.data()[i];
 
        j = 0;
        while ((x & 0x80) == 0) {
            x <<= 1;
            j++;
        }
 
        return 8 * i + j;
    }

    void swapContentBit(size_t bit) { swapBit(content_, bit); }

    void swapFlagBit(size_t bit) { swapBit(flags_, bit); }

    void addPaddingContent(size_t size) { content_ = addPadding(content_, size); }
 
    void updateFlags()
    {
        /* Fill first known bit */
        auto csize = size_ - flags_.size() * 8;
        while (csize >= 8) {
            flags_.push_back(0xFF);
            csize -= 8;
        }
 
        /* if needed fill remaining bit */
        if (csize)
            flags_.push_back(0xFF << (8 - csize));
 
        /* Complet vector space missing */
        for (auto i = flags_.size(); i < content_.size(); i++)
            flags_.push_back(0xFF);
    }
 
    std::string toString() const;
 
    size_t size_ {0};
 
    /* Will contain flags according to content_.
       If flags_[i] == 0, then content_[i] is unknown
       else if flags_[i] == 1, then content_[i] is known */
    Blob flags_ {};
    Blob content_ {};
 
private:
    Blob addPadding(Blob toP, size_t size)
    {
        Blob copy = toP;
        for (auto i = copy.size(); i < size; i++)
            copy.push_back(0);
 
        swapBit(copy, size_ + 1);
        return copy;
    }

    bool isActiveBit(const Blob& b, size_t pos) const
    {
        if (pos >= content_.size() * 8)
            throw std::out_of_range("Can't detect active bit at pos, pos larger than prefix size or empty prefix");
 
        return ((b[pos / 8] >> (7 - (pos % 8))) & 1) == 1;
    }

    void swapBit(Blob& b, size_t bit)
    {
        if (bit >= b.size() * 8)
            throw std::out_of_range("bit larger than prefix size.");
 
        size_t offset_bit = (8 - bit) % 8;
        b[bit / 8] ^= (1 << offset_bit);
    }
};
 
using Value = std::pair<InfoHash, dht::Value::Id>;
struct OPENDHT_PUBLIC IndexEntry : public dht::Value::Serializable<IndexEntry>
{
    static const ValueType& TYPE;
 
    virtual void unpackValue(const dht::Value& v)
    {
        Serializable<IndexEntry>::unpackValue(v);
        name = v.user_type;
    }
 
    virtual dht::Value packValue() const
    {
        auto pack = Serializable<IndexEntry>::packValue();
        pack.user_type = name;
        return pack;
    }
 
    Blob prefix;
    Value value;
    std::string name;
    MSGPACK_DEFINE_MAP(prefix, value)
};
 
class OPENDHT_PUBLIC Pht
{
    static constexpr const char* INVALID_KEY = "Key does not match the PHT key spec.";
 
    /* Prefixes the user_type for all dht values put on the DHT */
    static constexpr const char* INDEX_PREFIX = "index.pht.";
 
public:
    /* This is the maximum number of entries per node. This parameter is
     * critical and influences the traffic a lot during a lookup operation.
     */
    static constexpr const size_t MAX_NODE_ENTRY_COUNT {16};
 
    /* A key for a an index entry */
    using Key = std::map<std::string, Blob>;
 
    /* Specifications of the keys. It defines the number, the length and the
     * serialization order of fields. */
    using KeySpec = std::map<std::string, size_t>;
    using LookupCallback = std::function<void(std::vector<std::shared_ptr<Value>>& values, const Prefix& p)>;
 
    typedef void (*LookupCallbackRaw)(std::vector<std::shared_ptr<Value>>* values, Prefix* p, void* user_data);
    static LookupCallback bindLookupCb(LookupCallbackRaw raw_cb, void* user_data)
    {
        if (not raw_cb)
            return {};
        return [=](std::vector<std::shared_ptr<Value>>& values, const Prefix& p) {
            raw_cb((std::vector<std::shared_ptr<Value>>*) &values, (Prefix*) &p, user_data);
        };
    }
    using LookupCallbackSimple = std::function<void(std::vector<std::shared_ptr<Value>>& values)>;
    typedef void (*LookupCallbackSimpleRaw)(std::vector<std::shared_ptr<Value>>* values, void* user_data);
    static LookupCallbackSimple bindLookupCbSimple(LookupCallbackSimpleRaw raw_cb, void* user_data)
    {
        if (not raw_cb)
            return {};
        return [=](std::vector<std::shared_ptr<Value>>& values) {
            raw_cb((std::vector<std::shared_ptr<Value>>*) &values, user_data);
        };
    }
 
    Pht(std::string name, KeySpec k_spec, std::shared_ptr<DhtRunner> dht)
        : name_(INDEX_PREFIX + name)
        , canary_(name_ + ".canary")
        , keySpec_(k_spec)
        , dht_(dht)
    {}
 
    virtual ~Pht() {}

    void lookup(Key k, LookupCallback cb = {}, DoneCallbackSimple done_cb = {}, bool exact_match = true);
    void lookup(Key k, LookupCallbackSimple cb = {}, DoneCallbackSimple done_cb = {}, bool exact_match = true)
    {
        lookup(
            k,
            [cb = std::move(cb)](std::vector<std::shared_ptr<Value>>& values, Prefix) { cb(values); },
            done_cb,
            exact_match);
    }

    void insert(Key k, Value v, DoneCallbackSimple done_cb = {})
    {
        Prefix p = linearize(k);
 
        auto lo = std::make_shared<int>(0);
        auto hi = std::make_shared<int>(p.size_);
 
        IndexEntry entry;
        entry.value = v;
        entry.prefix = p.content_;
        entry.name = name_;
 
        Pht::insert(p, entry, lo, hi, clock::now(), true, done_cb);
    }
 
private:
 
    void insert(const Prefix& kp,
                IndexEntry entry,
                std::shared_ptr<int> lo,
                std::shared_ptr<int> hi,
                time_point time_p,
                bool check_split,
                DoneCallbackSimple done_cb = {});
 
    class Cache
    {
    public:
        void insert(const Prefix& p);

 
        int lookup(const Prefix& p);
 
    private:
        static constexpr const size_t MAX_ELEMENT {1024};
        static constexpr const std::chrono::minutes NODE_EXPIRE_TIME {5};
 
        struct Node
        {
            time_point last_reply;           /* Made the assocation between leaves and leaves multimap */
            std::shared_ptr<Node> parent;    /* Share_ptr to the parent, it allow the self destruction of tree */
            std::weak_ptr<Node> left_child;  /* Left child, for bit equal to 1 */
            std::weak_ptr<Node> right_child; /* Right child, for bit equal to 0 */
        };
 
        std::weak_ptr<Node> root_; /* Root of the tree */

        std::multimap<time_point, std::shared_ptr<Node>> leaves_;
    };
 
    /* Callback used for insert value by using the pht */
    using RealInsertCallback = std::function<void(const Prefix& p, IndexEntry entry)>;
    using LookupCallbackWrapper = std::function<void(std::vector<std::shared_ptr<IndexEntry>>& values, const Prefix& p)>;

    void lookupStep(Prefix k,
                    std::shared_ptr<int> lo,
                    std::shared_ptr<int> hi,
                    std::shared_ptr<std::vector<std::shared_ptr<IndexEntry>>> vals,
                    LookupCallbackWrapper cb,
                    DoneCallbackSimple done_cb,
                    std::shared_ptr<unsigned> max_common_prefix_len,
                    int start = -1,
                    bool all_values = false);

    Prefix zcurve(const std::vector<Prefix>& all_prefix) const;

    virtual Prefix linearize(Key k) const;

    void getRealPrefix(const std::shared_ptr<Prefix>& p, IndexEntry entry, RealInsertCallback end_cb);

    void checkPhtUpdate(Prefix p, IndexEntry entry, time_point time_p);

    static size_t findSplitLocation(const Prefix& compared, const std::vector<std::shared_ptr<IndexEntry>>& vals)
    {
        for (size_t i = 0; i < compared.content_.size() * 8 - 1; i++)
            for (auto const& v : vals)
                if (Prefix(v->prefix).isContentBitActive(i) != compared.isContentBitActive(i))
                    return i + 1;
        return compared.content_.size() * 8 - 1;
    }

    void split(const Prefix& insert,
               const std::vector<std::shared_ptr<IndexEntry>>& vals,
               IndexEntry entry,
               RealInsertCallback end_cb);

    bool validKey(const Key& k) const
    {
        return k.size() == keySpec_.size()
               and std::equal(k.begin(),
                              k.end(),
                              keySpec_.begin(),
                              [&](const Key::value_type& key, const KeySpec::value_type& key_spec) {
                                  return key.first == key_spec.first and key.second.size() <= key_spec.second;
                              });
    }

    void updateCanary(Prefix p);
 
    const std::string name_;
    const std::string canary_;
    const KeySpec keySpec_;
    Cache cache_;
    std::shared_ptr<DhtRunner> dht_;
};
 
} // namespace indexation
} // namespace dht