diff options
Diffstat (limited to 'NohGooee/PiecePicker.py')
-rw-r--r-- | NohGooee/PiecePicker.py | 138 |
1 files changed, 138 insertions, 0 deletions
diff --git a/NohGooee/PiecePicker.py b/NohGooee/PiecePicker.py new file mode 100644 index 0000000..e22b13a --- /dev/null +++ b/NohGooee/PiecePicker.py @@ -0,0 +1,138 @@ +# The contents of this file are subject to the BitTorrent Open Source License +# Version 1.1 (the License). You may not copy or use this file, in either +# source code or executable form, except in compliance with the License. You +# may obtain a copy of the License at http://www.bittorrent.com/license/. +# +# Software distributed under the License is distributed on an AS IS basis, +# WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License +# for the specific language governing rights and limitations under the +# License. + +# Written by Bram Cohen + +from random import randrange, shuffle, choice + + +class PiecePicker(object): + + def __init__(self, numpieces, config): + self.config = config + self.numpieces = numpieces + self.interests = [range(numpieces)] + self.pos_in_interests = range(numpieces) + self.numinterests = [0] * numpieces + self.have = [False] * numpieces + self.crosscount = [numpieces] + self.started = [] + self.seedstarted = [] + self.numgot = 0 + self.scrambled = range(numpieces) + shuffle(self.scrambled) + + def got_have(self, piece): + numint = self.numinterests[piece] + self.crosscount[numint + self.have[piece]] -= 1 + self.numinterests[piece] += 1 + try: + self.crosscount[numint + 1 + self.have[piece]] += 1 + except IndexError: + self.crosscount.append(1) + if self.have[piece]: + return + if numint == len(self.interests) - 1: + self.interests.append([]) + self._shift_over(piece, self.interests[numint], self.interests[numint + 1]) + + def lost_have(self, piece): + numint = self.numinterests[piece] + self.crosscount[numint + self.have[piece]] -= 1 + self.numinterests[piece] -= 1 + self.crosscount[numint - 1 + self.have[piece]] += 1 + if self.have[piece]: + return + self._shift_over(piece, self.interests[numint], self.interests[numint - 1]) + + def _shift_over(self, piece, l1, l2): + p = self.pos_in_interests[piece] + l1[p] = l1[-1] + self.pos_in_interests[l1[-1]] = p + del l1[-1] + newp = randrange(len(l2) + 1) + if newp == len(l2): + self.pos_in_interests[piece] = len(l2) + l2.append(piece) + else: + old = l2[newp] + self.pos_in_interests[old] = len(l2) + l2.append(old) + l2[newp] = piece + self.pos_in_interests[piece] = newp + + def requested(self, piece, seed = False): + if piece not in self.started: + self.started.append(piece) + if seed and piece not in self.seedstarted: + self.seedstarted.append(piece) + + def complete(self, piece): + assert not self.have[piece] + self.have[piece] = True + self.crosscount[self.numinterests[piece]] -= 1 + try: + self.crosscount[self.numinterests[piece] + 1] += 1 + except IndexError: + self.crosscount.append(1) + self.numgot += 1 + l = self.interests[self.numinterests[piece]] + p = self.pos_in_interests[piece] + l[p] = l[-1] + self.pos_in_interests[l[-1]] = p + del l[-1] + try: + self.started.remove(piece) + self.seedstarted.remove(piece) + except ValueError: + pass + + def next(self, havefunc, seed = False): + bests = None + bestnum = 2 ** 30 + if seed: + s = self.seedstarted + else: + s = self.started + for i in s: + if havefunc(i): + if self.numinterests[i] < bestnum: + bests = [i] + bestnum = self.numinterests[i] + elif self.numinterests[i] == bestnum: + bests.append(i) + if bests: + return choice(bests) + if self.numgot < self.config['rarest_first_cutoff']: + for i in self.scrambled: + if havefunc(i): + return i + return None + for i in xrange(1, min(bestnum, len(self.interests))): + for j in self.interests[i]: + if havefunc(j): + return j + return None + + def am_I_complete(self): + return self.numgot == self.numpieces + + def bump(self, piece): + l = self.interests[self.numinterests[piece]] + pos = self.pos_in_interests[piece] + del l[pos] + l.append(piece) + for i in range(pos,len(l)): + self.pos_in_interests[l[i]] = i + try: + self.started.remove(piece) + self.seedstarted.remove(piece) + except ValueError: + pass |