mst.py

from dataclasses import dataclass
from abc import ABC, abstractstaticmethod
from typing import Tuple, Self, Optional, Any, Type, Iterable

# key type - could be anything comparable, in theory
KTYPE = str
# could be literally anything (will be CID object in atproto)
VTYPE = Any

# tuple helpers
def tuple_replace_at(original: tuple, i: int, value: Any) -> tuple:
	return original[:i] + (value,) + original[i + 1:]

def tuple_insert_at(original: tuple, i: int, value: Any) -> tuple:
	return original[:i] + (value,) + original[i:]

def tuple_remove_at(original: tuple, i: int) -> tuple:
	return original[:i] + original[i + 1:]


@dataclass(frozen=True) # frozen == immutable == win
class MSTNode(ABC):
	subtrees: Tuple[Optional[Self]]
	keys: Tuple[KTYPE]
	vals: Tuple[VTYPE]

	@abstractstaticmethod
	def key_height(key: str) -> int:
		pass

	# NB: __init__ is auto-generated by dataclass decorator

	# these checks should never fail, and could be skipped for performance
	def __post_init__(self) -> None:
		# TODO: maybe check that they're tuples here?
		# implicitly, the length of self.subtrees must be at least 1
		if len(self.subtrees) != len(self.keys) + 1:
			raise ValueError("Invalid subtree count")
		if len(self.keys) != len(self.vals):
			raise ValueError("Mismatched keys/vals lengths")

	@classmethod
	def empty_root(cls) -> Self:
		return cls(
			subtrees=(None,),
			keys=(),
			vals=()
		)

	@classmethod
	def _from_optional(cls, value: Optional[Self]) -> Self:
		"""
		turn None values into empty nodes
		"""
		if value is None:
			return cls.empty_root()
		return value
	
	def _to_optional(self) -> Optional[Self]:
		"""
		returns None if the node is empty
		"""
		if self.subtrees == (None,):
			return None
		return self
	
	def _squash_top(self, created: set) -> Self:
		"""
		strip empty nodes from the top of the tree
		"""
		if self.keys:
			return self
		if self.subtrees[0] is None:
			return self
		created.discard(self)
		return self.subtrees[0]._squash_top(created)


	# we're immutable, so this could be cached
	def height(self) -> int:
		# if there are keys at this level, query one directly
		if self.keys:
			return self.key_height(self.keys[0])
		
		# we're an empty tree
		if self.subtrees[0] is None:
			return 0
		
		# we're a non-empty tree with no keys at this level - look below, recursively
		return self.subtrees[0].height() + 1
	
	def _gte_index(self, key: KTYPE) -> int:
		"""
		find the index of the first key greater than or equal to the specified key
		if all keys are smaller, it returns len(keys)
		"""
		i = 0 # this loop could be a binary search but not worth it for small fanouts
		while i < len(self.keys) and key > self.keys[i]:
			i += 1
		return i
	
	def get(self, key: KTYPE, sentinel: Any=None) -> VTYPE | Any:
		key_height = self.key_height(key)
		tree_height = self.height()
		if key_height > tree_height:
			return sentinel
		if key_height < tree_height:
			subtree = self.subtrees[self._gte_index(key)]
			if subtree is None:
				return sentinel
			return subtree.get(key, sentinel)
		i = self._gte_index(key)
		if i == len(self.keys):
			return sentinel
		if self.keys[i] != key:
			return sentinel
		return self.vals[i]
	
	def get_range(self, key_min: KTYPE, key_max: KTYPE, reverse: bool=False) -> Iterable[Tuple[KTYPE, VTYPE]]:
		# currently inclusive, exclusive I thiiiink
		start, end = self._gte_index(key_min), self._gte_index(key_max)
		if reverse:
			if self.subtrees[end] is not None:
				yield from self.subtrees[end].get_range(key_min, key_max, reverse)
			for i in reversed(range(start, end)):
				yield self.keys[i], self.vals[i]
				if self.subtrees[i] is not None:
					yield from self.subtrees[i].get_range(key_min, key_max, reverse)
		else:
			for i in range(start, end):
				if self.subtrees[i] is not None:
					yield from self.subtrees[i].get_range(key_min, key_max, reverse)
				yield self.keys[i], self.vals[i]
			if self.subtrees[end] is not None:
				yield from self.subtrees[end].get_range(key_min, key_max, reverse)

	def put(self, key: KTYPE, val: VTYPE, created: set) -> Self:
		if self.subtrees == (None,): # special case for empty tree
			return self._put_here(key, val, created)

		return self._put_recursive(key, val, self.key_height(key), self.height(), created)

	
	def _put_recursive(self, key: KTYPE, val: VTYPE, key_height: int, tree_height: int, created: set) -> Self:
		cls = self.__class__ # maybe this could be a class method???

		if key_height > tree_height: # we need to grow the tree
			new = cls(
				subtrees=(self,),
				keys=(),
				vals=()
			)._put_recursive(key, val, key_height, tree_height + 1, created)
			created.add(new)
			return new
		
		if key_height < tree_height: # we need to look below
			i = self._gte_index(key)
			new = cls(
				subtrees=tuple_replace_at(
					self.subtrees, i,
					cls._from_optional(self.subtrees[i])._put_recursive(
						key, val, key_height, tree_height - 1, created
					)
				),
				keys=self.keys,
				vals=self.vals
			)
			created.add(new)
			return new
		
		# we can insert here
		assert(key_height == tree_height)
		return self._put_here(key, val, created)

	# insert a key/value record into this node. Caller is responsible for making
	# sure this is the right height to insert it at
	def _put_here(self, key: KTYPE, val: VTYPE, created: set) -> Self:
		cls = self.__class__ # maybe this could be a class method???
		
		i = self._gte_index(key)
		# the key is already present!
		if i < len(self.keys) and self.keys[i] == key:
			if self.vals[i] == val:
				return self # we can return our old self if there is no change
			new = cls(
				subtrees=self.subtrees,
				keys=self.keys,
				vals=tuple_replace_at(self.vals, i, val)
			)
			created.add(new)
			return new
		
		new = cls(
			subtrees = self.subtrees[:i] + \
				cls._split_on_key(self.subtrees[i], key, created) + \
				self.subtrees[i + 1:],
			keys=tuple_insert_at(self.keys, i, key),
			vals=tuple_insert_at(self.vals, i, val),
		)
		created.add(new)
		return new
	
	@classmethod
	def _split_on_key(cls, tree: Optional[Self], key: KTYPE, created: set) -> Tuple[Optional[Self], Optional[Self]]:
		if tree is None:
			return None, None
		i = tree._gte_index(key)
		lsub, rsub = cls._split_on_key(tree.subtrees[i], key, created)
		left = cls(
			subtrees=tree.subtrees[:i] + (lsub,),
			keys=tree.keys[:i],
			vals=tree.vals[:i]
		)._to_optional()
		right = cls(
			subtrees=(rsub,) + tree.subtrees[i + 1:],
			keys=tree.keys[i:],
			vals=tree.vals[i:]
		)._to_optional()
		if left is not None:
			created.add(left)
		if right is not None:
			created.add(right)
		return left, right

	def delete(self, key: KTYPE, created: set) -> Self:
		return self.__class__._from_optional(self._delete_recursive(key, self.key_height(key), self.height(), created))._squash_top(created)

	def _delete_recursive(self, key: KTYPE, key_height: int, tree_height: int, created: set) -> Optional[Self]:
		cls = self.__class__

		if key_height > tree_height: # the key cannot possibly be in this tree, no change needed
			return self
		elif key_height < tree_height: # the key must be deleted from a subtree
			i = self._gte_index(key)
			if self.subtrees[i] is None:
				return self # the key cannot be in this subtree, no change needed
			new = cls(
				subtrees=tuple_replace_at(
					self.subtrees,
					i,
					self.subtrees[i]._delete_recursive(key, key_height, tree_height - 1, created)
				),
				keys=self.keys,
				vals=self.vals
			)._to_optional()
			if new is not None:
				created.add(new)
			return new
		
		i = self._gte_index(key)
		if i == len(self.keys) or self.keys[i] != key:
			return self # key already not present
		
		assert(self.keys[i] == key) # sanity check (should always be true)

		new = cls(
			subtrees=self.subtrees[:i] + (
				cls._merge(self.subtrees[i], self.subtrees[i + 1], created),
			) + self.subtrees[i + 2:],
			keys=tuple_remove_at(self.keys, i),
			vals=tuple_remove_at(self.vals, i)
		)._to_optional()
		if new is not None:
			created.add(new)
		return new
	
	@classmethod
	def _merge(cls, left: Optional[Self], right: Optional[Self], created: set) -> Optional[Self]:
		if left is None:
			return right # includes the case where left == right == None
		if right is None:
			return left
		new = left.__class__(
			subtrees=left.subtrees[:-1] + (
				cls._merge(
					left.subtrees[-1],
					right.subtrees[0],
					created
				),
			 ) + right.subtrees[1:],
			keys=left.keys + right.keys,
			vals=left.vals + right.vals
		)._to_optional()
		if new is not None:
			created.add(new)
		return new


# Nodes are immutable, the Tree class wraps them and provides a mutable interface
@dataclass
class MST:
	root: MSTNode

	# maybe this should just be __init__ idk
	@classmethod
	def new_with(cls: Type[Self], node_type: Type[MSTNode]) -> Self:
		return cls(root=node_type.empty_root())

	def height(self) -> int:
		return self.root.height()
	
	def __setitem__(self, key: KTYPE, val: VTYPE) -> None:
		self.root = self.root.put(key, val, set())
	
	def __delitem__(self, key: KTYPE) -> None:
		prev_root = self.root
		self.root = self.root.delete(key, set())
		if self.root == prev_root: # if nothing changed it's because the key didn't exist in the first place
			raise KeyError(key)
	
	def get(self, key: KTYPE, sentinel: Any=None) -> VTYPE | Any:
		return self.root.get(key, sentinel)

	def get_range(self, key_min: KTYPE, key_max: KTYPE, reverse: bool=False) -> Iterable[Tuple[KTYPE, VTYPE]]:
		return self.root.get_range(key_min, key_max, reverse)

	def __getitem__(self, key: KTYPE) -> VTYPE:
		value = self.get(key)
		if value is None: # TODO: use a proper sentinel object
			raise KeyError(key)
		return value


if __name__ == "__main__":

	# using len(key) is convenient for testing
	class StrlenNode(MSTNode):
		@staticmethod
		def key_height(key: str) -> int:
			return len(key)
	
	# make a new tree
	tree = MST.new_with(StrlenNode)

	# store something
	tree["foo"] = "bar"

	# retrieve it back again
	assert(tree["foo"] == "bar")

	# the structure should look like this
	assert(tree.root == StrlenNode(
		subtrees=(None, None),
		keys=('foo',),
		vals=('bar',)
	))

	tree["f"] = "f" # insert a record "below"
	tree["foooooo"] = "foooooo" # insert a record above

	# it's getting a bit messy now...
	assert(tree.root == StrlenNode(
		subtrees=(StrlenNode(
			subtrees=(StrlenNode(
				subtrees=(StrlenNode(
					subtrees=(StrlenNode(
						subtrees=(StrlenNode(
							subtrees=(StrlenNode(
								subtrees=(None, None),
								keys=("f",),
								vals=("f",),
							),),
							keys=(),
							vals=(),
						), None),
						keys=("foo",),
						vals=("bar",),
					),),
					keys=(),
					vals=(),
				),),
				keys=(),
				vals=(),
			),),
			keys=(),
			vals=(),
		), None),
		keys=("foooooo",),
		vals=("foooooo",),
	))

	# add some more...
	tree["bar"] = "bar"
	tree["bat"] = "bat"

	# delete some
	del tree["foo"]
	del tree["bar"]
	del tree["bat"]

	# make a copy of the tree
	t2 = MST(tree.root)

	del t2["foooooo"]

	assert(t2.root == StrlenNode(subtrees=(None, None), keys=('f',), vals=('f',)))

	# make another copy
	t3 = MST(tree.root)

	del t3["f"]

	assert(t3.root == StrlenNode(subtrees=(None, None), keys=('foooooo',), vals=('foooooo',)))

	assert(t3 != t2)

	del t2["f"]
	del t3["foooooo"]

	# they should both be empty now, and identical
	assert(t3 == t2)

	assert(t3.root == StrlenNode(subtrees=(None,), keys=(), vals=()))


	# test inserting the same things in different orders
	words1 = ["foo", "bar", "hello", "world", "this", "is", "a", "test"]
	words2 = sorted(words1)

	assert(words1 != words2) # different orders!

	tree1 = MST.new_with(StrlenNode)
	for word in words1:
		tree1[word] = 123 # value itself is irrelevant

	tree2 = MST.new_with(StrlenNode)
	for word in words2:
		tree2[word] = 123 # value itself is irrelevant
	
	assert(tree1 == tree2)
	backup = MST(tree1.root)

	words3 = ["apple", "banana", "cherry", "grapefruit"]
	words4 = ["red", "green", "blue", "indigo"]

	for word in words3:
		tree1[word] = 123 # value itself is irrelevant

	for word in words4:
		tree2[word] = 123 # value itself is irrelevant
	
	# the trees should now be different (we inserted different things!)
	assert(tree1 != tree2)

	for word in reversed(words3): # order should not matter, lets reverse for fun
		del tree1[word]

	for word in words4:
		del tree2[word]
	
	# we should be back to identical trees again
	assert(tree1 == tree2)
	assert(tree1 == backup) # and identical to the copy we made earlier


	# testing range iteration
	tree = MST.new_with(StrlenNode)
	tree["0"] = None
	tree["01"] = None
	tree["02"] = None
	tree["1"] = None
	tree["12"] = None
	tree["13"] = None
	tree["2"] = None

	assert(list(tree.get_range("02", "3")) == [
		('02', None),
		('1', None),
		('12', None),
		('13', None),
		('2', None)
	])

	# check that reversed result is the same but... reversed
	assert(list(reversed(list(tree.get_range("02", "3")))) == list(tree.get_range("02", "3", reverse=True)))

	# end value is exclusive, not inclusive
	assert(list(tree.get_range("0", "13")) == [
		('0', None),
		('01', None),
		('02', None),
		('1', None),
		('12', None)
	])

	# protip, add a null char to get a pseudo-includive end value
	assert(list(tree.get_range("0", "13\0")) == [
		('0', None),
		('01', None),
		('02', None),
		('1', None),
		('12', None),
		('13', None)
	])