PieceSetBuckets.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132
  1. # The contents of this file are subject to the BitTorrent Open Source License
  2. # Version 1.1 (the License). You may not copy or use this file, in either
  3. # source code or executable form, except in compliance with the License. You
  4. # may obtain a copy of the License at http://www.bittorrent.com/license/.
  5. #
  6. # Software distributed under the License is distributed on an AS IS basis,
  7. # WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
  8. # for the specific language governing rights and limitations under the
  9. # License.
  10. # Written by Greg Hazel
  11. ## up
  12. # p.add(piece, p.remove(piece) + 1)
  13. ## down
  14. # p.add(piece, p.remove(piece) - 1)
  15. from BTL.sparse_set import SparseSet
  16. class PieceSetBuckets(object):
  17. """A PieceBuckets object is an array of arrays. ith bucket contains
  18. pieces that have i known instances within the network. Pieces
  19. within each bucket are randomly ordered."""
  20. def __init__(self):
  21. # [SparseSet(piece)]
  22. self.buckets = []
  23. # {piece: (bucket)}
  24. self.place_in_buckets = {}
  25. def get_position(self, piece): # returns which bucket piece is in.
  26. return self.place_in_buckets[piece]
  27. def __contains__(self, piece):
  28. return piece in self.place_in_buckets
  29. def add(self, piece, bucketindex):
  30. assert not self.place_in_buckets.has_key(piece)
  31. while len(self.buckets) <= bucketindex:
  32. self.buckets.append(SparseSet())
  33. bucket = self.buckets[bucketindex]
  34. bucket.add(piece)
  35. self.place_in_buckets[piece] = bucketindex
  36. def remove(self, piece):
  37. bucketindex = self.place_in_buckets.pop(piece)
  38. bucket = self.buckets[bucketindex]
  39. bucket.subtract(piece)
  40. while len(self.buckets) > 0 and len(self.buckets[-1]) == 0:
  41. del self.buckets[-1]
  42. return bucketindex
  43. def prepend_bucket(self):
  44. # it' possible we had everything to begin with
  45. if len(self.buckets) == 0:
  46. return
  47. self.buckets.insert(0, SparseSet())
  48. # bleh.
  49. for piece in self.place_in_buckets:
  50. self.place_in_buckets[piece] += 1
  51. def popleft_bucket(self):
  52. # it' possible we had everything to begin with
  53. if len(self.buckets) == 0:
  54. return
  55. self.buckets.pop(0)
  56. # bleh.
  57. for piece in self.place_in_buckets:
  58. self.place_in_buckets[piece] -= 1
  59. import array
  60. import bisect
  61. def resolve_typecode(n):
  62. if n < 32768:
  63. return 'h'
  64. return 'l'
  65. class SortedPieceBuckets(object):
  66. """A PieceBuckets object is an array of arrays. ith bucket contains
  67. pieces that have i known instances within the network. Pieces
  68. within each bucket are randomly ordered."""
  69. def __init__(self, typecode):
  70. self.typecode = typecode
  71. # [[piece]]
  72. self.buckets = []
  73. # {piece: (bucket, bucketpos)}
  74. self.place_in_buckets = {}
  75. def get_position(self, piece): # returns which bucket piece is in.
  76. return self.place_in_buckets[piece][0]
  77. def __contains__(self, piece):
  78. return piece in self.place_in_buckets
  79. def add(self, piece, bucketindex):
  80. assert not self.place_in_buckets.has_key(piece)
  81. while len(self.buckets) <= bucketindex:
  82. self.buckets.append(array.array(self.typecode))
  83. bucket = self.buckets[bucketindex]
  84. newspot = bisect.bisect_right(bucket, piece)
  85. bucket.insert(newspot, piece)
  86. self.place_in_buckets[piece] = (bucketindex, newspot)
  87. def remove(self, piece):
  88. bucketindex, bucketpos = self.place_in_buckets.pop(piece)
  89. bucket = self.buckets[bucketindex]
  90. newspot = bisect.bisect_left(bucket, piece)
  91. del bucket[newspot]
  92. while len(self.buckets) > 0 and len(self.buckets[-1]) == 0:
  93. del self.buckets[-1]
  94. return bucketindex
  95. def prepend_bucket(self):
  96. # it' possible we had everything to begin with
  97. if len(self.buckets) == 0:
  98. return
  99. self.buckets.insert(0, array.array(self.typecode))
  100. # bleh.
  101. for piece in self.place_in_buckets:
  102. index, pos = self.place_in_buckets[piece]
  103. self.place_in_buckets[piece] = (index + 1, pos)
  104. def popleft_bucket(self):
  105. # it' possible we had everything to begin with
  106. if len(self.buckets) == 0:
  107. return
  108. self.buckets.pop(0)
  109. # bleh.
  110. for piece in self.place_in_buckets:
  111. index, pos = self.place_in_buckets[piece]
  112. self.place_in_buckets[piece] = (index - 1, pos)