lokinet/llarp/dht/bucket.hpp

#pragma once

#include "kademlia.hpp"
#include "key.hpp"

#include <llarp/util/status.hpp>

#include <map>
#include <set>
#include <vector>

namespace llarp::dht
{
  template <typename Val_t>
  struct Bucket
  {
    using BucketStorage_t = std::map<Key_t, Val_t, XorMetric>;
    using Random_t = std::function<uint64_t()>;

    Bucket(const Key_t& us, Random_t r) : nodes(XorMetric(us)), random(std::move(r))
    {}

    util::StatusObject
    ExtractStatus() const
    {
      util::StatusObject obj{};
      for (const auto& item : nodes)
      {
        obj[item.first.ToString()] = item.second.ExtractStatus();
      }
      return obj;
    }

    size_t
    size() const
    {
      return nodes.size();
    }

    struct SetIntersector
    {
      bool
      operator()(const typename BucketStorage_t::value_type& lhs, const Key_t& rhs)
      {
        return lhs.first < rhs;
      }

      bool
      operator()(const Key_t& lhs, const typename BucketStorage_t::value_type& rhs)
      {
        return lhs < rhs.first;
      }
    };

    bool
    GetRandomNodeExcluding(Key_t& result, const std::set<Key_t>& exclude) const
    {
      std::vector<typename BucketStorage_t::value_type> candidates;
      std::set_difference(
          nodes.begin(),
          nodes.end(),
          exclude.begin(),
          exclude.end(),
          std::back_inserter(candidates),
          SetIntersector());

      if (candidates.empty())
      {
        return false;
      }
      result = candidates[random() % candidates.size()].first;
      return true;
    }

    bool
    FindClosest(const Key_t& target, Key_t& result) const
    {
      Key_t mindist;
      mindist.Fill(0xff);
      for (const auto& item : nodes)
      {
        auto curDist = item.first ^ target;
        if (curDist < mindist)
        {
          mindist = curDist;
          result = item.first;
        }
      }
      return nodes.size() > 0;
    }

    bool
    GetManyRandom(std::set<Key_t>& result, size_t N) const
    {
      if (nodes.size() < N || nodes.empty())
      {
        llarp::LogWarn("Not enough dht nodes, have ", nodes.size(), " want ", N);
        return false;
      }
      if (nodes.size() == N)
      {
        std::transform(
            nodes.begin(), nodes.end(), std::inserter(result, result.end()), [](const auto& a) {
              return a.first;
            });

        return true;
      }
      size_t expecting = N;
      size_t sz = nodes.size();
      while (N)
      {
        auto itr = nodes.begin();
        std::advance(itr, random() % sz);
        if (result.insert(itr->first).second)
        {
          --N;
        }
      }
      return result.size() == expecting;
    }

    bool
    FindCloseExcluding(const Key_t& target, Key_t& result, const std::set<Key_t>& exclude) const
    {
      Key_t maxdist;
      maxdist.Fill(0xff);
      Key_t mindist;
      mindist.Fill(0xff);
      for (const auto& item : nodes)
      {
        if (exclude.count(item.first))
        {
          continue;
        }

        auto curDist = item.first ^ target;
        if (curDist < mindist)
        {
          mindist = curDist;
          result = item.first;
        }
      }
      return mindist < maxdist;
    }

    bool
    GetManyNearExcluding(
        const Key_t& target,
        std::set<Key_t>& result,
        size_t N,
        const std::set<Key_t>& exclude) const
    {
      std::set<Key_t> s(exclude.begin(), exclude.end());

      Key_t peer;
      while (N--)
      {
        if (!FindCloseExcluding(target, peer, s))
        {
          return false;
        }
        s.insert(peer);
        result.insert(peer);
      }
      return true;
    }

    void
    PutNode(const Val_t& val)
    {
      auto itr = nodes.find(val.ID);
      if (itr == nodes.end() || itr->second < val)
      {
        nodes[val.ID] = val;
      }
    }

    void
    DelNode(const Key_t& key)
    {
      auto itr = nodes.find(key);
      if (itr != nodes.end())
      {
        nodes.erase(itr);
      }
    }

    bool
    HasNode(const Key_t& key) const
    {
      return nodes.find(key) != nodes.end();
    }

    // remove all nodes who's key matches a predicate
    template <typename Predicate>
    void
    RemoveIf(Predicate pred)
    {
      auto itr = nodes.begin();
      while (itr != nodes.end())
      {
        if (pred(itr->first))
          itr = nodes.erase(itr);
        else
          ++itr;
      }
    }

    template <typename Visit_t>
    void
    ForEachNode(Visit_t visit)
    {
      for (const auto& item : nodes)
      {
        visit(item.second);
      }
    }

    void
    Clear()
    {
      nodes.clear();
    }

    BucketStorage_t nodes;
    Random_t random;
  };
}  // namespace llarp::dht