mirror of
https://github.com/element-hq/synapse
synced 2024-08-20 10:50:21 +00:00
bindings
This commit is contained in:
parent
4874d6320a
commit
18ac015ecd
3 changed files with 146 additions and 1 deletions
|
@ -27,6 +27,7 @@ fn synapse_rust(py: Python<'_>, m: &PyModule) -> PyResult<()> {
|
|||
m.add_function(wrap_pyfunction!(get_rust_file_digest, m)?)?;
|
||||
|
||||
push::register_module(py, m)?;
|
||||
tree_cache::binding::register_module(py, m)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
|
128
rust/src/tree_cache/binding.rs
Normal file
128
rust/src/tree_cache/binding.rs
Normal file
|
@ -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!()
|
||||
}
|
||||
}
|
|
@ -2,6 +2,8 @@ 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>>),
|
||||
|
@ -114,17 +116,25 @@ impl<'a, K: Eq + Hash + 'a, V> TreeCacheNode<K, V> {
|
|||
}
|
||||
}
|
||||
|
||||
impl<K, V> Default for TreeCacheNode<K, V> {
|
||||
fn default() -> Self {
|
||||
TreeCacheNode::new_branch()
|
||||
}
|
||||
}
|
||||
|
||||
pub struct TreeCache<K, V> {
|
||||
root: TreeCacheNode<K, V>,
|
||||
}
|
||||
|
||||
impl<'a, K: Eq + Hash + 'a, V> TreeCache<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)?;
|
||||
|
||||
|
@ -224,6 +234,12 @@ impl<'a, K: Eq + Hash + 'a, V> TreeCache<K, V> {
|
|||
}
|
||||
}
|
||||
|
||||
impl<K, V> Default for TreeCache<K, V> {
|
||||
fn default() -> Self {
|
||||
TreeCache::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use std::collections::BTreeSet;
|
Loading…
Reference in a new issue