diff --git a/rust/src/tree_cache/binding.rs b/rust/src/tree_cache/binding.rs
new file mode 100644
index 0000000000..70207f8781
--- /dev/null
+++ b/rust/src/tree_cache/binding.rs
@@ -0,0 +1,128 @@
+use std::hash::Hash;
+
+use anyhow::Error;
+use pyo3::{
+ pyclass, pymethods, types::PyModule, IntoPy, PyAny, PyObject, PyResult, Python, ToPyObject,
+};
+
+use super::TreeCache;
+
+pub fn register_module(py: Python<'_>, m: &PyModule) -> PyResult<()> {
+ let child_module = PyModule::new(py, "tree_cache")?;
+ child_module.add_class::<PythonTreeCache>()?;
+
+ m.add_submodule(child_module)?;
+
+ // We need to manually add the module to sys.modules to make `from
+ // synapse.synapse_rust import push` work.
+ py.import("sys")?
+ .getattr("modules")?
+ .set_item("synapse.synapse_rust.tree_cache", child_module)?;
+
+ Ok(())
+}
+
+struct HashablePyObject {
+ obj: PyObject,
+ hash: isize,
+}
+
+impl HashablePyObject {
+ pub fn new(obj: &PyAny) -> Result<Self, Error> {
+ let hash = obj.hash()?;
+
+ Ok(HashablePyObject {
+ obj: obj.to_object(obj.py()),
+ hash,
+ })
+ }
+}
+
+impl IntoPy<PyObject> for &HashablePyObject {
+ fn into_py(self, _: Python<'_>) -> PyObject {
+ self.obj.clone()
+ }
+}
+
+impl Hash for HashablePyObject {
+ fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+ self.hash.hash(state);
+ }
+}
+
+impl PartialEq for HashablePyObject {
+ fn eq(&self, other: &Self) -> bool {
+ let equal = Python::with_gil(|py| {
+ let result = self.obj.as_ref(py).eq(other.obj.as_ref(py));
+ result.unwrap_or(false)
+ });
+
+ equal
+ }
+}
+
+impl Eq for HashablePyObject {}
+
+#[pyclass]
+struct PythonTreeCache(TreeCache<HashablePyObject, PyObject>);
+
+#[pymethods]
+impl PythonTreeCache {
+ #[new]
+ fn new() -> Self {
+ PythonTreeCache(Default::default())
+ }
+
+ pub fn set(&mut self, key: &PyAny, value: PyObject) -> Result<(), Error> {
+ let v: Vec<HashablePyObject> = key
+ .iter()?
+ .map(|obj| HashablePyObject::new(obj?))
+ .collect::<Result<_, _>>()?;
+
+ self.0.set(v, value)?;
+
+ Ok(())
+ }
+
+ // pub fn get_node(&self, key: &PyAny) -> Result<Option<&TreeCacheNode<K, PyObject>>, Error> {
+ // todo!()
+ // }
+
+ pub fn get(&self, key: &PyAny) -> Result<Option<&PyObject>, Error> {
+ let v: Vec<HashablePyObject> = key
+ .iter()?
+ .map(|obj| HashablePyObject::new(obj?))
+ .collect::<Result<_, _>>()?;
+
+ Ok(self.0.get(&v)?)
+ }
+
+ // pub fn pop_node(&mut self, key: &PyAny) -> Result<Option<TreeCacheNode<K, PyObject>>, Error> {
+ // todo!()
+ // }
+
+ pub fn pop(&mut self, key: &PyAny) -> Result<Option<PyObject>, Error> {
+ let v: Vec<HashablePyObject> = key
+ .iter()?
+ .map(|obj| HashablePyObject::new(obj?))
+ .collect::<Result<_, _>>()?;
+
+ Ok(self.0.pop(&v)?)
+ }
+
+ pub fn clear(&mut self) {
+ self.0.clear()
+ }
+
+ pub fn len(&self) -> usize {
+ self.0.len()
+ }
+
+ pub fn values(&self) -> Vec<&PyObject> {
+ self.0.values().collect()
+ }
+
+ pub fn items(&self) -> Vec<(Vec<&HashablePyObject>, &PyObject)> {
+ todo!()
+ }
+}
diff --git a/rust/src/tree_cache/mod.rs b/rust/src/tree_cache/mod.rs
new file mode 100644
index 0000000000..0a4905b881
--- /dev/null
+++ b/rust/src/tree_cache/mod.rs
@@ -0,0 +1,360 @@
+use std::{collections::HashMap, hash::Hash};
+
+use anyhow::{bail, Error};
+
+pub mod binding;
+
+pub enum TreeCacheNode<K, V> {
+ Leaf(V),
+ Branch(usize, HashMap<K, TreeCacheNode<K, V>>),
+}
+
+impl<K, V> TreeCacheNode<K, V> {
+ pub fn new_branch() -> Self {
+ TreeCacheNode::Branch(0, Default::default())
+ }
+
+ fn len(&self) -> usize {
+ match self {
+ TreeCacheNode::Leaf(_) => 1,
+ TreeCacheNode::Branch(size, _) => *size,
+ }
+ }
+}
+
+impl<'a, K: Eq + Hash + 'a, V> TreeCacheNode<K, V> {
+ pub fn set(
+ &mut self,
+ mut key: impl Iterator<Item = K>,
+ value: V,
+ ) -> Result<(usize, usize), Error> {
+ if let Some(k) = key.next() {
+ match self {
+ TreeCacheNode::Leaf(_) => bail!("Given key is too long"),
+ TreeCacheNode::Branch(size, map) => {
+ let node = map.entry(k).or_insert_with(TreeCacheNode::new_branch);
+ let (added, removed) = node.set(key, value)?;
+
+ *size += added;
+ *size -= removed;
+
+ Ok((added, removed))
+ }
+ }
+ } else {
+ let added = if let TreeCacheNode::Branch(_, map) = self {
+ (1, map.len())
+ } else {
+ (0, 0)
+ };
+
+ *self = TreeCacheNode::Leaf(value);
+
+ Ok(added)
+ }
+ }
+
+ pub fn pop(
+ &mut self,
+ current_key: &K,
+ mut next_keys: impl Iterator<Item = &'a K>,
+ ) -> Result<Option<TreeCacheNode<K, V>>, Error> {
+ if let Some(next_key) = next_keys.next() {
+ match self {
+ TreeCacheNode::Leaf(_) => bail!("Given key is too long"),
+ TreeCacheNode::Branch(size, map) => {
+ let node = if let Some(node) = map.get_mut(current_key) {
+ node
+ } else {
+ return Ok(None);
+ };
+
+ if let Some(popped) = node.pop(next_key, next_keys)? {
+ *size -= node.len();
+
+ Ok(Some(popped))
+ } else {
+ Ok(None)
+ }
+ }
+ }
+ } else {
+ match self {
+ TreeCacheNode::Leaf(_) => bail!("Given key is too long"),
+ TreeCacheNode::Branch(size, map) => {
+ if let Some(node) = map.remove(current_key) {
+ *size -= node.len();
+
+ Ok(Some(node))
+ } else {
+ Ok(None)
+ }
+ }
+ }
+ }
+ }
+
+ pub fn items(&self) -> impl Iterator<Item = (Vec<&K>, &V)> {
+ let mut stack = vec![(vec![], self)];
+
+ std::iter::from_fn(move || {
+ while let Some((prefix, node)) = stack.pop() {
+ match node {
+ TreeCacheNode::Leaf(value) => return Some((prefix, value)),
+ TreeCacheNode::Branch(_, map) => {
+ stack.extend(map.iter().map(|(k, v)| {
+ let mut prefix = prefix.clone();
+ prefix.push(k);
+ (prefix, v)
+ }));
+ }
+ }
+ }
+
+ None
+ })
+ }
+}
+
+impl<K, V> Default for TreeCacheNode<K, V> {
+ fn default() -> Self {
+ TreeCacheNode::new_branch()
+ }
+}
+
+pub struct TreeCache<K, V> {
+ root: TreeCacheNode<K, V>,
+}
+
+impl<K, V> TreeCache<K, V> {
+ pub fn new() -> Self {
+ TreeCache {
+ root: TreeCacheNode::new_branch(),
+ }
+ }
+}
+
+impl<'a, K: Eq + Hash + 'a, V> TreeCache<K, V> {
+ pub fn set(&mut self, key: impl IntoIterator<Item = K>, value: V) -> Result<(), Error> {
+ self.root.set(key.into_iter(), value)?;
+
+ Ok(())
+ }
+
+ pub fn get_node(
+ &self,
+ key: impl IntoIterator<Item = &'a K>,
+ ) -> Result<Option<&TreeCacheNode<K, V>>, Error> {
+ let mut node = &self.root;
+
+ for k in key {
+ match node {
+ TreeCacheNode::Leaf(_) => bail!("Given key is too long"),
+ TreeCacheNode::Branch(_, map) => {
+ node = if let Some(node) = map.get(k) {
+ node
+ } else {
+ return Ok(None);
+ };
+ }
+ }
+ }
+
+ Ok(Some(node))
+ }
+
+ pub fn get(&self, key: impl IntoIterator<Item = &'a K>) -> Result<Option<&V>, Error> {
+ if let Some(node) = self.get_node(key)? {
+ match node {
+ TreeCacheNode::Leaf(value) => Ok(Some(value)),
+ TreeCacheNode::Branch(_, _) => bail!("Given key is too short"),
+ }
+ } else {
+ Ok(None)
+ }
+ }
+
+ pub fn pop_node(
+ &mut self,
+ key: impl IntoIterator<Item = &'a K>,
+ ) -> Result<Option<TreeCacheNode<K, V>>, Error> {
+ let mut key_iter = key.into_iter();
+
+ let k = if let Some(k) = key_iter.next() {
+ k
+ } else {
+ let node = std::mem::replace(&mut self.root, TreeCacheNode::new_branch());
+ return Ok(Some(node));
+ };
+
+ self.root.pop(k, key_iter)
+ }
+
+ pub fn pop(&mut self, key: &[K]) -> Result<Option<V>, Error> {
+ if let Some(node) = self.pop_node(key)? {
+ match node {
+ TreeCacheNode::Leaf(value) => Ok(Some(value)),
+ TreeCacheNode::Branch(_, _) => bail!("Given key is too short"),
+ }
+ } else {
+ Ok(None)
+ }
+ }
+
+ pub fn clear(&mut self) {
+ self.root = TreeCacheNode::new_branch();
+ }
+
+ pub fn len(&self) -> usize {
+ match self.root {
+ TreeCacheNode::Leaf(_) => 1,
+ TreeCacheNode::Branch(size, _) => size,
+ }
+ }
+
+ pub fn values(&self) -> impl Iterator<Item = &V> {
+ let mut stack = vec![&self.root];
+
+ std::iter::from_fn(move || {
+ while let Some(node) = stack.pop() {
+ match node {
+ TreeCacheNode::Leaf(value) => return Some(value),
+ TreeCacheNode::Branch(_, map) => {
+ stack.extend(map.values());
+ }
+ }
+ }
+
+ None
+ })
+ }
+
+ pub fn items(&self) -> impl Iterator<Item = (Vec<&K>, &V)> {
+ self.root.items()
+ }
+}
+
+impl<K, V> Default for TreeCache<K, V> {
+ fn default() -> Self {
+ TreeCache::new()
+ }
+}
+
+#[cfg(test)]
+mod test {
+ use std::collections::BTreeSet;
+
+ use super::*;
+
+ #[test]
+ fn get_set() -> Result<(), Error> {
+ let mut cache = TreeCache::new();
+
+ cache.set(vec!["a", "b"], "c")?;
+
+ assert_eq!(cache.get(&["a", "b"])?, Some(&"c"));
+
+ let node = cache.get_node(&["a"])?.unwrap();
+
+ match node {
+ TreeCacheNode::Leaf(_) => bail!("expected branch"),
+ TreeCacheNode::Branch(_, map) => {
+ assert_eq!(map.len(), 1);
+ assert!(map.contains_key("b"));
+ }
+ }
+
+ Ok(())
+ }
+
+ #[test]
+ fn length() -> Result<(), Error> {
+ let mut cache = TreeCache::new();
+
+ cache.set(vec!["a", "b"], "c")?;
+
+ assert_eq!(cache.len(), 1);
+
+ cache.set(vec!["a", "b"], "d")?;
+
+ assert_eq!(cache.len(), 1);
+
+ cache.set(vec!["e", "f"], "g")?;
+
+ assert_eq!(cache.len(), 2);
+
+ cache.set(vec!["e", "h"], "i")?;
+
+ assert_eq!(cache.len(), 3);
+
+ cache.set(vec!["e"], "i")?;
+
+ assert_eq!(cache.len(), 2);
+
+ cache.pop_node(&["a"])?;
+
+ assert_eq!(cache.len(), 1);
+
+ Ok(())
+ }
+
+ #[test]
+ fn clear() -> Result<(), Error> {
+ let mut cache = TreeCache::new();
+
+ cache.set(vec!["a", "b"], "c")?;
+
+ assert_eq!(cache.len(), 1);
+
+ cache.clear();
+
+ assert_eq!(cache.len(), 0);
+
+ assert_eq!(cache.get(&["a", "b"])?, None);
+
+ Ok(())
+ }
+
+ #[test]
+ fn pop() -> Result<(), Error> {
+ let mut cache = TreeCache::new();
+
+ cache.set(vec!["a", "b"], "c")?;
+ assert_eq!(cache.pop(&["a", "b"])?, Some("c"));
+ assert_eq!(cache.pop(&["a", "b"])?, None);
+
+ Ok(())
+ }
+
+ #[test]
+ fn values() -> Result<(), Error> {
+ let mut cache = TreeCache::new();
+
+ cache.set(vec!["a", "b"], "c")?;
+
+ let expected = ["c"].iter().collect();
+ assert_eq!(cache.values().collect::<BTreeSet<_>>(), expected);
+
+ cache.set(vec!["d", "e"], "f")?;
+
+ let expected = ["c", "f"].iter().collect();
+ assert_eq!(cache.values().collect::<BTreeSet<_>>(), expected);
+
+ Ok(())
+ }
+
+ #[test]
+ fn items() -> Result<(), Error> {
+ let mut cache = TreeCache::new();
+
+ cache.set(vec!["a", "b"], "c")?;
+ cache.set(vec!["d", "e"], "f")?;
+
+ let expected = [(vec![&"a", &"b"], &"c"), (vec![&"d", &"e"], &"f")]
+ .into_iter()
+ .collect();
+ assert_eq!(cache.items().collect::<BTreeSet<_>>(), expected);
+
+ Ok(())
+ }
+}
|
