mirror of
https://mau.dev/maunium/synapse.git
synced 2024-08-25 03:35:28 +00:00
![Erik Johnston](/assets/img/avatar_default.png)
During the migration the automated script to update the copyright headers accidentally got rid of some of the existing copyright lines. Reinstate them.
190 lines
5.8 KiB
Python
190 lines
5.8 KiB
Python
#
|
|
# This file is licensed under the Affero General Public License (AGPL) version 3.
|
|
#
|
|
# Copyright 2016-2021 The Matrix.org Foundation C.I.C.
|
|
# Copyright (C) 2023 New Vector, Ltd
|
|
#
|
|
# This program is free software: you can redistribute it and/or modify
|
|
# it under the terms of the GNU Affero General Public License as
|
|
# published by the Free Software Foundation, either version 3 of the
|
|
# License, or (at your option) any later version.
|
|
#
|
|
# See the GNU Affero General Public License for more details:
|
|
# <https://www.gnu.org/licenses/agpl-3.0.html>.
|
|
#
|
|
# Originally licensed under the Apache License, Version 2.0:
|
|
# <http://www.apache.org/licenses/LICENSE-2.0>.
|
|
#
|
|
# [This file includes modifications made by New Vector Limited]
|
|
#
|
|
#
|
|
|
|
SENTINEL = object()
|
|
|
|
|
|
class TreeCacheNode(dict):
|
|
"""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.
|
|
"""
|
|
|
|
|
|
class TreeCache:
|
|
"""
|
|
Tree-based backing store for LruCache. Allows subtrees of data to be deleted
|
|
efficiently.
|
|
Keys must be tuples.
|
|
|
|
The data structure is a chain of TreeCacheNodes:
|
|
root = {key_1: {key_2: _value}}
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self.size: int = 0
|
|
self.root = TreeCacheNode()
|
|
|
|
def __setitem__(self, key, value) -> None:
|
|
self.set(key, value)
|
|
|
|
def __contains__(self, key) -> bool:
|
|
return self.get(key, SENTINEL) is not SENTINEL
|
|
|
|
def set(self, key, value) -> None:
|
|
if isinstance(value, TreeCacheNode):
|
|
# this would mean we couldn't tell where our tree ended and the value
|
|
# started.
|
|
raise ValueError("Cannot store TreeCacheNodes in a 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.
|
|
raise ValueError("value conflicts with an existing subtree")
|
|
node = next_node
|
|
|
|
node[key[-1]] = value
|
|
self.size += 1
|
|
|
|
def get(self, key, default=None):
|
|
"""When `key` is a full key, fetches the value for the given key (if
|
|
any).
|
|
|
|
If `key` is only a partial key (i.e. a truncated tuple) then returns a
|
|
`TreeCacheNode`, which can be passed to the `iterate_tree_cache_*`
|
|
functions to iterate over all entries in the cache with keys that start
|
|
with the given partial key.
|
|
"""
|
|
|
|
node = self.root
|
|
for k in key[:-1]:
|
|
node = node.get(k, None)
|
|
if node is None:
|
|
return default
|
|
return node.get(key[-1], default)
|
|
|
|
def clear(self) -> None:
|
|
self.size = 0
|
|
self.root = TreeCacheNode()
|
|
|
|
def pop(self, key, default=None):
|
|
"""Remove the given key, or subkey, from the cache
|
|
|
|
Args:
|
|
key: key or subkey to remove.
|
|
default: value to return if key is not found
|
|
|
|
Returns:
|
|
If the key is not found, 'default'. If the key is complete, the removed
|
|
value. If the key is partial, the TreeCacheNode corresponding to the part
|
|
of the tree that was removed.
|
|
"""
|
|
if not isinstance(key, tuple):
|
|
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 = []
|
|
|
|
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
|
|
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
|
|
|
|
# working back up the tree, clear out any nodes that are now empty
|
|
node_and_keys = list(zip(nodes, key))
|
|
node_and_keys.reverse()
|
|
node_and_keys.append((self.root, None))
|
|
|
|
for i in range(len(node_and_keys) - 1):
|
|
n, k = node_and_keys[i]
|
|
|
|
if n:
|
|
break
|
|
# found an empty node: remove it from its parent, and loop.
|
|
node_and_keys[i + 1][0].pop(k)
|
|
|
|
cnt = sum(1 for _ in iterate_tree_cache_entry(popped))
|
|
self.size -= cnt
|
|
return popped
|
|
|
|
def values(self):
|
|
return iterate_tree_cache_entry(self.root)
|
|
|
|
def items(self):
|
|
return iterate_tree_cache_items((), self.root)
|
|
|
|
def __len__(self) -> int:
|
|
return self.size
|
|
|
|
|
|
def iterate_tree_cache_entry(d):
|
|
"""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():
|
|
yield from iterate_tree_cache_entry(value_d)
|
|
else:
|
|
yield d
|
|
|
|
|
|
def iterate_tree_cache_items(key, value):
|
|
"""Helper function to iterate over the leaves of a tree, i.e. a dict of that
|
|
can contain dicts.
|
|
|
|
The provided key is a tuple that will get prepended to the returned keys.
|
|
|
|
Example:
|
|
|
|
cache = TreeCache()
|
|
cache[(1, 1)] = "a"
|
|
cache[(1, 2)] = "b"
|
|
cache[(2, 1)] = "c"
|
|
|
|
tree_node = cache.get((1,))
|
|
|
|
items = iterate_tree_cache_items((1,), tree_node)
|
|
assert list(items) == [((1, 1), "a"), ((1, 2), "b")]
|
|
|
|
Returns:
|
|
A generator yielding key/value pairs.
|
|
"""
|
|
if isinstance(value, TreeCacheNode):
|
|
for sub_key, sub_value in value.items():
|
|
yield from iterate_tree_cache_items((*key, sub_key), sub_value)
|
|
else:
|
|
# we've reached a leaf of the tree.
|
|
yield key, value
|