rust/hg-cpython/src/dirstate/dirstate_map.rs

rust-cpython: mark unsafe functions as such

It wasn't trivial to fix leak_immutable() to be safe since we have to
allow immutable operations (e.g. iter()) on the leaked reference. So
let's mark it unsafe for now. Callers must take care of the returned
object to guarantee the memory safety.

I'll revisit this later. I think $leaked<T: 'static> could have a function
that converts itself into $leaked<U: 'static> with a given FnOnce(&T) -> &U,
where T is $inner_struct, and U is $iterator_type for example.

// dirstate_map.rs
//
// Copyright 2019 Raphaël Gomès <[e-mail]>
//
// This software may be used and distributed according to the terms of the
// GNU General Public License version 2 or any later version.

//! Bindings for the `hg::dirstate::dirstate_map` file provided by the
//! `hg-core` package.

use std::cell::RefCell;
use std::convert::TryInto;
use std::time::Duration;

use cpython::{
    exc, ObjectProtocol, PyBool, PyBytes, PyClone, PyDict, PyErr, PyObject,
    PyResult, PyTuple, Python, PythonObject, ToPyObject,
};
use libc::c_char;

use crate::{
    dirstate::copymap::{CopyMap, CopyMapItemsIterator, CopyMapKeysIterator},
    dirstate::{decapsule_make_dirstate_tuple, dirs_multiset::Dirs},
    ref_sharing::{PySharedRefCell, PySharedState},
};
use hg::{
    DirsMultiset, DirstateEntry, DirstateMap as RustDirstateMap,
    DirstateParents, DirstateParseError, EntryState, PARENT_SIZE,
};

// TODO
//     This object needs to share references to multiple members of its Rust
//     inner struct, namely `copy_map`, `dirs` and `all_dirs`.
//     Right now `CopyMap` is done, but it needs to have an explicit reference
//     to `RustDirstateMap` which itself needs to have an encapsulation for
//     every method in `CopyMap` (copymapcopy, etc.).
//     This is ugly and hard to maintain.
//     The same logic applies to `dirs` and `all_dirs`, however the `Dirs`
//     `py_class!` is already implemented and does not mention
//     `RustDirstateMap`, rightfully so.
//     All attributes also have to have a separate refcount data attribute for
//     leaks, with all methods that go along for reference sharing.
py_class!(pub class DirstateMap |py| {
    data inner: PySharedRefCell<RustDirstateMap>;
    data py_shared_state: PySharedState;

    def __new__(_cls, _root: PyObject) -> PyResult<Self> {
        let inner = RustDirstateMap::default();
        Self::create_instance(
            py,
            PySharedRefCell::new(inner),
            PySharedState::default()
        )
    }

    def clear(&self) -> PyResult<PyObject> {
        self.borrow_mut(py)?.clear();
        Ok(py.None())
    }

    def get(
        &self,
        key: PyObject,
        default: Option<PyObject> = None
    ) -> PyResult<Option<PyObject>> {
        let key = key.extract::<PyBytes>(py)?;
        match self.inner(py).borrow().get(key.data(py)) {
            Some(entry) => {
                // Explicitly go through u8 first, then cast to
                // platform-specific `c_char`.
                let state: u8 = entry.state.into();
                Ok(Some(decapsule_make_dirstate_tuple(py)?(
                        state as c_char,
                        entry.mode,
                        entry.size,
                        entry.mtime,
                    )))
            },
            None => Ok(default)
        }
    }

    def addfile(
        &self,
        f: PyObject,
        oldstate: PyObject,
        state: PyObject,
        mode: PyObject,
        size: PyObject,
        mtime: PyObject
    ) -> PyResult<PyObject> {
        self.borrow_mut(py)?.add_file(
            f.extract::<PyBytes>(py)?.data(py),
            oldstate.extract::<PyBytes>(py)?.data(py)[0]
                .try_into()
                .map_err(|e: DirstateParseError| {
                    PyErr::new::<exc::ValueError, _>(py, e.to_string())
                })?,
            DirstateEntry {
                state: state.extract::<PyBytes>(py)?.data(py)[0]
                    .try_into()
                    .map_err(|e: DirstateParseError| {
                        PyErr::new::<exc::ValueError, _>(py, e.to_string())
                    })?,
                mode: mode.extract(py)?,
                size: size.extract(py)?,
                mtime: mtime.extract(py)?,
            },
        );
        Ok(py.None())
    }

    def removefile(
        &self,
        f: PyObject,
        oldstate: PyObject,
        size: PyObject
    ) -> PyResult<PyObject> {
        self.borrow_mut(py)?
            .remove_file(
                f.extract::<PyBytes>(py)?.data(py),
                oldstate.extract::<PyBytes>(py)?.data(py)[0]
                    .try_into()
                    .map_err(|e: DirstateParseError| {
                        PyErr::new::<exc::ValueError, _>(py, e.to_string())
                    })?,
                size.extract(py)?,
            )
            .or_else(|_| {
                Err(PyErr::new::<exc::OSError, _>(
                    py,
                    "Dirstate error".to_string(),
                ))
            })?;
        Ok(py.None())
    }

    def dropfile(
        &self,
        f: PyObject,
        oldstate: PyObject
    ) -> PyResult<PyBool> {
        self.borrow_mut(py)?
            .drop_file(
                f.extract::<PyBytes>(py)?.data(py),
                oldstate.extract::<PyBytes>(py)?.data(py)[0]
                    .try_into()
                    .map_err(|e: DirstateParseError| {
                        PyErr::new::<exc::ValueError, _>(py, e.to_string())
                    })?,
            )
            .and_then(|b| Ok(b.to_py_object(py)))
            .or_else(|_| {
                Err(PyErr::new::<exc::OSError, _>(
                    py,
                    "Dirstate error".to_string(),
                ))
            })
    }

    def clearambiguoustimes(
        &self,
        files: PyObject,
        now: PyObject
    ) -> PyResult<PyObject> {
        let files: PyResult<Vec<Vec<u8>>> = files
            .iter(py)?
            .map(|filename| {
                Ok(filename?.extract::<PyBytes>(py)?.data(py).to_owned())
            })
            .collect();
        self.borrow_mut(py)?
            .clear_ambiguous_times(files?, now.extract(py)?);
        Ok(py.None())
    }

    // TODO share the reference
    def nonnormalentries(&self) -> PyResult<PyObject> {
        let (non_normal, other_parent) =
            self.inner(py).borrow().non_normal_other_parent_entries();

        let locals = PyDict::new(py);
        locals.set_item(
            py,
            "non_normal",
            non_normal
                .iter()
                .map(|v| PyBytes::new(py, &v))
                .collect::<Vec<PyBytes>>()
                .to_py_object(py),
        )?;
        locals.set_item(
            py,
            "other_parent",
            other_parent
                .iter()
                .map(|v| PyBytes::new(py, &v))
                .collect::<Vec<PyBytes>>()
                .to_py_object(py),
        )?;

        py.eval("set(non_normal), set(other_parent)", None, Some(&locals))
    }

    def hastrackeddir(&self, d: PyObject) -> PyResult<PyBool> {
        let d = d.extract::<PyBytes>(py)?;
        Ok(self.borrow_mut(py)?
            .has_tracked_dir(d.data(py))
            .to_py_object(py))
    }

    def hasdir(&self, d: PyObject) -> PyResult<PyBool> {
        let d = d.extract::<PyBytes>(py)?;
        Ok(self.borrow_mut(py)?
            .has_dir(d.data(py))
            .to_py_object(py))
    }

    def parents(&self, st: PyObject) -> PyResult<PyTuple> {
        self.borrow_mut(py)?
            .parents(st.extract::<PyBytes>(py)?.data(py))
            .and_then(|d| {
                Ok((PyBytes::new(py, &d.p1), PyBytes::new(py, &d.p2))
                    .to_py_object(py))
            })
            .or_else(|_| {
                Err(PyErr::new::<exc::OSError, _>(
                    py,
                    "Dirstate error".to_string(),
                ))
            })
    }

    def setparents(&self, p1: PyObject, p2: PyObject) -> PyResult<PyObject> {
        let p1 = extract_node_id(py, &p1)?;
        let p2 = extract_node_id(py, &p2)?;

        self.borrow_mut(py)?
            .set_parents(&DirstateParents { p1, p2 });
        Ok(py.None())
    }

    def read(&self, st: PyObject) -> PyResult<Option<PyObject>> {
        match self.borrow_mut(py)?
            .read(st.extract::<PyBytes>(py)?.data(py))
        {
            Ok(Some(parents)) => Ok(Some(
                (PyBytes::new(py, &parents.p1), PyBytes::new(py, &parents.p2))
                    .to_py_object(py)
                    .into_object(),
            )),
            Ok(None) => Ok(Some(py.None())),
            Err(_) => Err(PyErr::new::<exc::OSError, _>(
                py,
                "Dirstate error".to_string(),
            )),
        }
    }
    def write(
        &self,
        p1: PyObject,
        p2: PyObject,
        now: PyObject
    ) -> PyResult<PyBytes> {
        let now = Duration::new(now.extract(py)?, 0);
        let parents = DirstateParents {
            p1: extract_node_id(py, &p1)?,
            p2: extract_node_id(py, &p2)?,
        };

        match self.borrow_mut(py)?.pack(parents, now) {
            Ok(packed) => Ok(PyBytes::new(py, &packed)),
            Err(_) => Err(PyErr::new::<exc::OSError, _>(
                py,
                "Dirstate error".to_string(),
            )),
        }
    }

    def filefoldmapasdict(&self) -> PyResult<PyDict> {
        let dict = PyDict::new(py);
        for (key, value) in
            self.borrow_mut(py)?.build_file_fold_map().iter()
        {
            dict.set_item(py, key, value)?;
        }
        Ok(dict)
    }

    def __len__(&self) -> PyResult<usize> {
        Ok(self.inner(py).borrow().len())
    }

    def __contains__(&self, key: PyObject) -> PyResult<bool> {
        let key = key.extract::<PyBytes>(py)?;
        Ok(self.inner(py).borrow().contains_key(key.data(py)))
    }

    def __getitem__(&self, key: PyObject) -> PyResult<PyObject> {
        let key = key.extract::<PyBytes>(py)?;
        let key = key.data(py);
        match self.inner(py).borrow().get(key) {
            Some(entry) => {
                // Explicitly go through u8 first, then cast to
                // platform-specific `c_char`.
                let state: u8 = entry.state.into();
                Ok(decapsule_make_dirstate_tuple(py)?(
                        state as c_char,
                        entry.mode,
                        entry.size,
                        entry.mtime,
                    ))
            },
            None => Err(PyErr::new::<exc::KeyError, _>(
                py,
                String::from_utf8_lossy(key),
            )),
        }
    }

    def keys(&self) -> PyResult<DirstateMapKeysIterator> {
        let (leak_handle, leaked_ref) = unsafe { self.leak_immutable(py)? };
        DirstateMapKeysIterator::from_inner(
            py,
            Some(leak_handle),
            Box::new(leaked_ref.iter()),
        )
    }

    def items(&self) -> PyResult<DirstateMapItemsIterator> {
        let (leak_handle, leaked_ref) = unsafe { self.leak_immutable(py)? };
        DirstateMapItemsIterator::from_inner(
            py,
            Some(leak_handle),
            Box::new(leaked_ref.iter()),
        )
    }

    def __iter__(&self) -> PyResult<DirstateMapKeysIterator> {
        let (leak_handle, leaked_ref) = unsafe { self.leak_immutable(py)? };
        DirstateMapKeysIterator::from_inner(
            py,
            Some(leak_handle),
            Box::new(leaked_ref.iter()),
        )
    }

    def getdirs(&self) -> PyResult<Dirs> {
        // TODO don't copy, share the reference
        self.borrow_mut(py)?.set_dirs();
        Dirs::from_inner(
            py,
            DirsMultiset::from_dirstate(
                &self.inner(py).borrow(),
                Some(EntryState::Removed),
            ),
        )
    }
    def getalldirs(&self) -> PyResult<Dirs> {
        // TODO don't copy, share the reference
        self.borrow_mut(py)?.set_all_dirs();
        Dirs::from_inner(
            py,
            DirsMultiset::from_dirstate(
                &self.inner(py).borrow(),
                None,
            ),
        )
    }

    // TODO all copymap* methods, see docstring above
    def copymapcopy(&self) -> PyResult<PyDict> {
        let dict = PyDict::new(py);
        for (key, value) in self.inner(py).borrow().copy_map.iter() {
            dict.set_item(py, PyBytes::new(py, key), PyBytes::new(py, value))?;
        }
        Ok(dict)
    }

    def copymapgetitem(&self, key: PyObject) -> PyResult<PyBytes> {
        let key = key.extract::<PyBytes>(py)?;
        match self.inner(py).borrow().copy_map.get(key.data(py)) {
            Some(copy) => Ok(PyBytes::new(py, copy)),
            None => Err(PyErr::new::<exc::KeyError, _>(
                py,
                String::from_utf8_lossy(key.data(py)),
            )),
        }
    }
    def copymap(&self) -> PyResult<CopyMap> {
        CopyMap::from_inner(py, self.clone_ref(py))
    }

    def copymaplen(&self) -> PyResult<usize> {
        Ok(self.inner(py).borrow().copy_map.len())
    }
    def copymapcontains(&self, key: PyObject) -> PyResult<bool> {
        let key = key.extract::<PyBytes>(py)?;
        Ok(self.inner(py).borrow().copy_map.contains_key(key.data(py)))
    }
    def copymapget(
        &self,
        key: PyObject,
        default: Option<PyObject>
    ) -> PyResult<Option<PyObject>> {
        let key = key.extract::<PyBytes>(py)?;
        match self.inner(py).borrow().copy_map.get(key.data(py)) {
            Some(copy) => Ok(Some(PyBytes::new(py, copy).into_object())),
            None => Ok(default),
        }
    }
    def copymapsetitem(
        &self,
        key: PyObject,
        value: PyObject
    ) -> PyResult<PyObject> {
        let key = key.extract::<PyBytes>(py)?;
        let value = value.extract::<PyBytes>(py)?;
        self.borrow_mut(py)?
            .copy_map
            .insert(key.data(py).to_vec(), value.data(py).to_vec());
        Ok(py.None())
    }
    def copymappop(
        &self,
        key: PyObject,
        default: Option<PyObject>
    ) -> PyResult<Option<PyObject>> {
        let key = key.extract::<PyBytes>(py)?;
        match self.borrow_mut(py)?.copy_map.remove(key.data(py)) {
            Some(_) => Ok(None),
            None => Ok(default),
        }
    }

    def copymapiter(&self) -> PyResult<CopyMapKeysIterator> {
        let (leak_handle, leaked_ref) = unsafe { self.leak_immutable(py)? };
        CopyMapKeysIterator::from_inner(
            py,
            Some(leak_handle),
            Box::new(leaked_ref.copy_map.iter()),
        )
    }

    def copymapitemsiter(&self) -> PyResult<CopyMapItemsIterator> {
        let (leak_handle, leaked_ref) = unsafe { self.leak_immutable(py)? };
        CopyMapItemsIterator::from_inner(
            py,
            Some(leak_handle),
            Box::new(leaked_ref.copy_map.iter()),
        )
    }

});

impl DirstateMap {
    fn translate_key(
        py: Python,
        res: (&Vec<u8>, &DirstateEntry),
    ) -> PyResult<Option<PyBytes>> {
        Ok(Some(PyBytes::new(py, res.0)))
    }
    fn translate_key_value(
        py: Python,
        res: (&Vec<u8>, &DirstateEntry),
    ) -> PyResult<Option<(PyBytes, PyObject)>> {
        let (f, entry) = res;

        // Explicitly go through u8 first, then cast to
        // platform-specific `c_char`.
        let state: u8 = entry.state.into();
        Ok(Some((
            PyBytes::new(py, f),
            decapsule_make_dirstate_tuple(py)?(
                state as c_char,
                entry.mode,
                entry.size,
                entry.mtime,
            ),
        )))
    }
}

py_shared_ref!(DirstateMap, RustDirstateMap, inner, DirstateMapLeakedRef,);

py_shared_mapping_iterator!(
    DirstateMapKeysIterator,
    DirstateMapLeakedRef,
    Vec<u8>,
    DirstateEntry,
    DirstateMap::translate_key,
    Option<PyBytes>
);

py_shared_mapping_iterator!(
    DirstateMapItemsIterator,
    DirstateMapLeakedRef,
    Vec<u8>,
    DirstateEntry,
    DirstateMap::translate_key_value,
    Option<(PyBytes, PyObject)>
);

fn extract_node_id(py: Python, obj: &PyObject) -> PyResult<[u8; PARENT_SIZE]> {
    let bytes = obj.extract::<PyBytes>(py)?;
    match bytes.data(py).try_into() {
        Ok(s) => Ok(s),
        Err(e) => Err(PyErr::new::<exc::ValueError, _>(py, e.to_string())),
    }
}