txgraph: Add Get{Ancestors,Descendants}Union functions (feature)

Like GetAncestors and GetDescendants, but for the union of multiple inputs.

Author: sipa

src/test/fuzz/txgraph.cpp

@@ -454,6 +454,28 @@ FUZZ_TARGET(txgraph)
                 auto expect_set = sel_sim.GetAncDesc(ref, alt);
                 assert(result_set == expect_set);
                 break;
+            } else if (!sel_sim.IsOversized() && command-- == 0) {
+                // GetAncestorsUnion/GetDescendantsUnion.
+                std::vector<TxGraph::Ref*> refs;
+                // Gather a list of up to 15 Ref pointers.
+                auto count = provider.ConsumeIntegralInRange<size_t>(0, 15);
+                refs.resize(count);
+                for (size_t i = 0; i < count; ++i) {
+                    refs[i] = pick_fn();
+                }
+                // Their order should not matter, shuffle them.
+                std::shuffle(refs.begin(), refs.end(), rng);
+                // Invoke the real function, and convert to SimPos set.
+                auto result = alt ? real->GetDescendantsUnion(refs, use_main)
+                                  : real->GetAncestorsUnion(refs, use_main);
+                auto result_set = sel_sim.MakeSet(result);
+                assert(result.size() == result_set.Count());
+                // Compute the expected result.
+                SimTxGraph::SetType expect_set;
+                for (TxGraph::Ref* ref : refs) expect_set |= sel_sim.GetAncDesc(ref, alt);
+                // Compare.
+                assert(result_set == expect_set);
+                break;
             } else if (!sel_sim.IsOversized() && command-- == 0) {
                 // GetCluster.
                 auto ref = pick_fn();

src/txgraph.cpp

@@ -453,6 +453,8 @@ class TxGraphImpl final : public TxGraph
     std::vector<Ref*> GetCluster(const Ref& arg, bool main_only = false) noexcept final;
     std::vector<Ref*> GetAncestors(const Ref& arg, bool main_only = false) noexcept final;
     std::vector<Ref*> GetDescendants(const Ref& arg, bool main_only = false) noexcept final;
+    std::vector<Ref*> GetAncestorsUnion(std::span<const Ref* const> args, bool main_only = false) noexcept final;
+    std::vector<Ref*> GetDescendantsUnion(std::span<const Ref* const> args, bool main_only = false) noexcept final;
     GraphIndex GetTransactionCount(bool main_only = false) noexcept final;
     bool IsOversized(bool main_only = false) noexcept final;
     std::strong_ordering CompareMainOrder(const Ref& a, const Ref& b) noexcept final;
@@ -1581,6 +1583,70 @@ std::vector<TxGraph::Ref*> TxGraphImpl::GetDescendants(const Ref& arg, bool main
     return ret;
 }
 
+std::vector<TxGraph::Ref*> TxGraphImpl::GetAncestorsUnion(std::span<const Ref* const> args, bool main_only) noexcept
+{
+    // Apply all dependencies, as the result might be incorrect otherwise.
+    size_t level = GetSpecifiedLevel(main_only);
+    ApplyDependencies(level);
+    // Ancestry cannot be known if unapplied dependencies remain.
+    Assume(GetClusterSet(level).m_deps_to_add.empty());
+
+    // Translate args to matches.
+    std::vector<std::pair<Cluster*, DepGraphIndex>> matches;
+    matches.reserve(args.size());
+    for (auto arg : args) {
+        // Skip empty Refs.
+        if (GetRefGraph(*arg) == nullptr) continue;
+        Assume(GetRefGraph(*arg) == this);
+        // Find the Cluster the argument is in, and skip if none is found.
+        auto cluster = FindCluster(GetRefIndex(*arg), level);
+        if (cluster == nullptr) continue;
+        // Append to matches.
+        matches.emplace_back(cluster, m_entries[GetRefIndex(*arg)].m_locator[cluster->m_level].index);
+    }
+    // Group by Cluster.
+    std::sort(matches.begin(), matches.end(), [](auto& a, auto& b) noexcept { return std::less{}(a.first, b.first); });
+    // Dispatch to the Clusters.
+    std::span match_span(matches);
+    std::vector<TxGraph::Ref*> ret;
+    while (!match_span.empty()) {
+        match_span.front().first->GetAncestorRefs(*this, match_span, ret);
+    }
+    return ret;
+}
+
+std::vector<TxGraph::Ref*> TxGraphImpl::GetDescendantsUnion(std::span<const Ref* const> args, bool main_only) noexcept
+{
+    // Apply all dependencies, as the result might be incorrect otherwise.
+    size_t level = GetSpecifiedLevel(main_only);
+    ApplyDependencies(level);
+    // Ancestry cannot be known if unapplied dependencies remain.
+    Assume(GetClusterSet(level).m_deps_to_add.empty());
+
+    // Translate args to matches.
+    std::vector<std::pair<Cluster*, DepGraphIndex>> matches;
+    matches.reserve(args.size());
+    for (auto arg : args) {
+        // Skip empty Refs.
+        if (GetRefGraph(*arg) == nullptr) continue;
+        Assume(GetRefGraph(*arg) == this);
+        // Find the Cluster the argument is in, and skip if none is found.
+        auto cluster = FindCluster(GetRefIndex(*arg), level);
+        if (cluster == nullptr) continue;
+        // Append to matches.
+        matches.emplace_back(cluster, m_entries[GetRefIndex(*arg)].m_locator[cluster->m_level].index);
+    }
+    // Group by Cluster.
+    std::sort(matches.begin(), matches.end(), [](auto& a, auto& b) noexcept { return std::less{}(a.first, b.first); });
+    // Dispatch to the Clusters.
+    std::span match_span(matches);
+    std::vector<TxGraph::Ref*> ret;
+    while (!match_span.empty()) {
+        match_span.front().first->GetDescendantRefs(*this, match_span, ret);
+    }
+    return ret;
+}
+
 std::vector<TxGraph::Ref*> TxGraphImpl::GetCluster(const Ref& arg, bool main_only) noexcept
 {
     // Return the empty vector if the Ref is empty (which may be indicative of the transaction

src/txgraph.h

@@ -142,6 +142,14 @@ class TxGraph
      *  queried; otherwise the main graph is queried. The queried graph must not be oversized.
      *  Returns {} if arg does not exist in the graph. */
     virtual std::vector<Ref*> GetDescendants(const Ref& arg, bool main_only = false) noexcept = 0;
+    /** Like GetAncestors, but return the Refs for all transactions in the union of the provided
+     *  arguments' ancestors (each transaction is only reported once). Refs that do not exist in
+     *  the queried graph are ignored. */
+    virtual std::vector<Ref*> GetAncestorsUnion(std::span<const Ref* const> args, bool main_only = false) noexcept = 0;
+    /** Like GetDescendants, but return the Refs for all transactions in the union of the provided
+     *  arguments' descendants (each transaction is only reported once). Refs that do not exist in
+     *  the queried graph are ignored. */
+    virtual std::vector<Ref*> GetDescendantsUnion(std::span<const Ref* const> args, bool main_only = false) noexcept = 0;
     /** Get the total number of transactions in the graph. If main_only is false and a staging
      *  graph exists, it is queried; otherwise the main graph is queried. This is available even
      *  for oversized graphs. */