p2pool/util/forest.py

   1 '''
   2 forest data structure
   3 '''
   4
   5 import itertools
   6 import weakref
   7
   8 from p2pool.util import skiplist, variable
   9 from p2pool.bitcoin import data as bitcoin_data
  10
  11
  12 class TrackerSkipList(skiplist.SkipList):
  13     def __init__(self, tracker):
  14         skiplist.SkipList.__init__(self)
  15         self.tracker = tracker
  16
  17         self_ref = weakref.ref(self, lambda _: tracker.removed.unwatch(watch_id))
  18         watch_id = self.tracker.removed.watch(lambda share: self_ref().forget_item(share.hash))
  19
  20     def previous(self, element):
  21         return self.tracker.shares[element].previous_hash
  22
  23
  24 class DistanceSkipList(TrackerSkipList):
  25     def get_delta(self, element):
  26         return element, 1, self.tracker.shares[element].previous_hash
  27
  28     def combine_deltas(self, (from_hash1, dist1, to_hash1), (from_hash2, dist2, to_hash2)):
  29         if to_hash1 != from_hash2:
  30             raise AssertionError()
  31         return from_hash1, dist1 + dist2, to_hash2
  32
  33     def initial_solution(self, start, (n,)):
  34         return 0, start
  35
  36     def apply_delta(self, (dist1, to_hash1), (from_hash2, dist2, to_hash2), (n,)):
  37         if to_hash1 != from_hash2:
  38             raise AssertionError()
  39         return dist1 + dist2, to_hash2
  40
  41     def judge(self, (dist, hash), (n,)):
  42         if dist > n:
  43             return 1
  44         elif dist == n:
  45             return 0
  46         else:
  47             return -1
  48
  49     def finalize(self, (dist, hash), (n,)):
  50         assert dist == n
  51         return hash
  52
  53
  54 class AttributeDelta(object):
  55     __slots__ = 'height work'.split(' ')
  56
  57     @classmethod
  58     def get_none(cls):
  59         return cls(0, 0)
  60
  61     @classmethod
  62     def from_element(cls, share):
  63         return cls(1, bitcoin_data.target_to_average_attempts(share.target))
  64
  65     def __init__(self, height, work):
  66         self.height, self.work = height, work
  67
  68     def __add__(self, other):
  69         return AttributeDelta(self.height + other.height, self.work + other.work)
  70
  71     def __sub__(self, other):
  72         return AttributeDelta(self.height - other.height, self.work - other.work)
  73
  74 class Tracker(object):
  75     def __init__(self, shares=[], delta_type=AttributeDelta):
  76         self.shares = {} # hash -> share
  77         self.reverse_shares = {} # previous_hash -> set of share_hashes
  78
  79         self.heads = {} # head hash -> tail_hash
  80         self.tails = {} # tail hash -> set of head hashes
  81
  82         self.deltas = {} # share_hash -> delta, ref
  83         self.reverse_deltas = {} # ref -> set of share_hashes
  84
  85         self.ref_generator = itertools.count()
  86         self.delta_refs = {} # ref -> delta, share_hash
  87         self.reverse_delta_refs = {} # share_hash -> ref
  88
  89         self.added = variable.Event()
  90         self.removed = variable.Event()
  91
  92         self.get_nth_parent_hash = DistanceSkipList(self)
  93
  94         self.delta_type = delta_type
  95
  96         for share in shares:
  97             self.add(share)
  98
  99     def add(self, share):
 100         assert not isinstance(share, (int, long, type(None)))
 101         if share.hash in self.shares:
 102             raise ValueError('share already present')
 103
 104         if share.hash in self.tails:
 105             heads = self.tails.pop(share.hash)
 106         else:
 107             heads = set([share.hash])
 108
 109         if share.previous_hash in self.heads:
 110             tail = self.heads.pop(share.previous_hash)
 111         else:
 112             tail = self.get_last(share.previous_hash)
 113
 114         self.shares[share.hash] = share
 115         self.reverse_shares.setdefault(share.previous_hash, set()).add(share.hash)
 116
 117         self.tails.setdefault(tail, set()).update(heads)
 118         if share.previous_hash in self.tails[tail]:
 119             self.tails[tail].remove(share.previous_hash)
 120
 121         for head in heads:
 122             self.heads[head] = tail
 123
 124         self.added.happened(share)
 125
 126     def remove(self, share_hash):
 127         assert isinstance(share_hash, (int, long, type(None)))
 128         if share_hash not in self.shares:
 129             raise KeyError()
 130
 131         share = self.shares[share_hash]
 132         del share_hash
 133
 134         children = self.reverse_shares.get(share.hash, set())
 135
 136         if share.hash in self.heads and share.previous_hash in self.tails:
 137             tail = self.heads.pop(share.hash)
 138             self.tails[tail].remove(share.hash)
 139             if not self.tails[share.previous_hash]:
 140                 self.tails.pop(share.previous_hash)
 141         elif share.hash in self.heads:
 142             tail = self.heads.pop(share.hash)
 143             self.tails[tail].remove(share.hash)
 144             if self.reverse_shares[share.previous_hash] != set([share.hash]):
 145                 pass # has sibling
 146             else:
 147                 self.tails[tail].add(share.previous_hash)
 148                 self.heads[share.previous_hash] = tail
 149         elif share.previous_hash in self.tails and len(self.reverse_shares[share.previous_hash]) <= 1:
 150             # move delta refs referencing children down to this, so they can be moved up in one step
 151             if share.previous_hash in self.reverse_delta_refs:
 152                 for x in list(self.reverse_deltas.get(self.reverse_delta_refs.get(share.hash, object()), set())):
 153                     self.get_last(x)
 154                 assert share.hash not in self.reverse_delta_refs, list(self.reverse_deltas.get(self.reverse_delta_refs.get(share.hash, None), set()))
 155
 156             heads = self.tails.pop(share.previous_hash)
 157             for head in heads:
 158                 self.heads[head] = share.hash
 159             self.tails[share.hash] = set(heads)
 160
 161             # move ref pointing to this up
 162             if share.previous_hash in self.reverse_delta_refs:
 163                 assert share.hash not in self.reverse_delta_refs, list(self.reverse_deltas.get(self.reverse_delta_refs.get(share.hash, object()), set()))
 164
 165                 ref = self.reverse_delta_refs[share.previous_hash]
 166                 cur_delta, cur_hash = self.delta_refs[ref]
 167                 assert cur_hash == share.previous_hash
 168                 self.delta_refs[ref] = cur_delta - self.delta_type.from_element(share), share.hash
 169                 del self.reverse_delta_refs[share.previous_hash]
 170                 self.reverse_delta_refs[share.hash] = ref
 171         else:
 172             raise NotImplementedError()
 173
 174         # delete delta entry and ref if it is empty
 175         if share.hash in self.deltas:
 176             delta1, ref = self.deltas.pop(share.hash)
 177             self.reverse_deltas[ref].remove(share.hash)
 178             if not self.reverse_deltas[ref]:
 179                 del self.reverse_deltas[ref]
 180                 delta2, ref_hash = self.delta_refs.pop(ref)
 181                 del self.reverse_delta_refs[ref_hash]
 182
 183         self.shares.pop(share.hash)
 184         self.reverse_shares[share.previous_hash].remove(share.hash)
 185         if not self.reverse_shares[share.previous_hash]:
 186             self.reverse_shares.pop(share.previous_hash)
 187
 188         self.removed.happened(share)
 189
 190     def get_height(self, share_hash):
 191         delta, last = self.get_delta(share_hash)
 192         return delta.height
 193
 194     def get_work(self, share_hash):
 195         delta, last = self.get_delta(share_hash)
 196         return delta.work
 197
 198     def get_last(self, share_hash):
 199         delta, last = self.get_delta(share_hash)
 200         return last
 201
 202     def get_height_and_last(self, share_hash):
 203         delta, last = self.get_delta(share_hash)
 204         return delta.height, last
 205
 206     def get_height_work_and_last(self, share_hash):
 207         delta, last = self.get_delta(share_hash)
 208         return delta.height, delta.work, last
 209
 210     def _get_delta(self, share_hash):
 211         if share_hash in self.deltas:
 212             delta1, ref = self.deltas[share_hash]
 213             delta2, share_hash = self.delta_refs[ref]
 214             return delta1 + delta2, share_hash
 215         else:
 216             return self.delta_type.from_element(self.shares[share_hash]), self.shares[share_hash].previous_hash
 217
 218     def _set_delta(self, share_hash, delta, other_share_hash):
 219         if other_share_hash not in self.reverse_delta_refs:
 220             ref = self.ref_generator.next()
 221             assert ref not in self.delta_refs
 222             self.delta_refs[ref] = self.delta_type.get_none(), other_share_hash
 223             self.reverse_delta_refs[other_share_hash] = ref
 224             del ref
 225
 226         ref = self.reverse_delta_refs[other_share_hash]
 227         ref_delta, ref_share_hash = self.delta_refs[ref]
 228         assert ref_share_hash == other_share_hash
 229
 230         if share_hash in self.deltas:
 231             prev_ref = self.deltas[share_hash][1]
 232             self.reverse_deltas[prev_ref].remove(share_hash)
 233             if not self.reverse_deltas[prev_ref] and prev_ref != ref:
 234                 self.reverse_deltas.pop(prev_ref)
 235                 _, x = self.delta_refs.pop(prev_ref)
 236                 self.reverse_delta_refs.pop(x)
 237         self.deltas[share_hash] = delta - ref_delta, ref
 238         self.reverse_deltas.setdefault(ref, set()).add(share_hash)
 239
 240     def get_delta(self, share_hash):
 241         assert isinstance(share_hash, (int, long, type(None)))
 242         delta = self.delta_type.get_none()
 243         updates = []
 244         while share_hash in self.shares:
 245             updates.append((share_hash, delta))
 246             this_delta, share_hash = self._get_delta(share_hash)
 247             delta += this_delta
 248         for update_hash, delta_then in updates:
 249             self._set_delta(update_hash, delta - delta_then, share_hash)
 250         return delta, share_hash
 251
 252     def get_chain(self, start_hash, length):
 253         assert length <= self.get_height(start_hash)
 254         for i in xrange(length):
 255             yield self.shares[start_hash]
 256             start_hash = self.shares[start_hash].previous_hash
 257
 258     def is_child_of(self, share_hash, possible_child_hash):
 259         height, last = self.get_height_and_last(share_hash)
 260         child_height, child_last = self.get_height_and_last(possible_child_hash)
 261         if child_last != last:
 262             return None # not connected, so can't be determined
 263         height_up = child_height - height
 264         return height_up >= 0 and self.get_nth_parent_hash(possible_child_hash, height_up) == share_hash