| 24 | |
| 25 | |
| 26 | class LRUCache(object): |
| 27 | _cnt: int = 0 |
| 28 | |
| 29 | def __init__(self, cache_size, keep_the_1st=True): |
| 30 | assert cache_size >= 2 |
| 31 | self.cache_size = cache_size |
| 32 | self.hash_map = dict() |
| 33 | self.keep_the_1st = keep_the_1st |
| 34 | self.queue = deque() |
| 35 | |
| 36 | def is_empty(self): |
| 37 | return len(self.hash_map) == 0 |
| 38 | |
| 39 | def is_full(self): |
| 40 | return len(self.hash_map) >= self.cache_size |
| 41 | |
| 42 | def pop(self): |
| 43 | if len(self.queue) == 0: |
| 44 | return None |
| 45 | pop_key = self.queue.pop() |
| 46 | value = self.hash_map.pop(pop_key) |
| 47 | del value |
| 48 | return pop_key |
| 49 | |
| 50 | def set(self, key, value): |
| 51 | new_key = None |
| 52 | old_key = None |
| 53 | if key in self.hash_map: |
| 54 | return new_key, old_key |
| 55 | |
| 56 | if self.is_full(): |
| 57 | old_key = self.pop() |
| 58 | assert old_key is not None, f"Cache size is {self.cache_size}, at least 2." |
| 59 | assert not self.is_full() |
| 60 | |
| 61 | if not (self.keep_the_1st and self.is_empty()): |
| 62 | self.queue.appendleft(key) |
| 63 | |
| 64 | value._oneflow_graph_cache_order = LRUCache._cnt |
| 65 | LRUCache._cnt += 1 |
| 66 | self.hash_map[key] = value |
| 67 | new_key = key |
| 68 | return new_key, old_key |
| 69 | |
| 70 | def get(self, key): |
| 71 | if key in self.hash_map: |
| 72 | if key in self.queue: |
| 73 | self.queue.remove(key) |
| 74 | self.queue.appendleft(key) |
| 75 | return self.hash_map[key] |
| 76 | |
| 77 | return None |
| 78 | |
| 79 | def items(self): |
| 80 | for (key, value) in self.hash_map.items(): |
| 81 | yield (key, value) |
| 82 | |
| 83 |
no outgoing calls
no test coverage detected