Optimize IpToPeerMap with a trie

ip_to_peer_map.cpp

@@ -4,6 +4,7 @@
 #include "ip_to_peer_map.h"
 #include "bit_ops.h"
 #include <string.h>
+#include <assert.h>
 
 IpToPeerMap::IpToPeerMap() {
 
@@ -12,10 +13,8 @@ IpToPeerMap::IpToPeerMap() {
 IpToPeerMap::~IpToPeerMap() {
 }
 
-bool IpToPeerMap::InsertV4(const void *addr, int cidr, void *peer) {
+bool IpToPeerMap::InsertV4(uint32 ip, int cidr, void *peer) {
-  uint32 mask = cidr == 32 ? 0xffffffff : ~(0xffffffff >> cidr);
+  ipv4_.Insert(ip, cidr, peer);
-  Entry4 e = {ReadBE32(addr) & mask, mask, peer};
-  ipv4_.push_back(e);
   return true;
 }
 
@@ -29,23 +28,11 @@ bool IpToPeerMap::InsertV6(const void *addr, int cidr, void *peer) {
 }
 
 void *IpToPeerMap::LookupV4(uint32 ip) {
-  uint32 best_mask = 0;
+  return ipv4_.Lookup(ip);
-  void *best_peer = NULL;
-  for (auto it = ipv4_.begin(); it != ipv4_.end(); ++it) {
-    if (it->ip == (ip & it->mask) && it->mask >= best_mask) {
-      best_mask = it->mask;
-      best_peer = it->peer;
-    }
-  }
-  return best_peer;
 }
 
 void *IpToPeerMap::LookupV4DefaultPeer() {
-  for (auto it = ipv4_.begin(); it != ipv4_.end(); ++it) {
+  return ipv4_.LookupExact(0, 0);
-    if (it->mask == 0)
-      return it->peer;
-  }
-  return NULL;
 }
 
 void *IpToPeerMap::LookupV6DefaultPeer() {
@@ -76,15 +63,8 @@ void *IpToPeerMap::LookupV6(const void *addr) {
 }
 
 void IpToPeerMap::RemovePeer(void *peer) {
-  {
+  assert(0);
-    size_t n = ipv4_.size();
+  // todo: remove peer also from ipv4_
-    Entry4 *r = &ipv4_[0], *w = r;
-    for (size_t i = 0; i != n; i++, r++) {
-      if (r->peer != peer)
-        *w++ = *r;
-    }
-    ipv4_.resize(w - &ipv4_[0]);
-  }
   {
     size_t n = ipv6_.size();
     Entry6 *r = &ipv6_[0], *w = r;
@@ -95,3 +75,678 @@ void IpToPeerMap::RemovePeer(void *peer) {
     ipv6_.resize(w - &ipv6_[0]);
   }
 }
+
+#pragma warning (disable: 4200)  // warning C4200: nonstandard extension used: zero-sized array in struct/union
+struct RoutingTrie32::Node {
+  uint32 key;
+  // bits == 0 if this is a leaf
+  uint8 pos, bits;
+  union {
+    struct {
+      Node *parent;
+      uint32 full_children, empty_children;
+      Node *child[0];
+    };
+    struct {
+      Node *leaf_next;
+      Value leaf_value;
+    };
+  };
+};
+
+
+static inline uint32 prefix_mismatch(uint32 key, RoutingTrie32::Node *n) {
+  uint32 prefix = n->key;
+  return (key ^ prefix) & (prefix | (uint32)-(int32)prefix);
+}
+
+static uint32 make_cidr_mask(uint8 cidr) {
+  return cidr == 0 ? 0 : 0xffffffff << (32 - cidr);
+}
+
+#define IS_LEAF(n) (n->bits == 0)
+#define GET_INDEX(k, n) ((n->key ^ k) >> n->pos)
+
+#define NODE_IS_OLEAF(n) ((intptr_t)(n) & 1)
+#define NODE_IS_NULL_OR_OLEAF(n) ((n) == 0 || NODE_IS_OLEAF(n))
+#define VALUE_TO_OLEAF(n) ((Node*)((intptr_t)(n) + 1))
+#define VALUE_FROM_OLEAF(n) ((void*)((intptr_t)(n) - 1))
+
+static RoutingTrie32::Node *NewNode(uint32 key, int pos, int bits) {
+  RoutingTrie32::Node *n = (RoutingTrie32::Node *)malloc(offsetof(RoutingTrie32::Node, child[(uint32)(1U << bits)]));
+  if (n) {
+    n->parent = NULL;
+    n->pos = pos;
+    n->bits = bits;
+    n->full_children = 0;
+    n->empty_children = 1U << bits;
+    uint32 s = pos + bits;
+    n->key = (s < 32) ? key >> s << s : 0;
+    memset(n->child, 0, n->empty_children * sizeof(RoutingTrie32::Node*));
+  }
+  return n;
+}
+
+static RoutingTrie32::Node *NewLeaf(uint32 key, uint8 leaf_pos, RoutingTrie32::Value value) {
+  RoutingTrie32::Node *n = (RoutingTrie32::Node *)malloc(sizeof(RoutingTrie32::Node));
+  if (n) {
+    n->key = key;
+    n->bits = 0;
+    n->pos = leaf_pos;
+    n->leaf_value = value;
+    n->leaf_next = NULL;
+  }
+  return n;
+}
+
+static void FreeNode(RoutingTrie32::Node *n) {
+  free(n);
+}
+
+static void RecursiveFreeNode(RoutingTrie32::Node *n) {
+  RoutingTrie32::Node *cn;
+
+  if (n->bits == 0) {
+    while ((cn = n->leaf_next) != NULL) {
+      n->leaf_next = cn->leaf_next;
+      FreeNode(cn);
+    }
+  } else {
+    uint32 items = 1 << n->bits;
+    for (uint32 i = 0; i != items; i++) {
+      RoutingTrie32::Node *cn = n->child[i];
+      if (!NODE_IS_NULL_OR_OLEAF(cn))
+        RecursiveFreeNode(cn);
+    }
+  }
+  FreeNode(n);
+}
+
+
+RoutingTrie32::RoutingTrie32()
+  : root_(NULL) {
+}
+
+RoutingTrie32::~RoutingTrie32() {
+  if (root_)
+    RecursiveFreeNode(root_);
+}
+
+RoutingTrie32::Value RoutingTrie32::Lookup(uint32 ip) {
+  uint32 key = ip;
+  Node *n = root_, *pn = n, *ppn;
+  int cindex = 0;
+  if (!n)
+    return NULL;
+  // Find the longest prefix match
+  for (;;) {
+    uint32 index = GET_INDEX(key, n);
+    if (index >> n->bits)
+      break; // mismatch in skipped bits
+    if (IS_LEAF(n))
+      return n->leaf_value;
+    pn = n;
+    cindex = index;
+    n = n->child[index];
+    if (NODE_IS_NULL_OR_OLEAF(n)) {
+      if (!n)
+        goto backtrace;
+      // node is an optimized leaf
+      return VALUE_FROM_OLEAF(n);
+    }
+  }
+  // backtrace for longest prefix
+  for (;;) {
+    if (prefix_mismatch(key, n))
+      goto backtrace;
+    if (IS_LEAF(n)) {
+      for (;;) {
+        if (((n->key ^ key) >> n->pos) == 0)
+          return n->leaf_value;
+        if (n->leaf_next == NULL) {
+          if (n->pos == 32)
+            return n->leaf_value;
+          break;
+        }
+        n = n->leaf_next;
+      }
+      goto backtrace;
+    }
+
+    ppn = n;
+    n = n->child[0];
+    if (NODE_IS_NULL_OR_OLEAF(n)) {
+      if (n) {
+        if (((ppn->key ^ key) >> ppn->pos) == 0)
+          return VALUE_FROM_OLEAF(n);
+      }
+      for (;;) {
+backtrace:
+        // step up to previous parent when we used all bits in current
+        while (cindex == 0) {
+          uint32 pkey = pn->key;
+          pn = pn->parent;
+          if (!pn)
+            return 0;
+          cindex = (pn->key ^ pkey) >> pn->pos;
+        }
+        // strip lsb of cindex and find child
+        cindex &= cindex - 1;
+        assert(cindex < (1 << pn->bits));
+        n = pn->child[cindex];
+        if (!NODE_IS_NULL_OR_OLEAF(n))
+          break;
+        if (n) {
+          uint32 nkey = pn->key + (cindex << pn->pos);
+          if (((nkey ^ key) >> pn->pos) == 0)
+            return VALUE_FROM_OLEAF(n);
+        }
+      }
+    }
+  }
+}
+
+bool RoutingTrie32::InsertLeafInto(Node **nn, uint8 leaf_pos, Value value) {
+  // put higher cidr higher up
+  Node *n = *nn;
+  assert(IS_LEAF(n));
+  uint32 key = n->key;
+  do {
+    if (leaf_pos < n->pos)
+      break;
+    if (leaf_pos == n->pos) {
+      n->leaf_value = value;
+      return true;
+    }
+    nn = &n->leaf_next;
+  } while ((n = *nn) != NULL);
+  Node *leaf = NewLeaf(key, leaf_pos, value);
+  if (leaf == NULL)
+    return false;
+  leaf->leaf_next = *nn;
+  *nn = leaf;
+  return true;
+}
+
+static inline bool IsFull(RoutingTrie32::Node *pn, RoutingTrie32::Node *n) {
+  return !NODE_IS_NULL_OR_OLEAF(n) && (n->pos + n->bits) == pn->pos && !IS_LEAF(n);
+}
+
+void RoutingTrie32::PutChild(Node *pn, uint32 i, Node *n) {
+  Node *on = pn->child[i];
+  pn->child[i] = n;
+
+  pn->empty_children += (n == NULL) - (on == NULL);
+  pn->full_children += IsFull(pn, n) - IsFull(pn, on);
+
+  assert(pn->empty_children < 0x80000000);
+  assert(pn->full_children < 0x80000000);
+}
+
+bool RoutingTrie32::Insert(uint32 ip, int cidr, Value value) {
+  uint32 key = ip;
+  Node **nn = &root_, *n = root_, *pn = NULL, *leaf, *tn = NULL, *leaf_to_free = NULL;
+  uint8 leaf_pos = 32 - cidr;
+
+  if (n == NULL) {
+    root_ = NewLeaf(key, leaf_pos, value);
+    return (root_ != NULL);
+  }
+  assert(!NODE_IS_OLEAF(n));
+
+  for (;;) {
+    uint32 index = GET_INDEX(key, n);
+    if (index >> n->bits) {
+force_add:
+      // n is a node and the key doesn't match, allocate a new node
+      // with two elements and insert it.
+      if (!(tn = NewNode(key, FindLastSetBit32(key ^ n->key), 1)))
+        return false;
+      tn->parent = pn;
+      // can convert leaf node to oleaf?
+      if (IS_LEAF(n)) {
+        if (tn->pos == n->pos && n->leaf_next == NULL) {
+          leaf_to_free = n;
+          n = VALUE_TO_OLEAF(n->leaf_value);
+        }
+      }
+      PutChild(tn, GET_INDEX(key, tn) ^ 1, n);
+      break;
+    }
+    if (IS_LEAF(n)) {
+      if (key != n->key)
+        goto force_add;
+      return InsertLeafInto(nn, leaf_pos, value);
+    }
+    pn = n;
+    nn = &n->child[index];
+    if ((n = *nn) == NULL) {
+      tn = pn;
+      break;
+    }
+    if (NODE_IS_OLEAF(n)) {
+      if (!(n = NewLeaf(pn->key + (index << pn->pos), pn->pos, VALUE_FROM_OLEAF(n))))
+        return false;
+      *nn = n;
+    }
+  }
+  // Create either leaf or oleaf
+  if (tn->pos == leaf_pos) {
+    leaf = VALUE_TO_OLEAF(value);
+  } else if (!(leaf = NewLeaf(key, leaf_pos, value))) {
+    if (tn != pn)
+      FreeNode(tn);
+    return false;
+  }
+
+  // -- Start making irreversible changes here
+  if (leaf_to_free)
+    FreeNode(leaf_to_free);
+
+  if (tn != pn) {
+    if (!NODE_IS_OLEAF(n) && !IS_LEAF(n))
+      n->parent = tn;
+
+    if (pn) {
+      PutChild(pn, GET_INDEX(key, pn), tn);
+    } else {
+      root_ = tn;
+    }
+    pn = tn;
+  }
+
+  PutChild(pn, GET_INDEX(key, pn), leaf);
+
+  Rebalance(pn);
+  return true;
+}
+
+bool RoutingTrie32::Delete(uint32 ip, int cidr) {
+  uint32 key = ip;
+
+  Node *n = root_, *pn = NULL;
+  uint32 pn_index = 0;
+
+  if (n == NULL)
+    return false;
+
+  uint8 leaf_pos = 32 - cidr;
+
+  for (;;) {
+    uint32 index = GET_INDEX(key, n);
+    if (index >> n->bits)
+      return false;
+    if (IS_LEAF(n)) {
+      if (n->key != key)
+        return false;
+      if (n->pos == leaf_pos) {
+        if (pn == NULL) {
+          root_ = n->leaf_next;
+        } else {
+          PutChild(pn, pn_index, n->leaf_next);
+          if (n->leaf_next == NULL)
+            Rebalance(pn);
+        }
+        FreeNode(n);
+        return true;
+      }
+      Node **nn = &n->leaf_next;
+      while (*nn) {
+        if ((*nn)->pos == leaf_pos) {
+          *nn = (*nn)->leaf_next;
+          FreeNode(*nn);
+          return true;
+        }
+        nn = &(*nn)->leaf_next;
+      }
+      return false;
+    }
+    pn = n;
+    pn_index = index;
+    n = n->child[index];
+    if (NODE_IS_NULL_OR_OLEAF(n)) {
+      if (n && key == pn->key + (index << pn->pos) && pn->pos == leaf_pos) {
+        PutChild(pn, index, NULL);
+        Rebalance(pn);
+        return true;
+      }
+      return false;
+    }
+  }
+}
+
+RoutingTrie32::Value RoutingTrie32::LookupExact(uint32 ip, int cidr) {
+  uint32 key = ip;
+  Node *n = root_, *pn;
+  if (n == NULL)
+    return NULL;
+  uint8 leaf_pos = 32 - cidr;
+  for (;;) {
+    uint32 index = GET_INDEX(key, n);
+    if (index >> n->bits)
+      return NULL;
+    if (IS_LEAF(n)) {
+      if (n->key != key)
+        return NULL;
+      do {
+        if (n->pos == leaf_pos)
+          return n->leaf_value;
+        n = n->leaf_next;
+      } while (n);
+      return NULL;
+    }
+    pn = n;
+    n = n->child[index];
+    if (NODE_IS_NULL_OR_OLEAF(n))
+      return (n && key == pn->key + (index << pn->pos) && pn->pos == leaf_pos) ? n->leaf_value : NULL;
+  }
+}
+
+void RoutingTrie32::Rebalance(Node *n) {
+  // Always resize |n| and its parent.  For each parent where
+  // Resize returns true resize also its parent.
+  Node *np = n->parent;
+  Resize(n);
+  while (np) {
+    n = np;
+    np = n->parent;
+    if (!Resize(n))
+      break;
+  }
+}
+
+void RoutingTrie32::ResizeChildren(Node *pn) {
+  for (uint32 i = 0, i_end = 1U << pn->bits; i != i_end; i++) {
+    Node *n = pn->child[i];
+    if (IsFull(pn, n))
+      Resize(n);
+  }
+}
+
+enum {
+  kHalveThreshold = 25,
+  kInflateThreshold = 50,
+  kHalveThresholdRoot = 15,
+  kInflateThresholdRoot = 30,
+};
+
+bool RoutingTrie32::Resize(Node *n) {
+  assert(!IS_LEAF(n));
+
+  Node **pn = n->parent ? &n->parent->child[GET_INDEX(n->key, n->parent)] : &root_;
+  bool did_work = false;
+
+  if (n->empty_children >= (1U << n->bits) - 1) {
+    Collapse(pn);
+    n = *pn;
+    if (n == NULL || IS_LEAF(n))
+      return true;
+    did_work = true;
+  }
+
+  bool did_inflate = false;
+
+  // Double as long as the resulting node has a number of
+  // nonempty nodes that are above the threshold.
+  while ((n->full_children > 0 && 50 * (n->full_children + (1U << n->bits) - n->empty_children) >=
+    (n->parent && !(n->pos == 9) ? kInflateThreshold : kInflateThresholdRoot) * (1U << n->bits)) && n->bits < 16) {
+    if (!Inflate(pn))
+      break;
+    n = *pn;
+    did_work = true;
+    did_inflate = true;
+  }
+
+  // Halve as long as the number of empty children in this
+  // node is above threshold.
+  while (n->bits > 1 && 100 * ((1U << n->bits) - n->empty_children) <
+    (n->parent && !(n->pos == 8) ? kHalveThreshold : kHalveThresholdRoot) * (1U << n->bits)) {
+    assert(!did_inflate);
+    if (!Halve(pn))
+      break;
+
+    n = *pn;
+    did_work = true;
+  }
+
+  if (n->empty_children >= (1U << n->bits) - 1) {
+    Collapse(pn);
+    n = *pn;
+    if (n == NULL || IS_LEAF(n))
+      return true;
+    did_work = true;
+  }
+
+  if (did_work)
+    ResizeChildren(n);
+
+  return did_work;
+}
+
+void RoutingTrie32::UpdateParent(Node *pn) {
+  uint32 i_end = 1U << pn->bits;
+  uint32 mask_of_halves = 0;
+
+  for (uint32 i = 0; i != i_end; i++) {
+    Node *n = pn->child[i];
+    if (n == NULL)
+      continue;
+
+    mask_of_halves |= (i & 1) + 1;
+
+    if (NODE_IS_NULL_OR_OLEAF(n) || IS_LEAF(n))
+      continue;
+    Node *op = n->parent;
+    n->parent = pn;
+    if (op == NULL)
+      UpdateParent(n);
+  }
+  // Collapse right away if there's too many children
+  if (pn->empty_children + 1 >= i_end) {
+    Collapse(pn->parent ? &pn->parent->child[GET_INDEX(pn->key, pn->parent)] : &root_);
+    return;
+  }
+
+  if (mask_of_halves != 3) {
+    // Only one half of the entries are actually used. Perform a halving operation.
+    Halve(pn->parent ? &pn->parent->child[GET_INDEX(pn->key, pn->parent)] : &root_);
+  }
+}
+
+RoutingTrie32::Node *RoutingTrie32::ConvertOleafToLeaf(Node *pn, uint32 i, Node *n) {
+  return NewLeaf(pn->key + (i << pn->pos), pn->pos, VALUE_FROM_OLEAF(n));
+}
+
+// The |parent| and |leaf_next| pointers are repurposed to hold
+// the next pointer in the free list.
+class FreeableNodeCollector {
+public:
+  FreeableNodeCollector() : ptr_(NULL) {}
+
+  void Add(RoutingTrie32::Node *n) {
+    n->leaf_next = ptr_;
+    ptr_ = n;
+  }
+
+  void Revert(RoutingTrie32::Node *reset_parent_to);
+  void Free();
+private:
+  RoutingTrie32::Node *ptr_;
+};
+
+void FreeableNodeCollector::Revert(RoutingTrie32::Node *reset_parent_to) {
+  for (RoutingTrie32::Node *p = ptr_, *pn; p != NULL; p = pn) {
+    pn = p->parent;
+    p->parent = reset_parent_to;
+  }
+}
+
+void FreeableNodeCollector::Free() {
+  for (RoutingTrie32::Node *p = ptr_, *pn; p != NULL; p = pn) {
+    pn = p->parent;
+    FreeNode(p);
+  }
+}
+
+void RoutingTrie32::ReplaceChild(RoutingTrie32::Node **pnp, RoutingTrie32::Node *n) {
+  RoutingTrie32::Node *pn = *pnp;
+  if (pn->parent)
+    RoutingTrie32::PutChild(pn->parent, (uint32)(pnp - pn->parent->child), n);
+  else
+    *pnp = n;
+}
+
+bool RoutingTrie32::Inflate(Node **pnp) {
+  Node *pn = *pnp, *n0, *n1;
+  FreeableNodeCollector free_on_failure, free_on_success, free_on_success_null;
+  Node *tn = NewNode(pn->key, pn->pos - 1, pn->bits + 1);
+  if (!tn)
+    return false;
+  tn->parent = pn->parent;
+
+
+  uint8 oleaf_compare_value = tn->pos;
+
+  for (uint32 i = 0, i_end = 1U << pn->bits; i != i_end; i++) {
+    Node *n = pn->child[i];
+    // An empty child
+    if (n == NULL)
+      continue;
+
+    if (NODE_IS_OLEAF(n)) {
+      // Convert oleaf to leaf as parent's |pos| changed
+      if (!(n = ConvertOleafToLeaf(pn, i, n)))
+        goto nomem;
+      free_on_failure.Add(n);
+      goto insert_child;
+    }
+
+    if (IS_LEAF(n)) {
+      // Check whether the leaf can be converted to an oleaf.
+      if (n->pos == oleaf_compare_value && n->leaf_next == NULL) {
+        free_on_success_null.Add(n);
+        PutChild(tn, GET_INDEX(n->key, tn), VALUE_TO_OLEAF(n->leaf_value));
+        continue;
+      }
+      goto insert_child;
+    }
+
+    // A leaf or an internal node with skipped bits
+    if ((n->pos + n->bits) != pn->pos) {
+insert_child:
+      PutChild(tn, GET_INDEX(n->key, tn), n);
+      continue;
+    }
+    free_on_success.Add(n);
+    // Copying oleaf from here is ok as pos is unchanged.
+    if (n->bits == 1) {
+      // An internal node with exactly two children
+      n0 = n->child[0];
+      n1 = n->child[1];
+    } else {
+      // An internal node with more than two children
+      if (!(n1 = NewNode(n->key | (1 << tn->pos), n->pos, n->bits - 1)))
+        goto nomem;
+      free_on_failure.Add(n1);
+      if (!(n0 = NewNode(n->key, n->pos, n->bits - 1)))
+        goto nomem;
+      free_on_failure.Add(n0);
+      uint32 j_end = 1U << (n->bits - 1);
+      for (uint32 j = 0; j != j_end; j++) {
+        PutChild(n0, j, n->child[j]);
+        PutChild(n1, j, n->child[j + j_end]);
+      }
+    }
+    PutChild(tn, 2 * i + 0, n0);
+    PutChild(tn, 2 * i + 1, n1);
+  }
+
+  free_on_success.Free();
+  free_on_success_null.Free();
+  free_on_failure.Revert(NULL);
+
+  ReplaceChild(pnp, tn);
+  UpdateParent(tn);
+  FreeNode(pn);
+  return true;
+
+nomem:
+  free_on_success.Revert(pn);
+  free_on_success_null.Revert(NULL);
+  free_on_failure.Free();
+  FreeNode(tn);
+  return false;
+}
+
+bool RoutingTrie32::Halve(Node **pnp) {
+  Node *pn = *pnp, *n;
+  Node *tn = NewNode(pn->key, pn->pos + 1, pn->bits - 1);
+  FreeableNodeCollector free_on_failure, free_on_success_null;
+  if (!tn)
+    return false;
+  tn->parent = pn->parent;
+
+  uint8 oleaf_compare_value = tn->pos;
+
+  for (uint32 i = 0, i_end = 1U << pn->bits; i != i_end; i += 2) {
+    Node *n0 = pn->child[i + 0];
+    Node *n1 = pn->child[i + 1];
+
+    if (n0 == NULL || n1 == NULL) {
+      // At least one of the children is empty.
+      n = n0 ? n0 : n1;
+
+      if (NODE_IS_OLEAF(n)) {
+        // Convert oleaf to leaf as parent's |pos| changed
+        if (!(n = ConvertOleafToLeaf(pn, i + (n0 == NULL), n)))
+          goto nomem;
+        free_on_failure.Add(n);
+      } else if (n && IS_LEAF(n) && n->pos == oleaf_compare_value && n->leaf_next == NULL) {
+        // The leaf can be converted to an oleaf.
+        free_on_success_null.Add(n);
+        n = VALUE_TO_OLEAF(n->leaf_value);
+      }
+    } else {
+      // Two nonempty children
+      if (!(n = NewNode(pn->key + (i << pn->pos), pn->pos, 1)))
+        goto nomem;
+      free_on_failure.Add(n);
+      PutChild(n, 0, n0);
+      PutChild(n, 1, n1);
+    }
+    PutChild(tn, i / 2, n);
+  }
+  free_on_failure.Revert(NULL);
+  free_on_success_null.Free();
+  ReplaceChild(pnp, tn);
+  UpdateParent(tn);
+  FreeNode(pn);
+  return true;
+
+nomem:
+  free_on_failure.Free();
+  free_on_success_null.Revert(NULL);
+  FreeNode(tn);
+  return false;
+}
+
+void RoutingTrie32::Collapse(Node **pnp) {
+  Node *pn = *pnp, *n = NULL;
+
+  if (pn->empty_children != (1U << pn->bits)) {
+    for (uint32 i = 0; ; i++) {
+      n = pn->child[i];
+      if (n) {
+        if (NODE_IS_OLEAF(n)) {
+          if (!(n = ConvertOleafToLeaf(pn, i, n)))
+            return;
+        } else if (!IS_LEAF(n)) {
+          n->parent = pn->parent;
+        }
+        break;
+      }
+    }
+  }
+  ReplaceChild(pnp, n);
+  FreeNode(pn);
+}

ip_to_peer_map.h

@@ -5,15 +5,45 @@
 #include "tunsafe_types.h"
 #include <vector>
 
+class RoutingTrie32 {
+  friend class FreeableNodeCollector;
+public:
+  typedef void *Value;
+  struct Node;
+
+  RoutingTrie32();
+  ~RoutingTrie32();
+  NOINLINE Value Lookup(uint32 ip);
+  NOINLINE Value LookupExact(uint32 ip, int cidr);
+  bool Insert(uint32 ip, int cidr, Value value);
+  bool Delete(uint32 ip, int cidr);
+
+private:
+  Node *root_;
+
+  void Rebalance(Node *n);
+  bool Resize(Node *n);
+  bool Inflate(Node **n);
+  bool Halve(Node **n);
+  void UpdateParent(Node *n);
+  void ResizeChildren(Node *n);
+  static void Collapse(Node **n);
+  static void PutChild(Node *pn, uint32 i, Node *n);
+  static void ReplaceChild(Node **pnp, Node *n);
+  static Node *ConvertOleafToLeaf(Node *pn, uint32 i, Node *n);
+  static bool InsertLeafInto(Node **n, uint8 leaf_pos, RoutingTrie32::Value value);
+};
+
+
 // Maps CIDR addresses to a peer, always returning the longest match
-// Slow O(n) implementation
+// IPv6 has a slow O(n) implementation
 class IpToPeerMap {
 public:
   IpToPeerMap();
   ~IpToPeerMap();
 
   // Inserts an IP address of a given CIDR length into the lookup table, pointing to peer.
-  bool InsertV4(const void *addr, int cidr, void *peer);
+  bool InsertV4(uint32 ip, int cidr, void *peer);
   bool InsertV6(const void *addr, int cidr, void *peer);
 
   // Lookup the peer matching the IP Address
@@ -26,16 +56,12 @@ public:
   // Remove a peer from the table
   void RemovePeer(void *peer);
 private:
-  struct Entry4 {
-    uint32 ip;
-    uint32 mask;
-    void *peer;
-  };
   struct Entry6 {
     uint8 ip[16];
     uint8 cidr_len;
     void *peer;
   };
-  std::vector<Entry4> ipv4_;
   std::vector<Entry6> ipv6_;
+
+  RoutingTrie32 ipv4_;
 };

wireguard_proto.cpp

@@ -1155,7 +1155,7 @@ bool WgPeer::AddIp(const WgCidrAddr &cidr_addr) {
     if (cidr_addr.cidr > 32)
       return false;
     WG_ACQUIRE_RWLOCK_EXCLUSIVE(dev_->ip_to_peer_map_lock_);
-    dev_->ip_to_peer_map_.InsertV4(cidr_addr.addr, cidr_addr.cidr, this);
+    dev_->ip_to_peer_map_.InsertV4(ReadBE32(cidr_addr.addr), cidr_addr.cidr, this);
     WG_RELEASE_RWLOCK_EXCLUSIVE(dev_->ip_to_peer_map_lock_);
     allowed_ips_.push_back(cidr_addr);
     return true;