mirror of
https://github.com/element-hq/synapse
synced 2024-10-05 18:52:45 +00:00
Implement {get,pop}_node
This commit is contained in:
parent
b9cdf3d85e
commit
c03785e121
3 changed files with 110 additions and 18 deletions
|
@ -3,3 +3,7 @@
|
|||
|
||||
[workspace]
|
||||
members = ["rust"]
|
||||
|
||||
[profile.dbgrelease]
|
||||
inherits = "release"
|
||||
debug = true
|
||||
|
|
|
@ -3,7 +3,7 @@ use std::hash::Hash;
|
|||
use anyhow::Error;
|
||||
use pyo3::{
|
||||
pyclass, pymethods,
|
||||
types::{PyIterator, PyModule},
|
||||
types::{PyModule, PyTuple},
|
||||
IntoPy, PyAny, PyObject, PyResult, Python, ToPyObject,
|
||||
};
|
||||
|
||||
|
@ -25,6 +25,7 @@ pub fn register_module(py: Python<'_>, m: &PyModule) -> PyResult<()> {
|
|||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct HashablePyObject {
|
||||
obj: PyObject,
|
||||
hash: isize,
|
||||
|
@ -41,12 +42,24 @@ impl HashablePyObject {
|
|||
}
|
||||
}
|
||||
|
||||
impl IntoPy<PyObject> for HashablePyObject {
|
||||
fn into_py(self, _: Python<'_>) -> PyObject {
|
||||
self.obj.clone()
|
||||
}
|
||||
}
|
||||
|
||||
impl IntoPy<PyObject> for &HashablePyObject {
|
||||
fn into_py(self, _: Python<'_>) -> PyObject {
|
||||
self.obj.clone()
|
||||
}
|
||||
}
|
||||
|
||||
impl ToPyObject for HashablePyObject {
|
||||
fn to_object(&self, _py: Python<'_>) -> PyObject {
|
||||
self.obj.clone()
|
||||
}
|
||||
}
|
||||
|
||||
impl Hash for HashablePyObject {
|
||||
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
|
||||
self.hash.hash(state);
|
||||
|
@ -87,9 +100,31 @@ impl PythonTreeCache {
|
|||
Ok(())
|
||||
}
|
||||
|
||||
// pub fn get_node(&self, key: &PyAny) -> Result<Option<&TreeCacheNode<K, PyObject>>, Error> {
|
||||
// todo!()
|
||||
// }
|
||||
pub fn get_node<'a>(
|
||||
&'a self,
|
||||
py: Python<'a>,
|
||||
key: &'a PyAny,
|
||||
) -> Result<Option<Vec<(&'a PyTuple, &'a PyObject)>>, Error> {
|
||||
let v: Vec<HashablePyObject> = key
|
||||
.iter()?
|
||||
.map(|obj| HashablePyObject::new(obj?))
|
||||
.collect::<Result<_, _>>()?;
|
||||
|
||||
let Some(node) = self.0.get_node(v.clone())? else {
|
||||
return Ok(None)
|
||||
};
|
||||
|
||||
let items = node
|
||||
.items()
|
||||
.map(|(k, value)| {
|
||||
let vec = v.iter().chain(k.iter().map(|a| *a)).collect::<Vec<_>>();
|
||||
let nk = PyTuple::new(py, vec);
|
||||
(nk, value)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
Ok(Some(items))
|
||||
}
|
||||
|
||||
pub fn get(&self, key: &PyAny) -> Result<Option<&PyObject>, Error> {
|
||||
let v: Vec<HashablePyObject> = key
|
||||
|
@ -100,9 +135,31 @@ impl PythonTreeCache {
|
|||
Ok(self.0.get(&v)?)
|
||||
}
|
||||
|
||||
// pub fn pop_node(&mut self, key: &PyAny) -> Result<Option<TreeCacheNode<K, PyObject>>, Error> {
|
||||
// todo!()
|
||||
// }
|
||||
pub fn pop_node<'a>(
|
||||
&'a mut self,
|
||||
py: Python<'a>,
|
||||
key: &'a PyAny,
|
||||
) -> Result<Option<Vec<(&'a PyTuple, PyObject)>>, Error> {
|
||||
let v: Vec<HashablePyObject> = key
|
||||
.iter()?
|
||||
.map(|obj| HashablePyObject::new(obj?))
|
||||
.collect::<Result<_, _>>()?;
|
||||
|
||||
let Some(node) = self.0.pop_node(v.clone())? else {
|
||||
return Ok(None)
|
||||
};
|
||||
|
||||
let items = node
|
||||
.into_items()
|
||||
.map(|(k, value)| {
|
||||
let vec = v.iter().chain(k.iter()).collect::<Vec<_>>();
|
||||
let nk = PyTuple::new(py, vec);
|
||||
(nk, value)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
Ok(Some(items))
|
||||
}
|
||||
|
||||
pub fn pop(&mut self, key: &PyAny) -> Result<Option<PyObject>, Error> {
|
||||
let v: Vec<HashablePyObject> = key
|
||||
|
|
|
@ -54,16 +54,20 @@ impl<'a, K: Eq + Hash + 'a, V> TreeCacheNode<K, V> {
|
|||
}
|
||||
}
|
||||
|
||||
pub fn pop(
|
||||
pub fn pop<Q>(
|
||||
&mut self,
|
||||
current_key: &K,
|
||||
mut next_keys: impl Iterator<Item = &'a K>,
|
||||
) -> Result<Option<TreeCacheNode<K, V>>, Error> {
|
||||
current_key: Q,
|
||||
mut next_keys: impl Iterator<Item = Q>,
|
||||
) -> Result<Option<TreeCacheNode<K, V>>, Error>
|
||||
where
|
||||
Q: Borrow<K>,
|
||||
Q: Hash + Eq + 'a,
|
||||
{
|
||||
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) {
|
||||
let node = if let Some(node) = map.get_mut(current_key.borrow()) {
|
||||
node
|
||||
} else {
|
||||
return Ok(None);
|
||||
|
@ -82,7 +86,7 @@ impl<'a, K: Eq + Hash + 'a, V> TreeCacheNode<K, V> {
|
|||
match self {
|
||||
TreeCacheNode::Leaf(_) => bail!("Given key is too long"),
|
||||
TreeCacheNode::Branch(size, map) => {
|
||||
if let Some(node) = map.remove(current_key) {
|
||||
if let Some(node) = map.remove(current_key.borrow()) {
|
||||
*size -= node.len();
|
||||
|
||||
Ok(Some(node))
|
||||
|
@ -94,8 +98,8 @@ impl<'a, K: Eq + Hash + 'a, V> TreeCacheNode<K, V> {
|
|||
}
|
||||
}
|
||||
|
||||
pub fn items(&self) -> impl Iterator<Item = (Vec<&K>, &V)> {
|
||||
let mut stack = vec![(vec![], self)];
|
||||
pub fn items(&'a self) -> impl Iterator<Item = (Vec<&K>, &V)> {
|
||||
let mut stack = vec![(Vec::new(), self)];
|
||||
|
||||
std::iter::from_fn(move || {
|
||||
while let Some((prefix, node)) = stack.pop() {
|
||||
|
@ -116,6 +120,29 @@ impl<'a, K: Eq + Hash + 'a, V> TreeCacheNode<K, V> {
|
|||
}
|
||||
}
|
||||
|
||||
impl<'a, K: Clone + Eq + Hash + 'a, V> TreeCacheNode<K, V> {
|
||||
pub fn into_items(self) -> impl Iterator<Item = (Vec<K>, V)> {
|
||||
let mut stack = vec![(Vec::new(), 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.into_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()
|
||||
|
@ -182,10 +209,14 @@ impl<'a, K: Eq + Hash + 'a, V> TreeCache<K, V> {
|
|||
}
|
||||
}
|
||||
|
||||
pub fn pop_node(
|
||||
pub fn pop_node<Q>(
|
||||
&mut self,
|
||||
key: impl IntoIterator<Item = &'a K>,
|
||||
) -> Result<Option<TreeCacheNode<K, V>>, Error> {
|
||||
key: impl IntoIterator<Item = Q>,
|
||||
) -> Result<Option<TreeCacheNode<K, V>>, Error>
|
||||
where
|
||||
Q: Borrow<K>,
|
||||
Q: Hash + Eq + 'a,
|
||||
{
|
||||
let mut key_iter = key.into_iter();
|
||||
|
||||
let k = if let Some(k) = key_iter.next() {
|
||||
|
|
Loading…
Reference in a new issue