← Back to branch summary

~ubuntu-branches/ubuntu/intrepid/miro/intrepid

« back to all changes in this revision

Viewing changes to portable/BitTorrent/PiecePicker.py

  • Committer: Bazaar Package Importer
  • Author(s): Christopher James Halse Rogers
  • Date: 2008-02-09 13:37:10 UTC
  • mfrom: (1.1.2 upstream)
  • Revision ID: james.westby@ubuntu.com-20080209133710-9rs90q6gckvp1b6i
Tags: 1.1.2-0ubuntu1
New upstream release

Show diffs side-by-side

added added

removed removed

Lines of Context:
portable/BitTorrent/PiecePicker.py
1
 
# Written by Bram Cohen
2
 
# see LICENSE.txt for license information
3
 
 
4
 
from random import randrange, shuffle, choice
5
 
 
6
 
class PiecePicker:
7
 
    def __init__(self, numpieces, rarest_first_cutoff = 1):
8
 
        self.rarest_first_cutoff = rarest_first_cutoff
9
 
        self.numpieces = numpieces
10
 
        self.interests = [range(numpieces)]
11
 
        self.pos_in_interests = range(numpieces)
12
 
        self.numinterests = [0] * numpieces
13
 
        self.started = []
14
 
        self.seedstarted = []
15
 
        self.numgot = 0
16
 
        self.scrambled = range(numpieces)
17
 
        shuffle(self.scrambled)
18
 
 
19
 
    def got_have(self, piece):
20
 
        if self.numinterests[piece] is None:
21
 
            return
22
 
        numint = self.numinterests[piece]
23
 
        if numint == len(self.interests) - 1:
24
 
            self.interests.append([])
25
 
        self.numinterests[piece] += 1
26
 
        self._shift_over(piece, self.interests[numint], self.interests[numint + 1])
27
 
 
28
 
    def lost_have(self, piece):
29
 
        if self.numinterests[piece] is None:
30
 
            return
31
 
        numint = self.numinterests[piece]
32
 
        self.numinterests[piece] -= 1
33
 
        self._shift_over(piece, self.interests[numint], self.interests[numint - 1])
34
 
 
35
 
    def _shift_over(self, piece, l1, l2):
36
 
        p = self.pos_in_interests[piece]
37
 
        l1[p] = l1[-1]
38
 
        self.pos_in_interests[l1[-1]] = p
39
 
        del l1[-1]
40
 
        newp = randrange(len(l2) + 1)
41
 
        if newp == len(l2):
42
 
            self.pos_in_interests[piece] = len(l2)
43
 
            l2.append(piece)
44
 
        else:
45
 
            old = l2[newp]
46
 
            self.pos_in_interests[old] = len(l2)
47
 
            l2.append(old)
48
 
            l2[newp] = piece
49
 
            self.pos_in_interests[piece] = newp
50
 
 
51
 
    def requested(self, piece, seed = False):
52
 
        if piece not in self.started:
53
 
            self.started.append(piece)
54
 
        if seed and piece not in self.seedstarted:
55
 
            self.seedstarted.append(piece)
56
 
 
57
 
    def complete(self, piece):
58
 
        assert self.numinterests[piece] is not None
59
 
        self.numgot += 1
60
 
        l = self.interests[self.numinterests[piece]]
61
 
        p = self.pos_in_interests[piece]
62
 
        l[p] = l[-1]
63
 
        self.pos_in_interests[l[-1]] = p
64
 
        del l[-1]
65
 
        self.numinterests[piece] = None
66
 
        try:
67
 
            self.started.remove(piece)
68
 
            self.seedstarted.remove(piece)
69
 
        except ValueError:
70
 
            pass
71
 
 
72
 
    def next(self, havefunc, seed = False):
73
 
        bests = None
74
 
        bestnum = 2 ** 30
75
 
        if seed:
76
 
            s = self.seedstarted
77
 
        else:
78
 
            s = self.started
79
 
        for i in s:
80
 
            if havefunc(i):
81
 
                if self.numinterests[i] < bestnum:
82
 
                    bests = [i]
83
 
                    bestnum = self.numinterests[i]
84
 
                elif self.numinterests[i] == bestnum:
85
 
                    bests.append(i)
86
 
        if bests:
87
 
            return choice(bests)
88
 
        if self.numgot < self.rarest_first_cutoff:
89
 
            for i in self.scrambled:
90
 
                if havefunc(i):
91
 
                    return i
92
 
            return None
93
 
        for i in xrange(1, min(bestnum, len(self.interests))):
94
 
            for j in self.interests[i]:
95
 
                if havefunc(j):
96
 
                    return j
97
 
        return None
98
 
 
99
 
    def am_I_complete(self):
100
 
        return self.numgot == self.numpieces
101
 
 
102
 
    def bump(self, piece):
103
 
        l = self.interests[self.numinterests[piece]]
104
 
        pos = self.pos_in_interests[piece]
105
 
        del l[pos]
106
 
        l.append(piece)
107
 
        for i in range(pos,len(l)):
108
 
            self.pos_in_interests[l[i]] = i
109
 
 
110
 
def test_requested():
111
 
    p = PiecePicker(9)
112
 
    p.complete(8)
113
 
    p.got_have(0)
114
 
    p.got_have(2)
115
 
    p.got_have(4)
116
 
    p.got_have(6)
117
 
    p.requested(1)
118
 
    p.requested(1)
119
 
    p.requested(3)
120
 
    p.requested(0)
121
 
    p.requested(6)
122
 
    v = _pull(p)
123
 
    assert v[:2] == [1, 3] or v[:2] == [3, 1]
124
 
    assert v[2:4] == [0, 6] or v[2:4] == [6, 0]
125
 
    assert v[4:] == [2, 4] or v[4:] == [4, 2]
126
 
 
127
 
def test_change_interest():
128
 
    p = PiecePicker(9)
129
 
    p.complete(8)
130
 
    p.got_have(0)
131
 
    p.got_have(2)
132
 
    p.got_have(4)
133
 
    p.got_have(6)
134
 
    p.lost_have(2)
135
 
    p.lost_have(6)
136
 
    v = _pull(p)
137
 
    assert v == [0, 4] or v == [4, 0]
138
 
 
139
 
def test_change_interest2():
140
 
    p = PiecePicker(9)
141
 
    p.complete(8)
142
 
    p.got_have(0)
143
 
    p.got_have(2)
144
 
    p.got_have(4)
145
 
    p.got_have(6)
146
 
    p.lost_have(2)
147
 
    p.lost_have(6)
148
 
    v = _pull(p)
149
 
    assert v == [0, 4] or v == [4, 0]
150
 
 
151
 
def test_complete():
152
 
    p = PiecePicker(1)
153
 
    p.got_have(0)
154
 
    p.complete(0)
155
 
    assert _pull(p) == []
156
 
    p.got_have(0)
157
 
    p.lost_have(0)
158
 
 
159
 
def test_rarer_in_started_takes_priority():
160
 
    p = PiecePicker(3)
161
 
    p.complete(2)
162
 
    p.requested(0)
163
 
    p.requested(1)
164
 
    p.got_have(1)
165
 
    p.got_have(0)
166
 
    p.got_have(0)
167
 
    assert _pull(p) == [1, 0]
168
 
 
169
 
def test_zero():
170
 
    assert _pull(PiecePicker(0)) == []
171
 
 
172
 
def _pull(pp):
173
 
    r = []
174
 
    def want(p, r = r):
175
 
        return p not in r
176
 
    while True:
177
 
        n = pp.next(want)
178
 
        if n is None:
179
 
            break
180
 
        r.append(n)
181
 
    return r