|
| 1 | +// Copyright 2016 The go-ethereum Authors |
| 2 | +// This file is part of the go-ethereum library. |
| 3 | +// |
| 4 | +// The go-ethereum library is free software: you can redistribute it and/or modify |
| 5 | +// it under the terms of the GNU Lesser General Public License as published by |
| 6 | +// the Free Software Foundation, either version 3 of the License, or |
| 7 | +// (at your option) any later version. |
| 8 | +// |
| 9 | +// The go-ethereum library is distributed in the hope that it will be useful, |
| 10 | +// but WITHOUT ANY WARRANTY; without even the implied warranty of |
| 11 | +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the |
| 12 | +// GNU Lesser General Public License for more details. |
| 13 | +// |
| 14 | +// You should have received a copy of the GNU Lesser General Public License |
| 15 | +// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>. |
| 16 | + |
| 17 | +// Package les implements the Light Ethereum Subprotocol. |
| 18 | +package les |
| 19 | + |
| 20 | +import ( |
| 21 | + "math/rand" |
| 22 | +) |
| 23 | + |
| 24 | +// wrsItem interface should be implemented by any entries that are to be selected from |
| 25 | +// a weightedRandomSelect set. Note that recalculating monotonously decreasing item |
| 26 | +// weights on-demand (without constantly calling update) is allowed |
| 27 | +type wrsItem interface { |
| 28 | + Weight() int64 |
| 29 | +} |
| 30 | + |
| 31 | +// weightedRandomSelect is capable of weighted random selection from a set of items |
| 32 | +type weightedRandomSelect struct { |
| 33 | + root *wrsNode |
| 34 | + idx map[wrsItem]int |
| 35 | +} |
| 36 | + |
| 37 | +// newWeightedRandomSelect returns a new weightedRandomSelect structure |
| 38 | +func newWeightedRandomSelect() *weightedRandomSelect { |
| 39 | + return &weightedRandomSelect{root: &wrsNode{maxItems: wrsBranches}, idx: make(map[wrsItem]int)} |
| 40 | +} |
| 41 | + |
| 42 | +// update updates an item's weight, adds it if it was non-existent or removes it if |
| 43 | +// the new weight is zero. Note that explicitly updating decreasing weights is not necessary. |
| 44 | +func (w *weightedRandomSelect) update(item wrsItem) { |
| 45 | + w.setWeight(item, item.Weight()) |
| 46 | +} |
| 47 | + |
| 48 | +// remove removes an item from the set |
| 49 | +func (w *weightedRandomSelect) remove(item wrsItem) { |
| 50 | + w.setWeight(item, 0) |
| 51 | +} |
| 52 | + |
| 53 | +// setWeight sets an item's weight to a specific value (removes it if zero) |
| 54 | +func (w *weightedRandomSelect) setWeight(item wrsItem, weight int64) { |
| 55 | + idx, ok := w.idx[item] |
| 56 | + if ok { |
| 57 | + w.root.setWeight(idx, weight) |
| 58 | + if weight == 0 { |
| 59 | + delete(w.idx, item) |
| 60 | + } |
| 61 | + } else { |
| 62 | + if weight != 0 { |
| 63 | + if w.root.itemCnt == w.root.maxItems { |
| 64 | + // add a new level |
| 65 | + newRoot := &wrsNode{sumWeight: w.root.sumWeight, itemCnt: w.root.itemCnt, level: w.root.level + 1, maxItems: w.root.maxItems * wrsBranches} |
| 66 | + newRoot.items[0] = w.root |
| 67 | + newRoot.weights[0] = w.root.sumWeight |
| 68 | + w.root = newRoot |
| 69 | + } |
| 70 | + w.idx[item] = w.root.insert(item, weight) |
| 71 | + } |
| 72 | + } |
| 73 | +} |
| 74 | + |
| 75 | +// choose randomly selects an item from the set, with a chance proportional to its |
| 76 | +// current weight. If the weight of the chosen element has been decreased since the |
| 77 | +// last stored value, returns it with a newWeight/oldWeight chance, otherwise just |
| 78 | +// updates its weight and selects another one |
| 79 | +func (w *weightedRandomSelect) choose() wrsItem { |
| 80 | + for { |
| 81 | + if w.root.sumWeight == 0 { |
| 82 | + return nil |
| 83 | + } |
| 84 | + val := rand.Int63n(w.root.sumWeight) |
| 85 | + choice, lastWeight := w.root.choose(val) |
| 86 | + weight := choice.Weight() |
| 87 | + if weight != lastWeight { |
| 88 | + w.setWeight(choice, weight) |
| 89 | + } |
| 90 | + if weight >= lastWeight || rand.Int63n(lastWeight) < weight { |
| 91 | + return choice |
| 92 | + } |
| 93 | + } |
| 94 | +} |
| 95 | + |
| 96 | +const wrsBranches = 8 // max number of branches in the wrsNode tree |
| 97 | + |
| 98 | +// wrsNode is a node of a tree structure that can store wrsItems or further wrsNodes. |
| 99 | +type wrsNode struct { |
| 100 | + items [wrsBranches]interface{} |
| 101 | + weights [wrsBranches]int64 |
| 102 | + sumWeight int64 |
| 103 | + level, itemCnt, maxItems int |
| 104 | +} |
| 105 | + |
| 106 | +// insert recursively inserts a new item to the tree and returns the item index |
| 107 | +func (n *wrsNode) insert(item wrsItem, weight int64) int { |
| 108 | + branch := 0 |
| 109 | + for n.items[branch] != nil && (n.level == 0 || n.items[branch].(*wrsNode).itemCnt == n.items[branch].(*wrsNode).maxItems) { |
| 110 | + branch++ |
| 111 | + if branch == wrsBranches { |
| 112 | + panic(nil) |
| 113 | + } |
| 114 | + } |
| 115 | + n.itemCnt++ |
| 116 | + n.sumWeight += weight |
| 117 | + n.weights[branch] += weight |
| 118 | + if n.level == 0 { |
| 119 | + n.items[branch] = item |
| 120 | + return branch |
| 121 | + } else { |
| 122 | + var subNode *wrsNode |
| 123 | + if n.items[branch] == nil { |
| 124 | + subNode = &wrsNode{maxItems: n.maxItems / wrsBranches, level: n.level - 1} |
| 125 | + n.items[branch] = subNode |
| 126 | + } else { |
| 127 | + subNode = n.items[branch].(*wrsNode) |
| 128 | + } |
| 129 | + subIdx := subNode.insert(item, weight) |
| 130 | + return subNode.maxItems*branch + subIdx |
| 131 | + } |
| 132 | +} |
| 133 | + |
| 134 | +// setWeight updates the weight of a certain item (which should exist) and returns |
| 135 | +// the change of the last weight value stored in the tree |
| 136 | +func (n *wrsNode) setWeight(idx int, weight int64) int64 { |
| 137 | + if n.level == 0 { |
| 138 | + oldWeight := n.weights[idx] |
| 139 | + n.weights[idx] = weight |
| 140 | + diff := weight - oldWeight |
| 141 | + n.sumWeight += diff |
| 142 | + if weight == 0 { |
| 143 | + n.items[idx] = nil |
| 144 | + n.itemCnt-- |
| 145 | + } |
| 146 | + return diff |
| 147 | + } |
| 148 | + branchItems := n.maxItems / wrsBranches |
| 149 | + branch := idx / branchItems |
| 150 | + diff := n.items[branch].(*wrsNode).setWeight(idx-branch*branchItems, weight) |
| 151 | + n.weights[branch] += diff |
| 152 | + n.sumWeight += diff |
| 153 | + if weight == 0 { |
| 154 | + n.itemCnt-- |
| 155 | + } |
| 156 | + return diff |
| 157 | +} |
| 158 | + |
| 159 | +// choose recursively selects an item from the tree and returns it along with its weight |
| 160 | +func (n *wrsNode) choose(val int64) (wrsItem, int64) { |
| 161 | + for i, w := range n.weights { |
| 162 | + if val < w { |
| 163 | + if n.level == 0 { |
| 164 | + return n.items[i].(wrsItem), n.weights[i] |
| 165 | + } else { |
| 166 | + return n.items[i].(*wrsNode).choose(val) |
| 167 | + } |
| 168 | + } else { |
| 169 | + val -= w |
| 170 | + } |
| 171 | + } |
| 172 | + panic(nil) |
| 173 | +} |
0 commit comments