Understanding the Stake-Based Weighted Selection Algorithm in Solidity

Oyster
1

Introduction

In blockchain applications, especially those involving decentralized networks, selecting participants based on their stake (or weight) is a common requirement. The TreeMapUpgradeable contract implements an efficient stake-based weighted selection algorithm using a binary tree structure. This post breaks down how it works, why it’s gas-efficient, and how it ensures fair selection proportional to stakes.

Core Problem: Weighted Random Selection

We need to:

  1. Store participants and their stakes (weights).
  2. Select N random participants, where the probability of selection is proportional to their stake.
  3. Ensure efficiency, as on-chain operations must minimize gas costs.

Solution: Binary Tree with Weight Aggregation

The contract uses a complete binary tree where:

  • Each node stores:
    • Its own value (stake).
    • leftSum (total stake in the left subtree).
    • rightSum (total stake in the right subtree).
  • The tree is always balanced (insertions follow a leftmost-empty-leaf approach).

Key Features

  1. Efficient Updates

    • When a stake changes, only O(log n) nodes need updating (since only ancestor nodes’ leftSum / rightSum are modified).

  2. Weighted Random Selection (_selectN)

    • Uses binary search over the tree to select nodes proportionally to their stake.
    • Excludes already-selected nodes dynamically.
    • Contract uses a persistent on-chain storage tree which contains all the nodes.
    • When selecting N random nodes, the algorithm builds a temporary in-memory tree (_selectedPathTree) to:
      i. Track which nodes have been selected.
      ii. Subtract their weights from future selections without modifying storage tree.

  3. Gas Optimization

  • Uses assembly for memory management to reduce gas costs.

    MemoryNode[] memory _selectedPathTree;
assembly {
    _selectedPathTree := mload(0x40)
    mstore(0x40, add(_selectedPathTree, 2688))
    mstore(_selectedPathTree, 83)
}

    In solidity, when we initialize an array with a fixed size, all the elements are automatically set to their default values (e.g., zero for integers). This initialization incurs a gas cost, which can become significant for large arrays.

    To optimize gas usage, we avoid default initialization by declaring the _selectedPathTree array without setting its size upfront. Instead, we use inline assembly to manually allocate 83 memory slots for it. This approach bypasses Solidity’s default zero-initialization, resulting in notable gas savings - over 16,000 units in this case.

    The size of 83 slots is determined based on the structure of the staking tree in Marlin, where the staking token is POND. Given the total POND supply, the tree can contain up to (215−1) nodes. From these, up to 6 nodes can be selected for the _selectedPathTree. If needed, the array length can be increased to select more nodes.

    By leveraging assembly for memory allocation, we significantly reduce gas costs while maintaining flexibility in the number of nodes selected.

  • The TreeMapUpgradeable contract makes a deliberate design choice to optimize for efficient weighted random selection (_selectN) at the cost of more expensive insertions/deletions.

    The tree stores leftSum and rightSum for each node, enabling:

    • O(log n) search time for weighted selection (traverse left/right based on precomputed sums).
    • O(log n) insertion/deletion time (must update ancestor sums).

    Without subtree sums, selecting a weighted random node would require:

    • O(n) traversal to compute weights dynamically.
    • Prohibitively expensive gas costs for large trees.

How the Algorithm Works

1. Tree Structure

The tree is stored as an array (index 1 = root, 2 = left child, 3 = right child, 4 = left-left child, 5 =left-right child, etc.).
Each node tracks:

struct Node {
    uint64 value;     // Node's own stake
    uint64 leftSum;   // Sum of left subtree
    uint64 rightSum;  // Sum of right subtree
}

2. Insertion & Deletion

  • Insertion (_insert_unchecked)
    • New nodes are added to the leftmost empty leaf.
    • Ancestor nodes update their leftSum / rightSum.
  • Deletion (_delete_unchecked)
    • The deleted node is replaced by the rightmost leaf to maintain balance.
    • Ancestor sums are updated.

3. Weighted Selection (_selectN)

This is the most complex part. Here’s how it works:

Step 1: Generate a Random Search Key

  • A random number _searchNumber is generated within the remaining stake range.

    _searchNumber = _randomizer % (_totalWeightInTree - _sumOfBalancesOfSelectedNodes);

Step 2: Traverse the Tree (_selectOne)

  • Start from the root (index = 1).
  • At each node, check:
    • Left subtree range: [0, leftSum)
    • Current node range: [leftSum, leftSum + value)
    • Right subtree range: [leftSum + value, leftSum + value + rightSum)
  • Depending on where _searchNumber falls:
    • Left subtree: Move to left child (index * 2).
    • Current node: Select it.
    • Right subtree: Move to right child (index * 2 + 1).

Step 3: Exclude Selected Node

  • Once a node is selected:
    • Its value is subtracted from future selections
      (_sumOfBalancesOfSelectedNodes += selectedValue).
    • The tree is not modified (only an in-memory _selectedPathTree tracks exclusions).
      (_selectedPathTree[index].value += value)

Step 4: Repeat Until N Selections

  • The loop continues until N nodes are picked.

Example Walkthrough

Assume the following tree (values = stakes):

                  Root (Index 1, Value=10)
                    /             \
              Left (2, 5)          Right (3, 5)
               /      \
 Left-Left (4, 3)    Left-Right (5, 2)

Total Stake = 25
We want to select N=2 nodes.

First Selection (_searchNumber = 15)

  1. Root (Index 1):
    • Left subtree ([0,10)), Root ([10,20)), Right ([20,25)).
    • 15 falls in [10,20), so Root (Value=10) is selected.
    • Remaining stake = 25 - 10 = 15.

Second Selection (_searchNumber = 7)

  1. Root (Index 1) is excluded, so we check:
    • Left subtree ([0,10)), Right ([10,15)).
    • 7 falls in [0,10), so move to Left (Index 2).
  2. Left (Index 2):
    • Left-Left ([0,3)), Left ([3,8)), Left-Right ([8,10)).
    • 7 falls in [3,8), so Left (Value=5) is selected.

Final Selection: [Root, Left] (Addresses [0x1, 0x2]).

Why This Algorithm is Efficient

  1. O(log n) Selection per Node
    • The tree depth is logarithmic, making each selection fast.
  2. No Restructuring Needed
    • Unlike heap-based approaches, the tree remains balanced without reordering.
  3. Gas Optimization
    • Uses assembly for memory allocation (_selectedPathTree).
    • Avoids storage writes during selection.

Use Cases

  1. Validator Selection (PoS blockchains)
  2. Lotteries with Weighted Tickets
  3. DAO Voting Power Sampling

Conclusion

The TreeMapUpgradeable contract provides a gas-efficient, fair, and scalable way to perform stake-based weighted selection in Solidity. By using a binary tree with aggregated weights, it ensures:

  • Proportional selection (higher stakes = higher chances).
  • Dynamic updates (stakes can change without restructuring).
  • Low gas costs (optimized storage and memory usage).

This makes it ideal for decentralized systems where fairness and efficiency are critical.