hyperswarm-peer / src / kademlia_table.h

#ifndef HYPERSWARM_KADEMLIA_TABLE_H
#define HYPERSWARM_KADEMLIA_TABLE_H

// Kademlia routing table — compatible with holepunchto/kademlia-routing-table.
// k=20, 32-byte node IDs, XOR distance, lazy row allocation.
// Rows maintain nodes sorted by ID for binary search.

#include "compact_encoding.h"

#include <cstring>

struct KademliaNode {
	NodeId id;
	IPv4Address address;
	uint64_t added_at; // Timestamp in ms when the node was added/updated.
};

// Fixed-size row with binary search. Max K nodes per row.
struct KademliaRow {
	static constexpr int K = 20;

	KademliaNode nodes[K];
	int count;
	int byte_offset; // Row index >> 3. Compare starts from this byte.

	KademliaRow() : count(0), byte_offset(0) {}

	// Compare two node IDs starting from byte_offset.
	// Returns <0, 0, or >0 (like memcmp).
	int compare(const uint8_t *a, const uint8_t *b) const {
		for (int i = byte_offset; i < 32; i++) {
			if (a[i] != b[i]) return a[i] < b[i] ? -1 : 1;
		}
		return 0;
	}

	// Binary search for a node ID. Returns index if found, or -1.
	// Sets insert_pos to the correct insertion point.
	int find(const uint8_t *id, int &insert_pos) const {
		int l = 0, r = count - 1;
		while (l <= r) {
			int m = (l + r) >> 1;
			int c = compare(id, nodes[m].id.bytes);
			if (c == 0) { insert_pos = m; return m; }
			if (c < 0) r = m - 1;
			else l = m + 1;
		}
		insert_pos = l;
		return -1;
	}

	// Add or update a node. Returns true if added (new or updated).
	bool add(const KademliaNode &node) {
		int pos;
		int idx = find(node.id.bytes, pos);

		if (idx >= 0) {
			// Update existing node.
			nodes[idx] = node;
			return true;
		}

		if (count >= K) return false; // Row full.

		// Insert at pos, shift right.
		for (int j = count; j > pos; j--) {
			nodes[j] = nodes[j - 1];
		}
		nodes[pos] = node;
		count++;
		return true;
	}

	// Remove a node by ID. Returns true if found and removed.
	bool remove(const uint8_t *id) {
		int pos;
		int idx = find(id, pos);
		if (idx < 0) return false;

		// Shift left.
		for (int j = idx; j < count - 1; j++) {
			nodes[j] = nodes[j + 1];
		}
		count--;
		return true;
	}

	// Get node by ID. Returns pointer or nullptr.
	KademliaNode *get(const uint8_t *id) {
		int pos;
		int idx = find(id, pos);
		return idx >= 0 ? &nodes[idx] : nullptr;
	}
};

class KademliaTable {
public:
	static constexpr int K = KademliaRow::K;
	static constexpr int NUM_ROWS = 256; // 32 bytes * 8 bits

	KademliaTable() : _size(0) {
		memset(_rows, 0, sizeof(_rows));
	}

	~KademliaTable() {
		for (int i = 0; i < NUM_ROWS; i++) {
			delete _rows[i];
		}
	}

	void init(const NodeId &self_id) {
		memcpy(_id.bytes, self_id.bytes, 32);
	}

	const NodeId &id() const { return _id; }
	int size() const { return _size; }

	// Add or update a node. Returns true if the table changed.
	bool add(const KademliaNode &node) {
		int i = _diff(node.id.bytes);

		KademliaRow *row = _rows[i];
		if (row == nullptr) {
			row = new KademliaRow();
			row->byte_offset = i >> 3;
			_rows[i] = row;
		}

		int old_count = row->count;
		if (!row->add(node)) return false;
		_size += row->count - old_count;
		return true;
	}

	// Remove a node by ID. Returns true if found and removed.
	bool remove(const uint8_t *id) {
		int i = _diff(id);
		KademliaRow *row = _rows[i];
		if (row == nullptr) return false;
		if (!row->remove(id)) return false;
		_size--;
		return true;
	}

	// Get a node by ID. Returns pointer or nullptr.
	KademliaNode *get(const uint8_t *id) {
		int i = _diff(id);
		KademliaRow *row = _rows[i];
		if (row == nullptr) return nullptr;
		return row->get(id);
	}

	bool has(const uint8_t *id) {
		return get(id) != nullptr;
	}

	// Find the k closest nodes to a target ID.
	// Returns the number of nodes written to out (up to max_k).
	int closest(const uint8_t *target, KademliaNode *out, int max_k = K) const {
		int d = _diff(target);
		int result_count = 0;

		// Fan out from the matching row: down first, then up.
		for (int i = d; i >= 0 && result_count < max_k; i--) {
			_push_nodes(i, max_k, out, result_count);
		}
		for (int i = d + 1; i < NUM_ROWS && result_count < max_k; i++) {
			_push_nodes(i, max_k, out, result_count);
		}

		return result_count;
	}

	// Get the closest nodes as IPv4Address array (for DHT responses).
	int closest_addresses(const uint8_t *target, IPv4Address *out, int max_k = K) const {
		KademliaNode nodes[K];
		int n = closest(target, nodes, max_k);
		for (int i = 0; i < n; i++) {
			out[i] = nodes[i].address;
		}
		return n;
	}

private:
	// Compute the XOR distance bit position.
	// Matches JS: Math.clz32(a ^ b) - 24 for the differing byte.
	int _diff(const uint8_t *id) const {
		for (int i = 0; i < 32; i++) {
			uint8_t a = id[i];
			uint8_t b = _id.bytes[i];
			if (a != b) {
				uint8_t diff = a ^ b;
				int clz = 0;
				if (diff < 0x10) { clz += 4; diff <<= 4; }
				if (diff < 0x40) { clz += 2; diff <<= 2; }
				if (diff < 0x80) { clz += 1; }
				return i * 8 + clz;
			}
		}
		// IDs are equal — place in the last row.
		return NUM_ROWS - 1;
	}

	void _push_nodes(int i, int max_k, KademliaNode *out, int &count) const {
		const KademliaRow *row = _rows[i];
		if (row == nullptr) return;
		int missing = max_k - count;
		if (missing > row->count) missing = row->count;
		for (int j = 0; j < missing; j++) {
			out[count++] = row->nodes[j];
		}
	}

	NodeId _id;
	int _size;
	KademliaRow *_rows[NUM_ROWS];
};

#endif // HYPERSWARM_KADEMLIA_TABLE_H