multiversx-sc 0.66.0

use unwrap_infallible::UnwrapInfallible;

use crate::{
    api::{
        ErrorApi, ManagedBufferApiImpl, ManagedTypeApi, ManagedTypeApiImpl, StorageReadApi,
        StorageReadApiImpl, StorageWriteApi, StorageWriteApiImpl, const_handles, use_raw_handle,
    },
    codec::*,
    contract_base::ExitCodecErrorHandler,
    err_msg,
    types::{BigInt, BigUint, ManagedBuffer, ManagedBufferBuilder, ManagedRef, ManagedType},
};

use super::StorageKey;

struct StorageSetOutput<'k, A>
where
    A: StorageWriteApi + ManagedTypeApi + ErrorApi + 'static,
{
    key: ManagedRef<'k, A, StorageKey<A>>,
}

impl<'k, A> StorageSetOutput<'k, A>
where
    A: StorageWriteApi + ManagedTypeApi + ErrorApi + 'static,
{
    #[inline]
    fn new(key: ManagedRef<'k, A, StorageKey<A>>) -> Self {
        StorageSetOutput { key }
    }

    fn set_managed_buffer(&self, managed_buffer: &ManagedBuffer<A>) {
        A::storage_write_api_impl().storage_store_managed_buffer_raw(
            self.key.buffer.get_handle(),
            managed_buffer.handle.clone(),
        );
    }
}

impl<A> TopEncodeOutput for StorageSetOutput<'_, A>
where
    A: StorageWriteApi + ManagedTypeApi + ErrorApi + 'static,
{
    type NestedBuffer = ManagedBufferBuilder<A>;

    fn set_slice_u8(self, bytes: &[u8]) {
        self.set_managed_buffer(&bytes.into())
    }

    fn set_u64(self, value: u64) {
        let handle: A::ManagedBufferHandle = use_raw_handle(const_handles::MBUF_TEMPORARY_1);
        A::managed_type_impl().mb_from_small_int_unsigned(handle.clone(), value as i64);
        let managed_buffer = unsafe { ManagedRef::wrap_handle(handle) };
        self.set_managed_buffer(&managed_buffer);
    }

    /// The version that uses the buggy VM hook. Activate this flag to test the VM.
    #[cfg(feature = "small-int-bug")]
    fn set_i64(self, value: i64) {
        let handle: A::ManagedBufferHandle = use_raw_handle(const_handles::MBUF_TEMPORARY_1);
        A::managed_type_impl().mb_from_small_int_signed(handle.clone(), value);
        let managed_buffer = unsafe { ManagedRef::wrap_handle(handle) };
        self.set_managed_buffer(&managed_buffer);
    }

    /// Workaround for `mb_from_small_int_signed` bug: converts via BigInt first.
    #[cfg(not(feature = "small-int-bug"))]
    fn set_i64(self, value: i64) {
        use crate::api::BigIntApiImpl;
        let api = A::managed_type_impl();
        let bi_handle: A::BigIntHandle = use_raw_handle(const_handles::BIG_INT_TEMPORARY_1);
        api.bi_set_int64(bi_handle.clone(), value);
        let mb_handle: A::ManagedBufferHandle = use_raw_handle(const_handles::MBUF_TEMPORARY_1);
        api.mb_from_big_int_signed(bi_handle, mb_handle.clone());
        let managed_buffer = unsafe { ManagedBuffer::from_handle(mb_handle) };
        self.set_managed_buffer(&managed_buffer);
    }

    #[inline]
    fn supports_specialized_type<T: TryStaticCast>() -> bool {
        T::type_eq::<ManagedBuffer<A>>() || T::type_eq::<BigUint<A>>() || T::type_eq::<BigInt<A>>()
    }

    #[inline]
    fn set_specialized<T, H>(self, value: &T, h: H) -> Result<(), H::HandledErr>
    where
        T: TryStaticCast,
        H: EncodeErrorHandler,
    {
        if let Some(managed_buffer) = value.try_cast_ref::<ManagedBuffer<A>>() {
            self.set_managed_buffer(managed_buffer);
            Ok(())
        } else if let Some(big_uint) = value.try_cast_ref::<BigUint<A>>() {
            self.set_managed_buffer(&big_uint.to_bytes_be_buffer());
            Ok(())
        } else if let Some(big_int) = value.try_cast_ref::<BigInt<A>>() {
            self.set_managed_buffer(&big_int.to_signed_bytes_be_buffer());
            Ok(())
        } else {
            Err(h.handle_error(EncodeError::UNSUPPORTED_OPERATION))
        }
    }

    fn start_nested_encode(&self) -> Self::NestedBuffer {
        ManagedBufferBuilder::new_from_slice(&[])
    }

    fn finalize_nested_encode(self, nb: Self::NestedBuffer) {
        self.set_managed_buffer(&nb.into_managed_buffer());
    }
}

// #[inline]
pub fn storage_set<A, T>(key: ManagedRef<'_, A, StorageKey<A>>, value: &T)
where
    T: TopEncode,
    A: StorageWriteApi + ManagedTypeApi + ErrorApi,
{
    value
        .top_encode_or_handle_err(
            StorageSetOutput::new(key),
            ExitCodecErrorHandler::<A>::from(err_msg::STORAGE_ENCODE_ERROR),
        )
        .unwrap_infallible()
}

/// Useful for storage mappers.
/// Also calls to it generated by macro.
pub fn storage_clear<A>(key: ManagedRef<'_, A, StorageKey<A>>)
where
    A: StorageWriteApi + ManagedTypeApi + ErrorApi,
{
    let value_handle: A::ManagedBufferHandle = use_raw_handle(const_handles::MBUF_CONST_EMPTY);
    A::managed_type_impl().mb_overwrite(value_handle.clone(), &[]);

    A::storage_write_api_impl().storage_store_managed_buffer_raw(key.get_handle(), value_handle);
}

/// Useful for storage mappers.
/// Replaces the content from a key to another without decoding.
pub fn storage_overwrite<A>(
    from_key: ManagedRef<'_, A, StorageKey<A>>,
    to_key: ManagedRef<'_, A, StorageKey<A>>,
) where
    A: StorageWriteApi + StorageReadApi + ManagedTypeApi + ErrorApi,
{
    let from_buffer_handle: A::ManagedBufferHandle =
        use_raw_handle(const_handles::MBUF_TEMPORARY_1);

    A::storage_read_api_impl()
        .storage_load_managed_buffer_raw(from_key.get_handle(), from_buffer_handle.clone());
    A::storage_write_api_impl()
        .storage_store_managed_buffer_raw(to_key.get_handle(), from_buffer_handle);
}