LedgerTrie.h

#pragma once
 
#include <xrpl/basics/ToString.h>
#include <xrpl/beast/utility/instrumentation.h>
#include <xrpl/json/json_value.h>
 
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <iomanip>
#include <map>
#include <memory>
#include <optional>
#include <ostream>
#include <sstream>
#include <stack>
#include <utility>
#include <vector>
 
namespace xrpl {

template <class Ledger>
class SpanTip
{
public:
    using Seq = Ledger::Seq;
    using ID = Ledger::ID;
 
    SpanTip(Seq s, ID i, Ledger const lgr) : seq{s}, id{i}, ledger_{std::move(lgr)}
    {
    }
 
    // The sequence number of the tip ledger
    Seq seq;
    // The ID of the tip ledger
    ID id;

    [[nodiscard]] ID
    ancestor(Seq const& s) const
    {
        XRPL_ASSERT(s <= seq, "xrpl::SpanTip::ancestor : valid input");
        return ledger_[s];
    }
 
private:
    Ledger const ledger_;
};
 
namespace ledger_trie_detail {
 
// Represents a span of ancestry of a ledger
template <class Ledger>
class Span
{
    using Seq = Ledger::Seq;
    using ID = Ledger::ID;
 
    // The span is the half-open interval [start,end) of ledger_
    Seq start_{0};
    Seq end_{1};
    Ledger ledger_;
 
public:
    Span() : ledger_{typename Ledger::MakeGenesis{}}
    {
        // Require default ledger to be genesis seq
        XRPL_ASSERT(ledger_.seq() == start_, "xrpl::Span::Span : ledger is genesis");
    }
 
    Span(Ledger ledger) : end_{ledger.seq() + Seq{1}}, ledger_{std::move(ledger)}
    {
    }
 
    Span(Span const& s) = default;
    Span(Span&& s) = default;
    Span&
    operator=(Span const&) = default;
    Span&
    operator=(Span&&) = default;
 
    [[nodiscard]] Seq
    start() const
    {
        return start_;
    }
 
    [[nodiscard]] Seq
    end() const
    {
        return end_;
    }
 
    // Return the Span from [spot,end_) or none if no such valid span
    [[nodiscard]] std::optional<Span>
    from(Seq spot) const
    {
        return sub(spot, end_);
    }
 
    // Return the Span from [start_,spot) or none if no such valid span
    [[nodiscard]] std::optional<Span>
    before(Seq spot) const
    {
        return sub(start_, spot);
    }
 
    // Return the ID of the ledger that starts this span
    [[nodiscard]] ID
    startID() const
    {
        return ledger_[start_];
    }
 
    // Return the ledger sequence number of the first possible difference
    // between this span and a given ledger.
    [[nodiscard]] Seq
    diff(Ledger const& o) const
    {
        return clamp(mismatch(ledger_, o));
    }
 
    //  The tip of this span
    [[nodiscard]] SpanTip<Ledger>
    tip() const
    {
        Seq const tipSeq{end_ - Seq{1}};
        return SpanTip<Ledger>{tipSeq, ledger_[tipSeq], ledger_};
    }
 
private:
    Span(Seq start, Seq end, Ledger l) : start_{start}, end_{end}, ledger_{std::move(l)}
    {
        // Spans cannot be empty
        XRPL_ASSERT(start < end, "xrpl::Span::Span : non-empty span input");
    }
 
    [[nodiscard]] Seq
    clamp(Seq val) const
    {
        return std::min(std::max(start_, val), end_);
    }
 
    // Return a span of this over the half-open interval [from,to)
    [[nodiscard]] std::optional<Span>
    sub(Seq from, Seq to) const
    {
        Seq const newFrom = clamp(from);
        Seq const newTo = clamp(to);
        if (newFrom < newTo)
            return Span(newFrom, newTo, ledger_);
        return std::nullopt;
    }
 
    friend std::ostream&
    operator<<(std::ostream& o, Span const& s)
    {
        return o << s.tip().id << "[" << s.start_ << "," << s.end_ << ")";
    }
 
    friend Span
    merge(Span const& a, Span const& b)
    {
        // Return combined span, using ledger_ from higher sequence span
        if (a.end_ < b.end_)
            return Span(std::min(a.start_, b.start_), b.end_, b.ledger_);
 
        return Span(std::min(a.start_, b.start_), a.end_, a.ledger_);
    }
};
 
// A node in the trie
template <class Ledger>
struct Node
{
    Node() = default;
 
    explicit Node(Ledger const& l) : span{l}, tipSupport{1}, branchSupport{1}
    {
    }
 
    explicit Node(Span<Ledger> s) : span{std::move(s)}
    {
    }
 
    Span<Ledger> span;
    std::uint32_t tipSupport = 0;
    std::uint32_t branchSupport = 0;
 
    std::vector<std::unique_ptr<Node>> children;
    Node* parent = nullptr;

    void
    erase(Node const* child)
    {
        auto it = std::ranges::find_if(
            children, [child](std::unique_ptr<Node> const& curr) { return curr.get() == child; });
        XRPL_ASSERT(it != children.end(), "xrpl::Node::erase : valid input");
        std::swap(*it, children.back());
        children.pop_back();
    }
 
    friend std::ostream&
    operator<<(std::ostream& o, Node const& s)
    {
        return o << s.span << "(T:" << s.tipSupport << ",B:" << s.branchSupport << ")";
    }
 
    [[nodiscard]] json::Value
    getJson() const
    {
        json::Value res;
        std::stringstream sps;
        sps << span;
        res["span"] = sps.str();
        res["startID"] = to_string(span.startID());
        res["seq"] = static_cast<std::uint32_t>(span.tip().seq);
        res["tipSupport"] = tipSupport;
        res["branchSupport"] = branchSupport;
        if (!children.empty())
        {
            json::Value& cs = (res["children"] = json::ValueType::Array);
            for (auto const& child : children)
            {
                cs.append(child->getJson());
            }
        }
        return res;
    }
};
}  // namespace ledger_trie_detail

template <class Ledger>
class LedgerTrie
{
    using Seq = Ledger::Seq;
    using ID = Ledger::ID;
 
    using Node = ledger_trie_detail::Node<Ledger>;
    using Span = ledger_trie_detail::Span<Ledger>;
 
    // The root of the trie. The root is allowed to break the no-single child
    // invariant.
    std::unique_ptr<Node> root_;
 
    // Count of the tip support for each sequence number
    std::map<Seq, std::uint32_t> seqSupport_;

    [[nodiscard]] std::pair<Node*, Seq>
    find(Ledger const& ledger) const
    {
        // NOLINTNEXTLINE(misc-const-correctness)
        Node* curr = root_.get();
 
        // Root is always defined and is in common with all ledgers
        XRPL_ASSERT(curr, "xrpl::LedgerTrie::find : non-null root");
        Seq pos = curr->span.diff(ledger);
 
        bool done = false;
 
        // Continue searching for a better span as long as the current position
        // matches the entire span
        while (!done && pos == curr->span.end())
        {
            done = true;
            // Find the child with the longest ancestry match
            for (std::unique_ptr<Node> const& child : curr->children)
            {
                auto const childPos = child->span.diff(ledger);
                if (childPos > pos)
                {
                    done = false;
                    pos = childPos;
                    curr = child.get();
                    break;
                }
            }
        }
        return std::make_pair(curr, pos);
    }

    Node*
    findByLedgerID(Ledger const& ledger, Node* parent = nullptr) const
    {
        if (parent == nullptr)
            parent = root_.get();
        if (ledger.id() == parent->span.tip().id)
            return parent;
        for (auto const& child : parent->children)
        {
            auto cl = findByLedgerID(ledger, child.get());
            if (cl)
                return cl;
        }
        return nullptr;
    }
 
    void
    dumpImpl(std::ostream& o, std::unique_ptr<Node> const& curr, int offset) const
    {
        if (curr)
        {
            if (offset > 0)
                o << std::setw(offset) << "|-";
 
            std::stringstream ss;
            ss << *curr;
            o << ss.str() << std::endl;
            for (std::unique_ptr<Node> const& child : curr->children)
                dumpImpl(o, child, offset + 1 + ss.str().size() + 2);
        }
    }
 
public:
    LedgerTrie() : root_{std::make_unique<Node>()}
    {
    }

    void
    insert(Ledger const& ledger, std::uint32_t count = 1)
    {
        auto const [loc, diffSeq] = find(ledger);
 
        // There is always a place to insert
        XRPL_ASSERT(loc, "xrpl::LedgerTrie::insert : valid input ledger");
 
        // Node from which to start incrementing branchSupport
        Node* incNode = loc;
 
        // loc->span has the longest common prefix with Span{ledger} of all
        // existing nodes in the trie. The optional<Span>'s below represent
        // the possible common suffixes between loc->span and Span{ledger}.
        //
        // loc->span
        //  a b c  | d e f
        //  prefix | oldSuffix
        //
        // Span{ledger}
        //  a b c  | g h i
        //  prefix | newSuffix
 
        std::optional<Span> prefix = loc->span.before(diffSeq);
        std::optional<Span> oldSuffix = loc->span.from(diffSeq);
        std::optional<Span> newSuffix = Span{ledger}.from(diffSeq);
 
        if (oldSuffix)
        {
            // Have
            //   abcdef -> ....
            // Inserting
            //   abc
            // Becomes
            //   abc -> def -> ...
 
            // Create oldSuffix node that takes over loc
            auto newNode = std::make_unique<Node>(*oldSuffix);
            newNode->tipSupport = loc->tipSupport;
            newNode->branchSupport = loc->branchSupport;
            newNode->children = std::move(loc->children);
            XRPL_ASSERT(loc->children.empty(), "xrpl::LedgerTrie::insert : moved-from children");
            for (std::unique_ptr<Node>& child : newNode->children)
                child->parent = newNode.get();
 
            // Loc truncates to prefix and newNode is its child
            XRPL_ASSERT(prefix, "xrpl::LedgerTrie::insert : prefix is set");
            loc->span = *prefix;  // NOLINT(bugprone-unchecked-optional-access) assert above
            newNode->parent = loc;
            loc->children.emplace_back(std::move(newNode));
            loc->tipSupport = 0;
        }
        if (newSuffix)
        {
            // Have
            //  abc -> ...
            // Inserting
            //  abcdef-> ...
            // Becomes
            //  abc -> ...
            //     \-> def
 
            auto newNode = std::make_unique<Node>(*newSuffix);
            newNode->parent = loc;
            // increment support starting from the new node
            incNode = newNode.get();
            loc->children.push_back(std::move(newNode));
        }
 
        incNode->tipSupport += count;
        while (incNode)
        {
            incNode->branchSupport += count;
            incNode = incNode->parent;
        }
 
        seqSupport_[ledger.seq()] += count;
    }

    bool
    remove(Ledger const& ledger, std::uint32_t count = 1)
    {
        Node* loc = findByLedgerID(ledger);
        // Must be exact match with tip support
        if ((loc == nullptr) || loc->tipSupport == 0)
            return false;
 
        // found our node, remove it
        count = std::min(count, loc->tipSupport);
        loc->tipSupport -= count;
 
        auto const it = seqSupport_.find(ledger.seq());
        XRPL_ASSERT(
            it != seqSupport_.end() && it->second >= count,
            "xrpl::LedgerTrie::remove : valid input ledger");
        it->second -= count;
        if (it->second == 0)
            seqSupport_.erase(it->first);
 
        Node* decNode = loc;
        while (decNode)
        {
            decNode->branchSupport -= count;
            decNode = decNode->parent;
        }
 
        while (loc->tipSupport == 0 && loc != root_.get())
        {
            Node* parent = loc->parent;
            if (loc->children.empty())
            {
                // this node can be erased
                parent->erase(loc);
            }
            else if (loc->children.size() == 1)
            {
                // This node can be combined with its child
                std::unique_ptr<Node> child = std::move(loc->children.front());
                child->span = merge(loc->span, child->span);
                child->parent = parent;
                parent->children.emplace_back(std::move(child));
                parent->erase(loc);
            }
            else
            {
                break;
            }
            loc = parent;
        }
        return true;
    }

    [[nodiscard]] std::uint32_t
    tipSupport(Ledger const& ledger) const
    {
        if (auto const* loc = findByLedgerID(ledger))
            return loc->tipSupport;
        return 0;
    }

    [[nodiscard]] std::uint32_t
    branchSupport(Ledger const& ledger) const
    {
        Node const* loc = findByLedgerID(ledger);
        if (loc == nullptr)
        {
            Seq diffSeq;
            std::tie(loc, diffSeq) = find(ledger);
            // Check that ledger is a proper prefix of loc
            if (!(diffSeq > ledger.seq() && ledger.seq() < loc->span.end()))
                loc = nullptr;
        }
        return loc ? loc->branchSupport : 0;
    }

    [[nodiscard]] std::optional<SpanTip<Ledger>>
    getPreferred(Seq const largestIssued) const
    {
        if (empty())
            return std::nullopt;
 
        Node* curr = root_.get();
 
        bool done = false;
 
        std::uint32_t uncommitted = 0;
        auto uncommittedIt = seqSupport_.begin();
 
        while (curr && !done)
        {
            // Within a single span, the preferred by branch strategy is simply
            // to continue along the span as long as the branch support of
            // the next ledger exceeds the uncommitted support for that ledger.
            {
                // Add any initial uncommitted support prior for ledgers
                // earlier than nextSeq or earlier than largestIssued
                Seq nextSeq = curr->span.start() + Seq{1};
                while (uncommittedIt != seqSupport_.end() &&
                       uncommittedIt->first < std::max(nextSeq, largestIssued))
                {
                    uncommitted += uncommittedIt->second;
                    uncommittedIt++;
                }
 
                // Advance nextSeq along the span
                while (nextSeq < curr->span.end() && curr->branchSupport > uncommitted)
                {
                    // Jump to the next seqSupport change
                    if (uncommittedIt != seqSupport_.end() &&
                        uncommittedIt->first < curr->span.end())
                    {
                        nextSeq = uncommittedIt->first + Seq{1};
                        uncommitted += uncommittedIt->second;
                        uncommittedIt++;
                    }
                    else
                    {  // otherwise we jump to the end of the span
                        nextSeq = curr->span.end();
                    }
                }
                // We did not consume the entire span, so we have found the
                // preferred ledger
                if (nextSeq < curr->span.end())
                {
                    // nextSeq within span guarantees before() is set
                    // NOLINTNEXTLINE(bugprone-unchecked-optional-access)
                    return curr->span.before(nextSeq)->tip();
                }
            }
 
            // We have reached the end of the current span, so we need to
            // find the best child
            Node* best = nullptr;
            std::uint32_t margin = 0;
            if (curr->children.size() == 1)
            {
                best = curr->children[0].get();
                margin = best->branchSupport;
            }
            else if (!curr->children.empty())
            {
                // Sort placing children with largest branch support in the
                // front, breaking ties with the span's starting ID
                std::partial_sort(
                    curr->children.begin(),
                    curr->children.begin() + 2,
                    curr->children.end(),
                    [](std::unique_ptr<Node> const& a, std::unique_ptr<Node> const& b) {
                        return std::make_tuple(a->branchSupport, a->span.startID()) >
                            std::make_tuple(b->branchSupport, b->span.startID());
                    });
 
                best = curr->children[0].get();
                margin = curr->children[0]->branchSupport - curr->children[1]->branchSupport;
 
                // If best holds the tie-breaker, gets one larger margin
                // since the second best needs additional branchSupport
                // to overcome the tie
                if (best->span.startID() > curr->children[1]->span.startID())
                    margin++;
            }
 
            // If the best child has margin exceeding the uncommitted support,
            // continue from that child, otherwise we are done
            if (best && ((margin > uncommitted) || (uncommitted == 0)))
            {
                curr = best;
            }
            else
            {  // current is the best
                done = true;
            }
        }
        return curr->span.tip();
    }

    [[nodiscard]] bool
    empty() const
    {
        return !root_ || root_->branchSupport == 0;
    }

    void
    dump(std::ostream& o) const
    {
        dumpImpl(o, root_, 0);
    }

    [[nodiscard]] json::Value
    getJson() const
    {
        json::Value res;
        res["trie"] = root_->getJson();
        res["seq_support"] = json::ValueType::Object;
        for (auto const& [seq, sup] : seqSupport_)
            res["seq_support"][to_string(seq)] = sup;
        return res;
    }

    [[nodiscard]] bool
    checkInvariants() const
    {
        std::map<Seq, std::uint32_t> expectedSeqSupport;
 
        std::stack<Node const*> nodes;
        nodes.push(root_.get());
        while (!nodes.empty())
        {
            Node const* curr = nodes.top();
            nodes.pop();
            if (curr == nullptr)
                continue;
 
            // Node with 0 tip support must have multiple children
            // unless it is the root node
            if (curr != root_.get() && curr->tipSupport == 0 && curr->children.size() < 2)
                return false;
 
            // branchSupport = tipSupport + sum(child->branchSupport)
            std::size_t support = curr->tipSupport;
            if (curr->tipSupport != 0)
                expectedSeqSupport[curr->span.end() - Seq{1}] += curr->tipSupport;
 
            for (auto const& child : curr->children)
            {
                if (child->parent != curr)
                    return false;
 
                support += child->branchSupport;
                nodes.push(child.get());
            }
            if (support != curr->branchSupport)
                return false;
        }
        return expectedSeqSupport == seqSupport_;
    }
};
 
}  // namespace xrpl