WIP: Type TreeCache erikj/tree_cache_typing

2 files changed, 125 insertions, 44 deletions

diff --git a/synapse/util/caches/lrucache.py b/synapse/util/caches/lrucache.py
index 6f95c1354e..80671daca3 100644
--- a/synapse/util/caches/lrucache.py
+++ b/synapse/util/caches/lrucache.py
@@ -417,7 +417,7 @@ class LruCache(Generic[KT, VT]):
         else:
             real_clock = clock
 
-        cache: Union[Dict[KT, _Node[KT, VT]], TreeCache] = cache_type()
+        cache: Union[Dict[KT, _Node[KT, VT]], TreeCache[_Node[KT, VT]]] = cache_type()
         self.cache = cache  # Used for introspection.
         self.apply_cache_factor_from_config = apply_cache_factor_from_config
 
diff --git a/synapse/util/caches/treecache.py b/synapse/util/caches/treecache.py
index c0136ea269..37dcc040f4 100644
--- a/synapse/util/caches/treecache.py
+++ b/synapse/util/caches/treecache.py
@@ -12,18 +12,59 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-SENTINEL = object()
+from enum import Enum
+from typing import (
+    Any,
+    Dict,
+    Generator,
+    Generic,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    TypeVar,
+    Union,
+    overload,
+)
 
 
-class TreeCacheNode(dict):
+class Sentinel(Enum):
+    sentinel = object()
+
+
+V = TypeVar("V")
+T = TypeVar("T")
+
+
+class TreeCacheNode(Generic[V]):
     """The type of nodes in our tree.
 
-    Has its own type so we can distinguish it from real dicts that are stored at the
-    leaves.
+    Either a leaf node or a branch node.
     """
 
+    __slots__ = ["leaf_value", "sub_tree"]
+
+    def __init__(
+        self,
+        leaf_value: Union[V, Literal[Sentinel.sentinel]] = Sentinel.sentinel,
+        sub_tree: Optional[Dict[Any, "TreeCacheNode[V]"]] = None,
+    ) -> None:
+        if leaf_value is Sentinel.sentinel and sub_tree is None:
+            raise Exception("One of leaf or sub tree must be set")
+
+        self.leaf_value: Union[V, Literal[Sentinel.sentinel]] = leaf_value
+        self.sub_tree: Optional[Dict[Any, "TreeCacheNode[V]"]] = sub_tree
+
+    @staticmethod
+    def leaf(value: V) -> "TreeCacheNode[V]":
+        return TreeCacheNode(leaf_value=value)
+
+    @staticmethod
+    def empty_branch() -> "TreeCacheNode[V]":
+        return TreeCacheNode(sub_tree={})
 
-class TreeCache:
+
+class TreeCache(Generic[V]):
     """
     Tree-based backing store for LruCache. Allows subtrees of data to be deleted
     efficiently.
@@ -35,15 +76,15 @@ class TreeCache:
 
     def __init__(self) -> None:
         self.size: int = 0
-        self.root = TreeCacheNode()
+        self.root: TreeCacheNode[V] = TreeCacheNode.empty_branch()
 
-    def __setitem__(self, key, value) -> None:
+    def __setitem__(self, key: tuple, value: V) -> None:
         self.set(key, value)
 
-    def __contains__(self, key) -> bool:
-        return self.get(key, SENTINEL) is not SENTINEL
+    def __contains__(self, key: tuple) -> bool:
+        return self.get(key, None) is not None
 
-    def set(self, key, value) -> None:
+    def set(self, key: tuple, value: V) -> None:
         if isinstance(value, TreeCacheNode):
             # this would mean we couldn't tell where our tree ended and the value
             # started.
@@ -51,31 +92,56 @@ class TreeCache:
 
         node = self.root
         for k in key[:-1]:
-            next_node = node.get(k, SENTINEL)
-            if next_node is SENTINEL:
-                next_node = node[k] = TreeCacheNode()
-            elif not isinstance(next_node, TreeCacheNode):
-                # this suggests that the caller is not being consistent with its key
-                # length.
+            sub_tree = node.sub_tree
+            if sub_tree is None:
                 raise ValueError("value conflicts with an existing subtree")
-            node = next_node
 
-        node[key[-1]] = value
+            next_node = sub_tree.get(k, None)
+            if next_node is None:
+                node = TreeCacheNode.empty_branch()
+                sub_tree[k] = node
+            else:
+                node = next_node
+
+        if node.sub_tree is None:
+            raise ValueError("value conflicts with an existing subtree")
+
+        node.sub_tree[key[-1]] = TreeCacheNode.leaf(value)
         self.size += 1
 
-    def get(self, key, default=None):
+    @overload
+    def get(self, key: tuple, default: Literal[None] = None) -> Union[None, V]:
+        ...
+
+    @overload
+    def get(self, key: tuple, default: T) -> Union[T, V]:
+        ...
+
+    def get(self, key: tuple, default: Optional[T] = None) -> Union[None, T, V]:
         node = self.root
-        for k in key[:-1]:
-            node = node.get(k, None)
-            if node is None:
+        for k in key:
+            sub_tree = node.sub_tree
+            if sub_tree is None:
+                raise ValueError("get() key too long")
+
+            next_node = sub_tree.get(k, None)
+            if next_node is None:
                 return default
-        return node.get(key[-1], default)
+
+            node = next_node
+
+        if node.leaf_value is Sentinel.sentinel:
+            raise ValueError("key points to a branch")
+
+        return node.leaf_value
 
     def clear(self) -> None:
         self.size = 0
         self.root = TreeCacheNode()
 
-    def pop(self, key, default=None):
+    def pop(
+        self, key: tuple, default: Optional[T] = None
+    ) -> Union[None, T, V, TreeCacheNode[V]]:
         """Remove the given key, or subkey, from the cache
 
         Args:
@@ -91,20 +157,25 @@ class TreeCache:
             raise TypeError("The cache key must be a tuple not %r" % (type(key),))
 
         # a list of the nodes we have touched on the way down the tree
-        nodes = []
+        nodes: List[TreeCacheNode[V]] = []
 
         node = self.root
         for k in key[:-1]:
-            node = node.get(k, None)
-            if node is None:
-                return default
-            if not isinstance(node, TreeCacheNode):
-                # we've gone off the end of the tree
+            sub_tree = node.sub_tree
+            if sub_tree is None:
                 raise ValueError("pop() key too long")
-            nodes.append(node)  # don't add the root node
-        popped = node.pop(key[-1], SENTINEL)
-        if popped is SENTINEL:
-            return default
+
+            next_node = sub_tree.get(k, None)
+            if next_node is None:
+                return default
+
+            node = next_node
+            nodes.append(node)
+
+        if node.sub_tree is None:
+            raise ValueError("pop() key too long")
+
+        popped = node.sub_tree.pop(key[-1])
 
         # working back up the tree, clear out any nodes that are now empty
         node_and_keys = list(zip(nodes, key))
@@ -116,8 +187,13 @@ class TreeCache:
 
             if n:
                 break
+
             # found an empty node: remove it from its parent, and loop.
-            node_and_keys[i + 1][0].pop(k)
+            node = node_and_keys[i + 1][0]
+
+            # We added it to the list so already know its a branch node.
+            assert node.sub_tree is not None
+            node.sub_tree.pop(k)
 
         cnt = sum(1 for _ in iterate_tree_cache_entry(popped))
         self.size -= cnt
@@ -130,26 +206,31 @@ class TreeCache:
         return self.size
 
 
-def iterate_tree_cache_entry(d):
+def iterate_tree_cache_entry(d: TreeCacheNode[V]) -> Generator[V, None, None]:
     """Helper function to iterate over the leaves of a tree, i.e. a dict of that
     can contain dicts.
     """
-    if isinstance(d, TreeCacheNode):
-        for value_d in d.values():
+
+    if d.sub_tree is not None:
+        for value_d in d.sub_tree.values():
             yield from iterate_tree_cache_entry(value_d)
     else:
-        yield d
+        assert d.leaf_value is not Sentinel.sentinel
+        yield d.leaf_value
 
 
-def iterate_tree_cache_items(key, value):
+def iterate_tree_cache_items(
+    key: tuple, value: TreeCacheNode[V]
+) -> Generator[Tuple[tuple, V], None, None]:
     """Helper function to iterate over the leaves of a tree, i.e. a dict of that
     can contain dicts.
 
     Returns:
         A generator yielding key/value pairs.
     """
-    if isinstance(value, TreeCacheNode):
-        for sub_key, sub_value in value.items():
+    if value.sub_tree is not None:
+        for sub_key, sub_value in value.sub_tree.items():
             yield from iterate_tree_cache_items((*key, sub_key), sub_value)
     else:
-        yield key, value
+        assert value.leaf_value is not Sentinel.sentinel
+        yield key, value.leaf_value