cairo-native 0.9.0-rc.2

//! # `Felt` dictionary entry libfuncs

use super::LibfuncHelper;
use crate::{
    error::{Error, Result},
    metadata::{runtime_bindings::RuntimeBindingsMeta, MetadataStorage},
    types::TypeBuilder,
    utils::ProgramRegistryExt,
};
use cairo_lang_sierra::{
    extensions::{
        core::{CoreLibfunc, CoreType},
        felt252_dict::Felt252DictEntryConcreteLibfunc,
        lib_func::SignatureAndTypeConcreteLibfunc,
        ConcreteLibfunc,
    },
    program_registry::ProgramRegistry,
};
use melior::{
    dialect::{llvm, scf},
    helpers::{ArithBlockExt, BuiltinBlockExt, LlvmBlockExt},
    ir::{r#type::IntegerType, Block, BlockLike, Location, Region},
    Context,
};
use std::cell::Cell;

/// Select and call the correct libfunc builder function from the selector.
pub fn build<'ctx, 'this>(
    context: &'ctx Context,
    registry: &ProgramRegistry<CoreType, CoreLibfunc>,
    entry: &'this Block<'ctx>,
    location: Location<'ctx>,
    helper: &LibfuncHelper<'ctx, 'this>,
    metadata: &mut MetadataStorage,
    selector: &Felt252DictEntryConcreteLibfunc,
) -> Result<()> {
    match selector {
        Felt252DictEntryConcreteLibfunc::Get(info) => {
            build_get(context, registry, entry, location, helper, metadata, info)
        }
        Felt252DictEntryConcreteLibfunc::Finalize(info) => {
            build_finalize(context, registry, entry, location, helper, metadata, info)
        }
    }
}

/// The felt252_dict_entry_get libfunc receives the dictionary and the key and
/// returns the associated dict entry, along with it's value.
///
/// The dict entry also contains a pointer to the dictionary.
///
/// If the key doesn't yet exist, it is created and the type's default value is returned.
///
/// # Cairo Signature
///
/// ```cairo
/// fn felt252_dict_entry_get<T>(dict: Felt252Dict<T>, key: felt252) -> (Felt252DictEntry<T>, T) nopanic;
/// ```
pub fn build_get<'ctx, 'this>(
    context: &'ctx Context,
    registry: &ProgramRegistry<CoreType, CoreLibfunc>,
    entry: &'this Block<'ctx>,
    location: Location<'ctx>,
    helper: &LibfuncHelper<'ctx, 'this>,
    metadata: &mut MetadataStorage,
    info: &SignatureAndTypeConcreteLibfunc,
) -> Result<()> {
    let (key_ty, key_layout) = registry.build_type_with_layout(
        context,
        helper,
        metadata,
        &info.param_signatures()[1].ty,
    )?;
    let entry_ty = registry.build_type(
        context,
        helper,
        metadata,
        &info.branch_signatures()[0].vars[0].ty,
    )?;
    let concrete_value_type = registry.get_type(&info.ty)?;
    let value_ty = concrete_value_type.build(context, helper, registry, metadata, &info.ty)?;

    let dict_ptr = entry.arg(0)?;
    let entry_key = entry.arg(1)?;

    let entry_key_ptr =
        helper
            .init_block()
            .alloca1(context, location, key_ty, key_layout.align())?;
    entry.store(context, location, entry_key_ptr, entry_key)?;

    let (is_present, value_ptr) = metadata
        .get_mut::<RuntimeBindingsMeta>()
        .ok_or(Error::MissingMetadata)?
        .dict_get(context, helper, entry, dict_ptr, entry_key_ptr, location)?;
    let is_present = entry.trunci(is_present, IntegerType::new(context, 1).into(), location)?;

    let value = entry.append_op_result(scf::r#if(
        is_present,
        &[value_ty],
        {
            let region = Region::new();
            let block = region.append_block(Block::new(&[]));

            // If the entry is present we can load the current value.
            let value = block.load(context, location, value_ptr, value_ty)?;

            block.append_operation(scf::r#yield(&[value], location));
            region
        },
        {
            let region = Region::new();
            let block = region.append_block(Block::new(&[]));

            let helper = LibfuncHelper {
                module: helper.module,
                init_block: helper.init_block,
                region: &region,
                blocks_arena: helper.blocks_arena,
                last_block: Cell::new(&block),
                branches: Vec::new(),
                results: Vec::new(),

                #[cfg(feature = "with-libfunc-profiling")]
                profiler: helper.profiler.clone(),
            };

            // When the entry is vacant we need to create the default value.
            let value = concrete_value_type.build_default(
                context, registry, &block, location, &helper, metadata, &info.ty,
            )?;

            block.append_operation(scf::r#yield(&[value], location));
            region
        },
        location,
    ))?;

    let dict_entry = entry.append_op_result(llvm::undef(entry_ty, location))?;
    let dict_entry = entry.insert_values(context, location, dict_entry, &[dict_ptr, value_ptr])?;

    // The `Felt252DictEntry<T>` holds both the `Felt252Dict<T>` and the pointer to the space where
    // the new value will be written when the entry is finalized. If the entry were to be dropped
    // (without being consumed by the finalizer), which shouldn't be possible under normal
    // conditions, and the type `T` requires a custom drop implementation (ex. arrays, dicts...),
    // it'll cause undefined behavior because when the value is moved out of the dictionary (on
    // `get`), the memory it occupied is not modified because we're expecting it to be overwritten
    // by the finalizer (in other words, the extracted element will be dropped twice).

    helper.br(entry, 0, &[dict_entry, value], location)
}

/// The felt252_dict_entry_finalize libfunc receives the dict entry and a new value,
/// inserts the new value in the entry, and returns the full dictionary.
///
/// # Cairo Signature
///
/// ```cairo
/// fn felt252_dict_entry_finalize<T>(dict_entry: Felt252DictEntry<T>, new_value: T) -> Felt252Dict<T> nopanic;
/// ```
pub fn build_finalize<'ctx, 'this>(
    context: &'ctx Context,
    _registry: &ProgramRegistry<CoreType, CoreLibfunc>,
    entry: &'this Block<'ctx>,
    location: Location<'ctx>,
    helper: &LibfuncHelper<'ctx, 'this>,
    _metadata: &mut MetadataStorage,
    _info: &SignatureAndTypeConcreteLibfunc,
) -> Result<()> {
    // Get the dict entry struct: `crate::types::felt252_dict_entry`.
    let dict_entry = entry.arg(0)?;
    let new_value = entry.arg(1)?;

    let dict_ptr = entry.extract_value(
        context,
        location,
        dict_entry,
        llvm::r#type::pointer(context, 0),
        0,
    )?;
    let value_ptr = entry.extract_value(
        context,
        location,
        dict_entry,
        llvm::r#type::pointer(context, 0),
        1,
    )?;

    entry.store(context, location, value_ptr, new_value)?;

    helper.br(entry, 0, &[dict_ptr], location)
}

#[cfg(test)]
mod test {
    use crate::{
        jit_dict,
        utils::testing::{get_compiled_program, run_program, run_program_assert_output},
    };

    #[test]
    fn run_dict_insert() {
        let program =
            get_compiled_program("test_data_artifacts/programs/libfuncs/felt252_dict_entry_insert");

        run_program_assert_output(&program, "run_test", &[], 1u32.into());
    }

    #[test]
    fn run_dict_insert_big() {
        let program = get_compiled_program(
            "test_data_artifacts/programs/libfuncs/felt252_dict_entry_insert_big",
        );

        run_program_assert_output(&program, "run_test", &[], 4u64.into());
    }

    #[test]
    fn run_dict_insert_ret_dict() {
        let program = get_compiled_program(
            "test_data_artifacts/programs/libfuncs/felt252_dict_entry_insert_ret_dict",
        );

        run_program_assert_output(
            &program,
            "run_test",
            &[],
            jit_dict!(
                2 => 1u32
            ),
        );
    }

    #[test]
    fn run_dict_insert_multiple() {
        let program = get_compiled_program(
            "test_data_artifacts/programs/libfuncs/felt252_dict_entry_insert_multiple",
        );

        run_program_assert_output(&program, "run_test", &[], 1345432_u32.into());
    }

    #[test]
    fn run_dict_clone_ptr_update() {
        let program = get_compiled_program(
            "test_data_artifacts/programs/libfuncs/felt252_dict_entry_clone_ptr_update",
        );

        run_program(&program, "run_test", &[]);
    }
}