-
Notifications
You must be signed in to change notification settings - Fork 11
/
lru_python_impl.py
221 lines (203 loc) · 7.84 KB
/
lru_python_impl.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
from _weakref import proxy as _proxy
class Op(object):
def __init__(self, op_type, prev_node=None, next_node=None, removed_node=None, key=None, value=None):
self.op_type = op_type
self.prev_node = prev_node
self.next_node = next_node
self.removed_node = removed_node
self.kv = (key, value)
class LinkNode(object):
__slots__ = 'prev', 'next', 'key', 'value', '__weakref__'
class LRU:
def __init__(self, capacity):
self.capacity = capacity
self.cur_idx = 0
self.cache = {} # map key to the double link list.
self.__hard_head = LinkNode()
self.__hard_tail = LinkNode()
self.__head = _proxy(self.__hard_head)
self.__head.key = '___head___'
self.__tail = _proxy(self.__hard_tail)
self.__tail.key = '___tail___'
self.__head.next = self.__tail
self.__tail.prev = self.__head
self.__head.prev = self.__tail.next = None
self.op_stack = []
def get(self, key):
if key in self.cache:
cur_node = self.cache[key]
value = cur_node.value
prev_node, next_node = cur_node.prev, cur_node.next
next_node.prev = prev_node
prev_node.next = next_node
# insert head
cur_node.next = self.__head.next
self.__head.next.prev = cur_node
cur_node.prev = self.__head
self.__head.next = cur_node
return value
else:
new_node = LinkNode()
new_node.key = key
if self.cur_idx < self.capacity: # not full
r = self.cur_idx
self.cache[key] = new_node
new_node.value = r
self.cur_idx += 1
new_node.next = self.__head.next
self.__head.next.prev = new_node
new_node.prev = self.__head
self.__head.next = new_node
return r
else:
free_node = self.__tail.prev
free_node.prev.next = self.__tail
self.__tail.prev = free_node.prev
old_key = free_node.key
r = free_node.value
free_node.prev = free_node.next = None
self.cache.pop(old_key)
new_node.next = self.__head.next
self.__head.next.prev = new_node
self.__head.next = new_node
new_node.prev = self.__head
new_node.value = r
self.cache[key] = new_node
return r
def __iter__(self):
cur = self.__head.next
while cur.key != self.__tail.key:
# print(cur.key)
yield cur.key, cur.value
cur = cur.next
def state_dict(self):
cur = self.__head.next
kvs = []
while cur.key != self.__tail.key:
# print(cur.key)
kvs.append((cur.key, cur.value))
cur = cur.next
return kvs
def restore(self, kvs):
assert len(kvs) <= self.capacity
assert self.cur_idx == 0
cur_node = self.__head
for kv in kvs:
assert kv[0] not in self.cache
new_node = LinkNode()
new_node.key = kv[0]
new_node.value = kv[1]
self.cache[kv[0]] = new_node
cur_node.next = new_node
new_node.prev = cur_node
cur_node = new_node
self.cur_idx += 1
cur_node.next = self.__tail
self.__tail.prev = cur_node
def clear(self):
self.cache.clear()
self.__head.next = self.__tail
self.__tail.prev = self.__head
def __contains__(self, key):
return key in self.cache
def view(self, key):
if key not in self.cache:
return -1
else:
return self.cache[key].value
def keys(self):
return self.cache.keys()
def try_get(self, key):
if key in self.cache:
cur_node = self.cache[key]
value = cur_node.value
prev_node, next_node = cur_node.prev, cur_node.next
op_records = Op('Get', prev_node, next_node)
self.op_stack.append(op_records)
next_node.prev = prev_node
prev_node.next = next_node
# insert head
cur_node.next = self.__head.next
self.__head.next.prev = cur_node
cur_node.prev = self.__head
self.__head.next = cur_node
return value
else:
new_node = LinkNode()
new_node.key = key
if self.cur_idx < self.capacity: # not full
r = self.cur_idx
self.cache[key] = new_node
new_node.value = r
self.cur_idx += 1
# insert head
new_node.next = self.__head.next
self.__head.next.prev = new_node
new_node.prev = self.__head
self.__head.next = new_node
op_records = Op('Add', new_node.prev, new_node.next, new_node)
self.op_stack.append(op_records)
return r
else:
free_node = self.__tail.prev
free_node.prev.next = self.__tail
self.__tail.prev = free_node.prev
old_key = free_node.key
r = free_node.value
op_records = Op(
'Overflow', free_node.prev, free_node.next, free_node, key=old_key, value=free_node.value
)
self.op_stack.append(op_records)
free_node.prev = free_node.next = None
self.cache.pop(old_key)
new_node.next = self.__head.next
self.__head.next.prev = new_node
self.__head.next = new_node
new_node.prev = self.__head
new_node.value = r
self.cache[key] = new_node
return r
def rollback_one_step(self, ):
if len(self.op_stack) > 0:
last_op = self.op_stack.pop()
if last_op.op_type == 'Add':
removed_node = last_op.removed_node
prev_node, next_node = last_op.prev_node, last_op.next_node
prev_node.next = next_node
next_node.prev = prev_node
removed_node.prev = removed_node.next = None
self.cache.pop(removed_node.key)
self.cur_idx -= 1
elif last_op.op_type == 'Overflow':
# first remove new node
new_node = self.__head.next
self.__head.next = new_node.next
new_node.next.prev = self.__head
new_node.next = new_node.prev = None
self.cache.pop(new_node.key)
# add old node
removed_node = last_op.removed_node
prev_node = last_op.prev_node
next_node = last_op.next_node
prev_node.next = removed_node
removed_node.prev = prev_node
removed_node.next = next_node
next_node.prev = removed_node
key = last_op.kv[0]
removed_node.key = key
removed_node.value = last_op.kv[1]
assert key not in self.cache
self.cache[key] = removed_node
else:
assert last_op.op_type == 'Get'
cur_node = self.__head.next
self.__head.next = cur_node.next
cur_node.next.prev = self.__head
last_op.prev_node.next = cur_node
cur_node.prev = last_op.prev_node
cur_node.next = last_op.next_node
last_op.next_node.prev = cur_node
def rollback_steps(self, steps):
max_steps = min(steps, len(self.op_stack))
for _ in range(max_steps):
self.rollback_one_step()