| 10 | |
| 11 | @dataclass |
| 12 | class PrefixTreeBlock: |
| 13 | idx: int |
| 14 | contents: tuple[int, ...] = field(default_factory=tuple) |
| 15 | last_used_at: float = 0.0 |
| 16 | seq_ids: set[str] = field(default_factory=set) |
| 17 | children: dict[tuple[int, ...], "PrefixTreeBlock"] = field(default_factory=dict) |
| 18 | parent: Optional["PrefixTreeBlock"] = None |
| 19 | |
| 20 | def is_root(self): |
| 21 | return self.idx == -1 |
| 22 | |
| 23 | def detach_from_parent(self): |
| 24 | assert self.parent is not None |
| 25 | assert self.parent.children[self.contents] == self |
| 26 | self.parent.children.pop(self.contents) |
| 27 | self.parent = None |
| 28 | |
| 29 | def wipe(self): |
| 30 | if self.parent is not None: |
| 31 | self.detach_from_parent() |
| 32 | self.children.clear() |
| 33 | self.seq_ids.clear() |
| 34 | self.contents = tuple() |
| 35 | self.last_used_at = 0.0 |
| 36 | |
| 37 | def __repr__(self): |
| 38 | return f"Block(idx={self.idx}, contents={self.contents}, children={len(self.children)} used_by={len(self.seq_ids)})" |
| 39 | |
| 40 | def __hash__(self): |
| 41 | return self.idx |
| 42 | |
| 43 | def __lt__(self, other: "PrefixTreeBlock"): |
| 44 | return self.last_used_at < other.last_used_at |
| 45 | |
| 46 | def tree_repr(self): |
| 47 | def indent(s: str, spacing: int): |
| 48 | lines = s.split("\n") |
| 49 | with_index = [" " * spacing + line for line in lines] |
| 50 | return "\n".join(with_index) |
| 51 | |
| 52 | out_lines = [repr(self)] |
| 53 | |
| 54 | for child in self.children.values(): |
| 55 | out_lines.append(indent(child.tree_repr(), 2)) |
| 56 | |
| 57 | return "\n".join(out_lines) |
| 58 | |
| 59 | |
| 60 | class NoSpaceException(ValueError): |