116 lines
3.1 KiB
Python
116 lines
3.1 KiB
Python
"""Stateful iterator for flat trees."""
|
|
|
|
__all__ = ['FlatTreeIterator']
|
|
|
|
import attr
|
|
|
|
from flat_tree.accessor import FlatTreeAccessor
|
|
|
|
|
|
@attr.s(auto_attribs=True)
|
|
class FlatTreeIterator:
|
|
"""Stateful iterator for flat trees."""
|
|
|
|
index: int = 0
|
|
offset: int = 0
|
|
factor: int = 0
|
|
accessor: FlatTreeAccessor = FlatTreeAccessor()
|
|
|
|
def __attrs_post_init__(self):
|
|
self.seek(self.index)
|
|
|
|
# TODO(decentral1se): Once we get to the point of actually using this
|
|
# module in hypercore or wherever it is needed we should consider whether
|
|
# or not to make this a real Python iterator using the protocol methods of
|
|
# __iter__ and __next__
|
|
def next(self) -> int:
|
|
"""The next index in the tree."""
|
|
self.offset += 1
|
|
self.index += self.factor
|
|
return self.index
|
|
|
|
def prev(self) -> int:
|
|
"""The previous index in the tree."""
|
|
if not self.offset:
|
|
return self.index
|
|
|
|
self.offset -= 1
|
|
self.index = self.factor
|
|
|
|
return self.index
|
|
|
|
def seek(self, index: int) -> None:
|
|
"""Move iterator to the given index.
|
|
|
|
:param index: The index to move to
|
|
"""
|
|
self.index = index
|
|
|
|
if self.index & 1:
|
|
self.offset = self.accessor.offset(index)
|
|
self.factor = 2 ** (self.accessor.depth(index) + 1)
|
|
else:
|
|
self.offset = int(index / 2)
|
|
self.factor = 2
|
|
|
|
def parent(self) -> int:
|
|
"""Move iterator to the parent index."""
|
|
if self.offset & 1:
|
|
self.index -= int(self.factor / 2)
|
|
self.offset = int((self.offset - 1) / 2)
|
|
else:
|
|
self.index += int(self.factor / 2)
|
|
self.offset = int(self.offset / 2)
|
|
|
|
self.factor *= 2
|
|
|
|
return self.index
|
|
|
|
def left_child(self) -> int:
|
|
"""Move iterator to the left child."""
|
|
if self.factor == 2:
|
|
return self.index
|
|
|
|
self.factor = int(self.factor / 2)
|
|
self.index -= int(self.factor / 2)
|
|
self.offset *= 2
|
|
|
|
return self.index
|
|
|
|
def right_child(self) -> int:
|
|
"""Move iterator to the right child."""
|
|
if self.factor == 2:
|
|
return self.index
|
|
|
|
self.factor = int(self.factor / 2)
|
|
self.index += int(self.factor / 2)
|
|
self.offset = (2 * self.offset) + 1
|
|
|
|
return self.index
|
|
|
|
def left_span(self) -> int:
|
|
"""Move iterator to the left span."""
|
|
self.index = int(self.index - (self.factor / 2)) + 1
|
|
self.offset = int(self.index / 2)
|
|
self.factor = 2
|
|
return self.index
|
|
|
|
def right_span(self) -> int:
|
|
"""Move iterator to the right span."""
|
|
self.index = int((self.index - self.factor) / 2) - 1
|
|
self.offset = int(self.index / 2)
|
|
self.factor = 2
|
|
return self.index
|
|
|
|
def sibling(self) -> int:
|
|
"""Move iterator to the sibling."""
|
|
return self.next() if self.is_left() else self.prev()
|
|
|
|
def is_left(self) -> bool:
|
|
"""Is this index a left sibling?"""
|
|
return not self.offset & 1
|
|
|
|
def is_right(self) -> bool:
|
|
"""Is this index a right sibling?"""
|
|
return not self.is_left()
|