diff --git a/llarp/nodedb.cpp b/llarp/nodedb.cpp index 39d9e8c05..91e8d4ed2 100644 --- a/llarp/nodedb.cpp +++ b/llarp/nodedb.cpp @@ -61,14 +61,16 @@ namespace llarp } std::optional> - NodeDB::get_n_random_rcs(size_t n) const + NodeDB::get_n_random_rcs(size_t n, bool exact) const { - std::vector rand{}; + auto rand = std::make_optional>(); + rand->reserve(n); - std::sample(known_rcs.begin(), known_rcs.end(), std::back_inserter(rand), n, csrng); - return rand.empty() ? std::nullopt : std::make_optional(rand); + std::sample(known_rcs.begin(), known_rcs.end(), std::back_inserter(*rand), n, csrng); + if (rand->size() < (exact ? n : 1)) + rand.reset(); + return rand; } - std::optional NodeDB::get_random_rc_conditional(std::function hook) const { @@ -99,10 +101,11 @@ namespace llarp } std::optional> - NodeDB::get_n_random_rcs_conditional(size_t n, std::function hook) const + NodeDB::get_n_random_rcs_conditional( + size_t n, std::function hook, bool exact) const { - std::vector selected; - selected.reserve(n); + auto selected = std::make_optional>(); + selected->reserve(n); size_t i = 0; @@ -115,17 +118,19 @@ namespace llarp // load the first n RC's that pass the condition into selected if (++i <= n) { - selected.push_back(rc); + selected->push_back(rc); continue; } // replace selections with decreasing probability per iteration size_t x = csrng() % (i + 1); if (x < n) - selected[x] = rc; + (*selected)[x] = rc; } - return selected.size() == n ? std::make_optional(selected) : std::nullopt; + if (selected->size() < (exact ? n : 1)) + selected.reset(); + return selected; } void diff --git a/llarp/nodedb.hpp b/llarp/nodedb.hpp index 3bff7be1f..57a835915 100644 --- a/llarp/nodedb.hpp +++ b/llarp/nodedb.hpp @@ -429,8 +429,11 @@ namespace llarp std::optional get_random_rc() const; + // Get `n` random RCs from all RCs we know about. If `exact` is true then we require n matches + // (and otherwise return nullopt); otherwise we return whatever we found, or nullopt if we find + // nothing at all. std::optional> - get_n_random_rcs(size_t n) const; + get_n_random_rcs(size_t n, bool exact = false) const; /** The following random conditional functions utilize a simple implementation of reservoir sampling to return either 1 or n random RC's using only one pass through the set of RC's. @@ -447,7 +450,7 @@ namespace llarp get_random_rc_conditional(std::function hook) const; std::optional> - get_n_random_rcs_conditional(size_t n, std::function hook) const; + get_n_random_rcs_conditional(size_t n, std::function hook, bool exact = false) const; // Updates `current` to not contain any of the elements of `replace` and resamples (up to // `target_size`) from population to refill it.