radix: Add a find method

Like the find method of standard containers, this returns an iterator to the requested key.

radix: Add a find method
6478af3c · Austin Clements · 0ee57d47 · 6478af3c · 6478af3c
--- a/include/radix.hh
+++ b/include/radix.hh
@@ -218,6 +218,8 @@ struct radix {
  radix_elem* search(u64 key);
  radix_range search_lock(u64 start, u64 size);

+  iterator find(u64 key);
+
  radix_iterator begin() const;
  radix_iterator end() const;

@@ -225,11 +227,11 @@ struct radix {
 };

 struct radix_iterator {
-  radix_iterator(const radix* r, u64 k, u64 limit)
+  radix_iterator(const radix* r, u64 k, u64 limit, bool skip_nulls)
    : r_(r), k_(k), key_limit_(limit)
  {
    if (k_ != key_limit_)
-      prime_path();
+      prime_path(skip_nulls);
  }
  radix_iterator() = default;
  radix_iterator(const radix_iterator &o) = default;
@@ -250,10 +252,11 @@ struct radix_iterator {

  // Advance the iterator until it points at a non-null entry or end.
  // If the current element is non-null, does nothing.
-  void skip_nulls()
+  radix_iterator &skip_nulls()
  {
    if (!**this)
      ++(*this);
+    return *this;
  }

  // Return the key of the iterator's current element.
@@ -291,7 +294,7 @@ private:
  u32 level_;

  // Prepare the initial path_ and level_ based on k_.
-  void prime_path();
+  void prime_path(bool skip_nulls);
  // Advance to the next leaf.  If skip_nulls is true, advances to the
  // next non-null leaf.  This assumes that k_ < key_limit_.
  void advance(bool skip_nulls = true);
@@ -299,20 +302,20 @@ private:

 inline radix_iterator
 radix_range::begin() const {
-  return radix_iterator(r_, start_, start_ + size_);
+  return radix_iterator(r_, start_, start_ + size_, true);
 }

 inline radix_iterator
 radix_range::end() const {
-  return radix_iterator(r_, start_ + size_, start_ + size_);
+  return radix_iterator(r_, start_ + size_, start_ + size_, true);
 }

 inline radix_iterator
 radix::begin() const {
-  return radix_iterator(this, 0, (u64)1 << key_bits);
+  return radix_iterator(this, 0, (u64)1 << key_bits, true);
 }

 inline radix_iterator
 radix::end() const {
-  return radix_iterator(this, (u64)1 << key_bits, (u64)1 << key_bits);
+  return radix_iterator(this, (u64)1 << key_bits, (u64)1 << key_bits, true);
 }
--- a/kernel/radix.cc
+++ b/kernel/radix.cc
@@ -180,6 +180,12 @@ radix::search_lock(u64 start, u64 size)
  return radix_range(this, start >> shift_, size >> shift_);
 }

+radix::iterator
+radix::find(u64 key)
+{
+  return radix::iterator(this, key >> shift_, (u64)1 << key_bits, false);
+}
+
 radix_range::radix_range(radix *r, u64 start, u64 size)
  : r_(r), start_(start), size_(size)
 {
@@ -261,8 +267,13 @@ radix_iterator::next_change() const
 }

 void
-radix_iterator::prime_path()
+radix_iterator::prime_path(bool skip_nulls)
 {
+  // XXX(amdragon) Do this lazily in operator*.  Then we don't need
+  // special-casing for the end iterator and callers can jump around
+  // the keyspace however they want (including, say, starting with the
+  // begin iterator of the whole radix tree and jumping to a
+  // particular key.)
  dprintf("%p: Made iterator with k = %lx\n", r_, k_);

  // Load the initial path
@@ -276,7 +287,8 @@ radix_iterator::prime_path()
  }

  // Find a real element
-  skip_nulls();
+  if (skip_nulls)
+    this->skip_nulls();

  dprintf("%p: Adjusted: k = %lx\n", r_, k_);
 }