trezor-firmware/core/embed/rust/src/micropython/map.rs

use core::{marker::PhantomData, mem::MaybeUninit, ops::Deref, ptr, slice};

use crate::{
    error::Error,
    micropython::{obj::Obj, qstr::Qstr},
};

use super::{ffi, runtime::catch_exception};

pub type Map = ffi::mp_map_t;
pub type MapElem = ffi::mp_map_elem_t;

impl Map {
    pub const fn from_fixed_static<const N: usize>(table: &'static [MapElem; N]) -> Self {
        // We can't use the regular accessors generated by bindgen for us, as they are
        // not `const`. Instead, we build the bitfield manually.

        // micropython/py/obj.h MP_DEFINE_CONST_DICT
        // .all_keys_are_qstrs = 1,
        // .is_fixed = 1,
        // .is_ordered = 1,
        // .used = MP_ARRAY_SIZE(table_name),
        // .alloc = MP_ARRAY_SIZE(table_name),
        let bits = 0b111 | N << 3;
        let bitfield = ffi::__BindgenBitfieldUnit::new(bits.to_ne_bytes());
        Self {
            _bitfield_align_1: [],
            _bitfield_1: bitfield,
            alloc: N,
            table: table.as_ptr() as *mut MapElem,
        }
    }

    pub const fn at(key: Qstr, value: Obj) -> MapElem {
        MapElem {
            key: key.to_obj(),
            value,
        }
    }

    pub const EMPTY: Map = Self::from_fixed_static::<0>(&[]);
}

impl Map {
    pub fn from_fixed(table: &[MapElem]) -> MapRef {
        let mut map = MaybeUninit::uninit();
        // SAFETY: `mp_map_init_fixed_table` completely initializes all fields of `map`.
        unsafe {
            // EXCEPTION: Does not raise.
            ffi::mp_map_init_fixed_table(map.as_mut_ptr(), table.len(), table.as_ptr().cast());
            MapRef::new(map.assume_init())
        }
    }

    pub fn with_capacity(capacity: usize) -> Result<Self, Error> {
        let mut map = MaybeUninit::uninit();
        // SAFETY: `mp_map_init` completely initializes all fields of `map`, allocating
        // the backing storage for `capacity` items on the heap.
        unsafe {
            // EXCEPTION: Will raise if allocation fails.
            catch_exception(|| {
                ffi::mp_map_init(map.as_mut_ptr(), capacity);
            })?;
            Ok(map.assume_init())
        }
    }

    pub fn len(&self) -> usize {
        self.used()
    }

    pub fn elems(&self) -> &[MapElem] {
        // SAFETY: `self.table` should always point to an array of `MapElem` of
        // `self.len()` items valid at least for the lifetime of `self`.
        unsafe { slice::from_raw_parts(self.table, self.len()) }
    }

    pub fn contains_key(&self, index: impl Into<Obj>) -> bool {
        self.get_obj(index.into()).is_ok()
    }

    pub fn get(&self, index: impl Into<Obj>) -> Result<Obj, Error> {
        self.get_obj(index.into())
    }

    pub fn get_obj(&self, index: Obj) -> Result<Obj, Error> {
        // SAFETY:
        //  - `mp_map_lookup` returns either NULL or a pointer to a `mp_map_elem_t`
        //    value with a lifetime valid for the whole lifetime of the passed immutable
        //    ref of `map`.
        //  - with `_mp_map_lookup_kind_t_MP_MAP_LOOKUP`, `map` stays unmodified and the
        //    cast to mut ptr is therefore safe.
        unsafe {
            let map = self as *const Self as *mut Self;
            // EXCEPTION: Does not raise for MP_MAP_LOOKUP.
            let elem = ffi::mp_map_lookup(map, index, ffi::_mp_map_lookup_kind_t_MP_MAP_LOOKUP)
                .as_ref()
                .ok_or(Error::KeyError(index))?;
            Ok(elem.value)
        }
    }

    pub fn get_or<T>(&self, index: impl Into<Obj>, default: T) -> Result<T, Error>
    where
        T: TryFrom<Obj, Error = Error>,
    {
        let res = self.get(index);
        match res {
            Ok(obj) => obj.try_into(),
            Err(Error::KeyError(_)) => Ok(default),
            Err(e) => Err(e),
        }
    }

    pub fn set(&mut self, index: impl Into<Obj>, value: impl Into<Obj>) -> Result<(), Error> {
        self.set_obj(index.into(), value.into())
    }

    pub fn set_obj(&mut self, index: Obj, value: Obj) -> Result<(), Error> {
        // SAFETY:
        //  - `mp_map_lookup` with `_mp_map_lookup_kind_t_MP_MAP_LOOKUP_ADD_IF_NOT_FOUND
        //    returns a pointer to a `mp_map_elem_t` value with a lifetime valid for the
        //    whole lifetime of `&mut self`.
        //  - adding an element is an allocation, so it might raise.
        //  - the original `elem.value` might be an `Obj` that will get GCd when we
        //    replace it.
        unsafe {
            let map = self as *mut Self;
            // EXCEPTION: Will raise if allocation fails.
            let elem = catch_exception(|| {
                ffi::mp_map_lookup(
                    map,
                    index,
                    ffi::_mp_map_lookup_kind_t_MP_MAP_LOOKUP_ADD_IF_NOT_FOUND,
                )
            })?
            .as_mut()
            .unwrap(); // `MP_MAP_LOOKUP_ADD_IF_NOT_FOUND` should always return a non-null pointer.
            elem.value = value;
        }
        Ok(())
    }

    pub fn delete(&mut self, index: impl Into<Obj>) {
        self.delete_obj(index.into())
    }

    pub fn delete_obj(&mut self, index: Obj) {
        unsafe {
            let map = self as *mut Self;
            // EXCEPTION: Does not raise for MP_MAP_LOOKUP_REMOVE_IF_FOUND.
            ffi::mp_map_lookup(
                map,
                index,
                ffi::_mp_map_lookup_kind_t_MP_MAP_LOOKUP_REMOVE_IF_FOUND,
            );
        }
    }
}

impl Map {
    pub fn try_clone(&self) -> Result<Self, Error> {
        let mut map = Self::with_capacity(self.len())?;
        unsafe {
            ptr::copy_nonoverlapping(self.table, map.table, self.len());
        }
        map.set_used(self.used());
        map.set_all_keys_are_qstrs(self.all_keys_are_qstrs());
        map.set_is_ordered(self.is_ordered());
        map.set_is_fixed(0);
        Ok(map)
    }
}

// SAFETY: We are in a single-threaded environment.
unsafe impl Sync for Map {}
unsafe impl Sync for MapElem {}

/// Wrapper over a private `Map` value with element lifetime `'a`.
///
/// `Map` itself does not have any lifetime parameter, so it can effectively
/// only reference element tables with `'static` lifetime. `MapRef` is tying a
/// private `Map` value to a lifetime `'a`, exposing it only by-ref within the
/// `'a` lifetime. It's not possible to move the value out from the reference,
/// and because `Map` is not `Clone`, it is also impossible to clone it. All
/// this only holds for immutable ref, therefore we only impl `Deref`, not
/// `DerefMut`.
//
// TODO: Better solution would be to make Map a wrapper type instead of a type
// alias and add a lifetime parameter.
pub struct MapRef<'a> {
    map: Map,
    _phantom: PhantomData<&'a Map>,
}

impl<'a> MapRef<'a> {
    fn new(map: Map) -> Self {
        Self {
            map,
            _phantom: PhantomData,
        }
    }
}

impl<'a> Deref for MapRef<'a> {
    type Target = Map;

    fn deref(&self) -> &Self::Target {
        &self.map
    }
}