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
  10
  11 class TrackerSkipList(skiplist.SkipList):
  12     def __init__(self, tracker):
  13         skiplist.SkipList.__init__(self)
  14         self.tracker = tracker
  15
  16         self_ref = weakref.ref(self, lambda _: tracker.removed.unwatch(watch_id))
  17         watch_id = self.tracker.removed.watch(lambda item: self_ref().forget_item(item.hash))
  18
  19     def previous(self, element):
  20         return self.tracker._delta_type.from_element(self.tracker.items[element]).tail
  21
  22
  23 class DistanceSkipList(TrackerSkipList):
  24     def get_delta(self, element):
  25         return element, 1, self.previous(element)
  26
  27     def combine_deltas(self, (from_hash1, dist1, to_hash1), (from_hash2, dist2, to_hash2)):
  28         if to_hash1 != from_hash2:
  29             raise AssertionError()
  30         return from_hash1, dist1 + dist2, to_hash2
  31
  32     def initial_solution(self, start, (n,)):
  33         return 0, start
  34
  35     def apply_delta(self, (dist1, to_hash1), (from_hash2, dist2, to_hash2), (n,)):
  36         if to_hash1 != from_hash2:
  37             raise AssertionError()
  38         return dist1 + dist2, to_hash2
  39
  40     def judge(self, (dist, hash), (n,)):
  41         if dist > n:
  42             return 1
  43         elif dist == n:
  44             return 0
  45         else:
  46             return -1
  47
  48     def finalize(self, (dist, hash), (n,)):
  49         assert dist == n
  50         return hash
  51
  52 def get_attributedelta_type(attrs): # attrs: {name: func}
  53     class ProtoAttributeDelta(object):
  54         __slots__ = ['head', 'tail'] + attrs.keys()
  55
  56         @classmethod
  57         def get_none(cls, element_id):
  58             return cls(element_id, element_id, **dict((k, 0) for k in attrs))
  59
  60         @classmethod
  61         def from_element(cls, item):
  62             return cls(item.hash, item.previous_hash, **dict((k, v(item)) for k, v in attrs.iteritems()))
  63
  64         def __init__(self, head, tail, **kwargs):
  65             self.head, self.tail = head, tail
  66             for k, v in kwargs.iteritems():
  67                 setattr(self, k, v)
  68
  69         def __add__(self, other):
  70             assert self.tail == other.head
  71             return self.__class__(self.head, other.tail, **dict((k, getattr(self, k) + getattr(other, k)) for k in attrs))
  72
  73         def __sub__(self, other):
  74             if self.head == other.head:
  75                 return self.__class__(other.tail, self.tail, **dict((k, getattr(self, k) - getattr(other, k)) for k in attrs))
  76             elif self.tail == other.tail:
  77                 return self.__class__(self.head, other.head, **dict((k, getattr(self, k) - getattr(other, k)) for k in attrs))
  78             else:
  79                 raise AssertionError()
  80
  81         def __repr__(self):
  82             return '%s(%r, %r%s)' % (self.__class__, self.head, self.tail, ''.join(', %s=%r' % (k, getattr(self, k)) for k in attrs))
  83     ProtoAttributeDelta.attrs = attrs
  84     return ProtoAttributeDelta
  85
  86 AttributeDelta = get_attributedelta_type(dict(
  87     height=lambda item: 1,
  88 ))
  89
  90 class Tracker(object):
  91     def __init__(self, items=[], delta_type=AttributeDelta):
  92         self.items = {} # hash -> item
  93         self.reverse = {} # delta.tail -> set of item_hashes
  94
  95         self.heads = {} # head hash -> tail_hash
  96         self.tails = {} # tail hash -> set of head hashes
  97
  98         self._deltas = {} # item_hash -> delta, ref
  99         self._reverse_deltas = {} # ref -> set of item_hashes
 100
 101         self._ref_generator = itertools.count()
 102         self._delta_refs = {} # ref -> delta
 103         self._reverse_delta_refs = {} # delta.tail -> ref
 104
 105         self.added = variable.Event()
 106         self.removed = variable.Event()
 107
 108         self.get_nth_parent_hash = DistanceSkipList(self)
 109
 110         self._delta_type = delta_type
 111
 112         for item in items:
 113             self.add(item)
 114
 115     def add(self, item):
 116         assert not isinstance(item, (int, long, type(None)))
 117         delta = self._delta_type.from_element(item)
 118
 119         if delta.head in self.items:
 120             raise ValueError('item already present')
 121
 122         if delta.head in self.tails:
 123             heads = self.tails.pop(delta.head)
 124         else:
 125             heads = set([delta.head])
 126
 127         if delta.tail in self.heads:
 128             tail = self.heads.pop(delta.tail)
 129         else:
 130             tail = self.get_last(delta.tail)
 131
 132         self.items[delta.head] = item
 133         self.reverse.setdefault(delta.tail, set()).add(delta.head)
 134
 135         self.tails.setdefault(tail, set()).update(heads)
 136         if delta.tail in self.tails[tail]:
 137             self.tails[tail].remove(delta.tail)
 138
 139         for head in heads:
 140             self.heads[head] = tail
 141
 142         self.added.happened(item)
 143
 144     def remove(self, item_hash):
 145         assert isinstance(item_hash, (int, long, type(None)))
 146         if item_hash not in self.items:
 147             raise KeyError()
 148
 149         item = self.items[item_hash]
 150         del item_hash
 151
 152         delta = self._delta_type.from_element(item)
 153
 154         children = self.reverse.get(delta.head, set())
 155
 156         if delta.head in self.heads and delta.tail in self.tails:
 157             tail = self.heads.pop(delta.head)
 158             self.tails[tail].remove(delta.head)
 159             if not self.tails[delta.tail]:
 160                 self.tails.pop(delta.tail)
 161         elif delta.head in self.heads:
 162             tail = self.heads.pop(delta.head)
 163             self.tails[tail].remove(delta.head)
 164             if self.reverse[delta.tail] != set([delta.head]):
 165                 pass # has sibling
 166             else:
 167                 self.tails[tail].add(delta.tail)
 168                 self.heads[delta.tail] = tail
 169         elif delta.tail in self.tails and len(self.reverse[delta.tail]) <= 1:
 170             # move delta refs referencing children down to this, so they can be moved up in one step
 171             if delta.tail in self._reverse_delta_refs:
 172                 for x in list(self._reverse_deltas.get(self._reverse_delta_refs.get(delta.head, object()), set())):
 173                     self.get_last(x)
 174                 assert delta.head not in self._reverse_delta_refs, list(self._reverse_deltas.get(self._reverse_delta_refs.get(delta.head, None), set()))
 175
 176             heads = self.tails.pop(delta.tail)
 177             for head in heads:
 178                 self.heads[head] = delta.head
 179             self.tails[delta.head] = set(heads)
 180
 181             # move ref pointing to this up
 182             if delta.tail in self._reverse_delta_refs:
 183                 assert delta.head not in self._reverse_delta_refs, list(self._reverse_deltas.get(self._reverse_delta_refs.get(delta.head, object()), set()))
 184
 185                 ref = self._reverse_delta_refs[delta.tail]
 186                 cur_delta = self._delta_refs[ref]
 187                 assert cur_delta.tail == delta.tail
 188                 self._delta_refs[ref] = cur_delta - self._delta_type.from_element(item)
 189                 assert self._delta_refs[ref].tail == delta.head
 190                 del self._reverse_delta_refs[delta.tail]
 191                 self._reverse_delta_refs[delta.head] = ref
 192         else:
 193             raise NotImplementedError()
 194
 195         # delete delta entry and ref if it is empty
 196         if delta.head in self._deltas:
 197             delta1, ref = self._deltas.pop(delta.head)
 198             self._reverse_deltas[ref].remove(delta.head)
 199             if not self._reverse_deltas[ref]:
 200                 del self._reverse_deltas[ref]
 201                 delta2 = self._delta_refs.pop(ref)
 202                 del self._reverse_delta_refs[delta2.tail]
 203
 204         self.items.pop(delta.head)
 205         self.reverse[delta.tail].remove(delta.head)
 206         if not self.reverse[delta.tail]:
 207             self.reverse.pop(delta.tail)
 208
 209         self.removed.happened(item)
 210
 211     def get_height(self, item_hash):
 212         return self.get_delta_to_last(item_hash).height
 213
 214     def get_work(self, item_hash):
 215         return self.get_delta_to_last(item_hash).work
 216
 217     def get_last(self, item_hash):
 218         return self.get_delta_to_last(item_hash).tail
 219
 220     def get_height_and_last(self, item_hash):
 221         delta = self.get_delta_to_last(item_hash)
 222         return delta.height, delta.tail
 223
 224     def _get_delta(self, item_hash):
 225         if item_hash in self._deltas:
 226             delta1, ref = self._deltas[item_hash]
 227             delta2 = self._delta_refs[ref]
 228             res = delta1 + delta2
 229         else:
 230             res = self._delta_type.from_element(self.items[item_hash])
 231         assert res.head == item_hash
 232         return res
 233
 234     def _set_delta(self, item_hash, delta):
 235         other_item_hash = delta.tail
 236         if other_item_hash not in self._reverse_delta_refs:
 237             ref = self._ref_generator.next()
 238             assert ref not in self._delta_refs
 239             self._delta_refs[ref] = self._delta_type.get_none(other_item_hash)
 240             self._reverse_delta_refs[other_item_hash] = ref
 241             del ref
 242
 243         ref = self._reverse_delta_refs[other_item_hash]
 244         ref_delta = self._delta_refs[ref]
 245         assert ref_delta.tail == other_item_hash
 246
 247         if item_hash in self._deltas:
 248             prev_ref = self._deltas[item_hash][1]
 249             self._reverse_deltas[prev_ref].remove(item_hash)
 250             if not self._reverse_deltas[prev_ref] and prev_ref != ref:
 251                 self._reverse_deltas.pop(prev_ref)
 252                 x = self._delta_refs.pop(prev_ref)
 253                 self._reverse_delta_refs.pop(x.tail)
 254         self._deltas[item_hash] = delta - ref_delta, ref
 255         self._reverse_deltas.setdefault(ref, set()).add(item_hash)
 256
 257     def get_delta_to_last(self, item_hash):
 258         assert isinstance(item_hash, (int, long, type(None)))
 259         delta = self._delta_type.get_none(item_hash)
 260         updates = []
 261         while delta.tail in self.items:
 262             updates.append((delta.tail, delta))
 263             this_delta = self._get_delta(delta.tail)
 264             delta += this_delta
 265         for update_hash, delta_then in updates:
 266             self._set_delta(update_hash, delta - delta_then)
 267         return delta
 268
 269     def get_delta(self, item, ancestor):
 270         assert self.is_child_of(ancestor, item)
 271         return self.get_delta_to_last(item) - self.get_delta_to_last(ancestor)
 272
 273     def get_chain(self, start_hash, length):
 274         assert length <= self.get_height(start_hash)
 275         for i in xrange(length):
 276             yield self.items[start_hash]
 277             start_hash = self._delta_type.from_element(self.items[start_hash]).tail
 278
 279     def is_child_of(self, item_hash, possible_child_hash):
 280         height, last = self.get_height_and_last(item_hash)
 281         child_height, child_last = self.get_height_and_last(possible_child_hash)
 282         if child_last != last:
 283             return None # not connected, so can't be determined
 284         height_up = child_height - height
 285         return height_up >= 0 and self.get_nth_parent_hash(possible_child_hash, height_up) == item_hash
 286
 287 class SubsetTracker(Tracker):
 288     def __init__(self, subset_of, **kwargs):
 289         Tracker.__init__(self, **kwargs)
 290         self.get_nth_parent_hash = subset_of.get_nth_parent_hash # overwrites Tracker.__init__'s
 291         self._subset_of = subset_of
 292
 293     def add(self, item):
 294         delta = self._delta_type.from_element(item)
 295         if self._subset_of is not None:
 296             assert delta.head in self._subset_of.items
 297         Tracker.add(self, item)
 298
 299     def remove(self, item_hash):
 300         if self._subset_of is not None:
 301             assert item_hash in self._subset_of.items
 302         Tracker.remove(self, item_hash)