cairo-native 0.9.0-rc.9

#![cfg(any(test, feature = "testing"))]

use crate::{
    context::NativeContext,
    error::{panic::ToNativeAssertError, Error, Result, SierraAssertError},
    execution_result::ExecutionResult,
    executor::JitNativeExecutor,
    ffi::{get_data_layout_rep, initialize_llvm_targets},
    metadata::MetadataStorage,
    starknet_stub::StubSyscallHandler,
    types::TypeBuilder,
    utils::*,
    values::Value,
};
use cairo_lang_compiler::{
    compile_prepared_db, db::RootDatabase, diagnostics::DiagnosticsReporter,
    project::setup_project, CompilerConfig,
};
use cairo_lang_filesystem::{db::init_dev_corelib, ids::CrateInput};
use cairo_lang_lowering::utils::InliningStrategy;
use cairo_lang_runner::SierraCasmRunner;
use cairo_lang_sierra::{
    extensions::core::{CoreLibfunc, CoreType, CoreTypeConcrete},
    program::{Program, VersionedProgram},
    program_registry::ProgramRegistry,
    ProgramParser,
};
use cairo_lang_starknet::{compile::compile_contract_in_prepared_db, starknet_plugin_suite};
use cairo_lang_starknet_classes::contract_class::ContractClass;
use itertools::Itertools;
use llvm_sys::{
    core::{
        LLVMContextCreate, LLVMContextDispose, LLVMDisposeModule, LLVMGetNamedGlobal,
        LLVMGlobalGetValueType, LLVMSetDataLayout,
    },
    prelude::LLVMModuleRef,
    target::{LLVMABIAlignmentOfType, LLVMABISizeOfType, LLVMGetModuleDataLayout},
};
use melior::{
    dialect::ods,
    ir::{
        attribute::{StringAttribute, TypeAttribute},
        Attribute, BlockLike, Location, Module, Region,
    },
};
use mlir_sys::mlirTranslateModuleToLLVMIR;
use starknet_types_core::felt::Felt;
use std::{env::var, ffi::CString, fs, path::Path, sync::Arc};

#[macro_export]
macro_rules! load_cairo {
        ( $( $program:tt )+ ) => {
            $crate::utils::testing::load_cairo_str(stringify!($($program)+))
        };
    }
#[macro_export]
macro_rules! load_starknet {
        ( $( $program:tt )+ ) => {
            $crate::utils::testing::load_starknet_str(stringify!($($program)+))
        };
    }
#[macro_export]
macro_rules! load_starknet_contract {
        ( $( $program:tt )+ ) => {
            $crate::utils::testing::load_starknet_contract_str(stringify!($($program)+))
        };
    }

// Helper macros for faster testing.
#[macro_export]
macro_rules! jit_struct {
        ($($y:expr),* $(,)? ) => {
            $crate::values::Value::Struct {
                fields: vec![$($y), *],
                debug_name: None
            }
        };
    }
#[macro_export]
macro_rules! jit_enum {
    ( $tag:expr, $value:expr ) => {
        $crate::values::Value::Enum {
            tag: $tag,
            value: Box::new($value),
            debug_name: None,
        }
    };
}
#[macro_export]
macro_rules! jit_dict {
        ( $($key:expr $(=>)+ $value:expr),* $(,)? ) => {
            $crate::values::Value::Felt252Dict {
                value: {
                    let mut map = std::collections::HashMap::new();
                    $(map.insert($key.into(), $value.into());)*
                    map
                },
                debug_name: None,
            }
        };
    }
#[macro_export]
macro_rules! jit_panic {
        ( $($value:expr)? ) => {
            $crate::jit_enum!(1, $crate::jit_struct!(
                $crate::jit_struct!(),
                [$($value), *].into()
            ))
        };
    }

#[macro_export]
macro_rules! jit_panic_byte_array {
    ( $value:expr ) => {
        $crate::jit_enum!(
            1,
            $crate::jit_struct!(
                $crate::jit_struct!(),
                $crate::utils::testing::panic_byte_array($value).into()
            )
        )
    };
}

/// Compile a cairo program found at the given path to sierra.
pub fn cairo_to_sierra(program: &Path) -> crate::error::Result<Arc<Program>> {
    if program
        .extension()
        .map(|x| {
            x.to_ascii_lowercase()
                .to_string_lossy()
                .eq_ignore_ascii_case("cairo")
        })
        .unwrap_or(false)
    {
        cairo_lang_compiler::compile_cairo_project_at_path(
            program,
            CompilerConfig {
                replace_ids: true,
                ..Default::default()
            },
            InliningStrategy::Default,
        )
        .map_err(|err| crate::error::Error::ProgramParser(err.to_string()))
    } else {
        let source = std::fs::read_to_string(program)?;
        ProgramParser::new()
            .parse(&source)
            .map_err(|err| crate::error::Error::ProgramParser(err.to_string()))
    }
    .map(Arc::new)
}

pub fn load_cairo_str(program_str: &str) -> (String, Program) {
    compile_program(program_str, RootDatabase::default())
}

pub fn load_starknet_str(program_str: &str) -> (String, Program) {
    compile_program(
        program_str,
        RootDatabase::builder()
            .with_default_plugin_suite(starknet_plugin_suite())
            .build()
            .unwrap(),
    )
}

pub fn load_starknet_contract_str(program_str: &str) -> (String, ContractClass) {
    compile_contract(
        program_str,
        RootDatabase::builder()
            .with_default_plugin_suite(starknet_plugin_suite())
            .build()
            .unwrap(),
    )
}

pub(crate) fn compile_contract(program_str: &str, mut db: RootDatabase) -> (String, ContractClass) {
    let mut program_file = tempfile::Builder::new()
        .prefix("test_")
        .suffix(".cairo")
        .tempfile()
        .unwrap();
    fs::write(&mut program_file, program_str).unwrap();

    init_dev_corelib(
        &mut db,
        Path::new(&var("CARGO_MANIFEST_DIR").unwrap()).join("corelib/src"),
    );
    let main_crate_ids = {
        let main_crate_inputs = setup_project(&mut db, program_file.path()).unwrap();
        CrateInput::into_crate_ids(&db, main_crate_inputs)
    };
    let contract = compile_contract_in_prepared_db(
        &db,
        None,
        main_crate_ids,
        CompilerConfig {
            diagnostics_reporter: DiagnosticsReporter::stderr(),
            replace_ids: true,
            ..Default::default()
        },
    )
    .unwrap();

    let module_name = program_file.path().with_extension("");
    let module_name = module_name.file_name().unwrap().to_str().unwrap();
    (module_name.to_string(), contract)
}

pub(crate) fn compile_program(program_str: &str, mut db: RootDatabase) -> (String, Program) {
    let mut program_file = tempfile::Builder::new()
        .prefix("test_")
        .suffix(".cairo")
        .tempfile()
        .unwrap();
    fs::write(&mut program_file, program_str).unwrap();

    init_dev_corelib(
        &mut db,
        Path::new(&var("CARGO_MANIFEST_DIR").unwrap()).join("corelib/src"),
    );
    let main_crate_ids = {
        let main_crate_inputs = setup_project(&mut db, program_file.path()).unwrap();
        CrateInput::into_crate_ids(&db, main_crate_inputs)
    };
    let sierra_program_with_dbg = compile_prepared_db(
        &db,
        main_crate_ids,
        CompilerConfig {
            diagnostics_reporter: DiagnosticsReporter::stderr(),
            replace_ids: true,
            ..Default::default()
        },
    )
    .unwrap();

    let module_name = program_file.path().with_extension("");
    let module_name = module_name.file_name().unwrap().to_str().unwrap();
    (module_name.to_string(), sierra_program_with_dbg.program)
}

pub fn run_program(
    program: &(String, Program),
    entry_point: &str,
    args: &[Value],
) -> ExecutionResult {
    let entry_point = format!("{0}::{0}::{1}", program.0, entry_point);
    let program = &program.1;

    let entry_point_id = &program
        .funcs
        .iter()
        .find(|x| x.id.debug_name.as_deref() == Some(&entry_point))
        .expect("Test program entry point not found.")
        .id;

    let context = NativeContext::new();

    let module = context
        .compile(program, false, Some(Default::default()), None)
        .expect("Could not compile test program to MLIR.");

    let executor = JitNativeExecutor::from_native_module(module, OptLevel::Less).unwrap();
    executor
        .invoke_dynamic_with_syscall_handler(
            entry_point_id,
            args,
            Some(u64::MAX),
            &mut StubSyscallHandler::default(),
        )
        .unwrap()
}

#[track_caller]
pub fn run_program_assert_output(
    program: &(String, Program),
    entry_point: &str,
    args: &[Value],
    output: Value,
) {
    let result = run_program(program, entry_point, args);
    assert_eq!(result.return_value, output);
}

/// Serializes a message into a vector of felts, the same way that Cairo
/// serializes byte arrays. Used for asserting panic message on tests.
///
/// https://github.com/starkware-libs/cairo/tree/v2.12.3/corelib/src/debug.cairo#L142
pub fn panic_byte_array(message: &str) -> Vec<Felt> {
    // Prepend byte array magic, used to identify serialized `ByteArray` variables.
    // https://github.com/starkware-libs/cairo/tree/v2.12.3/corelib/src/byte_array.cairo#L64
    let mut array = vec![Felt::from_hex_unchecked(
        "0x46a6158a16a947e5916b2a2ca68501a45e93d7110e81aa2d6438b1c57c879a3",
    )];

    let chunk_iter = message.bytes().chunks(31);
    let mut chunks = chunk_iter.into_iter().collect_vec();

    // Take last word as its serialized differently.
    let pending = chunks
        .pop()
        .map(|pendign| pendign.collect_vec())
        .unwrap_or_default();

    // Serialize length of the byte array.
    array.push(chunks.len().into());

    // Serialize each byte array element.
    for chunk in chunks {
        let chunk = chunk.collect_vec();
        array.push(Felt::from_bytes_be_slice(&chunk));
    }

    // Serialize last word with its length.
    array.extend_from_slice(&[Felt::from_bytes_be_slice(&pending), pending.len().into()]);

    array
}

// TODO: Eventually this function will be changed or will be replaced in a future refactor. Right now
// so as to not make a big unrelated change we keep it like this.
pub fn load_program_and_runner(name: &str) -> (String, Program, SierraCasmRunner) {
    let program = load_program(name);
    let runner = SierraCasmRunner::new(
        program.clone(),
        Some(Default::default()),
        Default::default(),
        None,
    )
    .unwrap();
    let entrypoint = name
        .split("/")
        .collect::<Vec<&str>>()
        .last()
        .unwrap()
        .to_string();
    (entrypoint, program, runner)
}

// TODO: Same as load_program_and_runner(). Both of the functions will be changed in a rafactor.
pub fn get_compiled_program(name: &str) -> (String, Program) {
    let program_path = format!(
        "{}/test_data_artifacts/{name}.sierra.json",
        env!("CARGO_MANIFEST_DIR"),
    );
    let program_content = fs::read_to_string(program_path)
        .expect("Failed to read the content of the program into a String");
    let versioned_program =
        serde_json::from_str::<cairo_lang_sierra::program::VersionedProgram>(&program_content)
            .unwrap();
    let program = versioned_program.into_v1().unwrap().program;
    let entrypoint = name
        .split("/")
        .collect::<Vec<&str>>()
        .last()
        .unwrap()
        .to_string();
    (entrypoint, program)
}

pub fn load_contract(path: &str) -> ContractClass {
    serde_json::from_str::<ContractClass>(
        &fs::read_to_string(format!(
            "{}/test_data_artifacts/{path}",
            env!("CARGO_MANIFEST_DIR")
        ))
        .unwrap(),
    )
    .unwrap()
}

pub fn load_program(path: &str) -> Program {
    let versioned_program = serde_json::from_str::<VersionedProgram>(
        &fs::read_to_string(format!(
            "{}/test_data_artifacts/{}.sierra.json",
            env!("CARGO_MANIFEST_DIR"),
            path
        ))
        .unwrap(),
    )
    .unwrap();
    versioned_program.into_v1().unwrap().program
}

/// Checks that, for every type declared in `program`, the Rust-side
/// [`TypeBuilder::layout`] agrees with the size and alignment LLVM assigns to
/// the MLIR type emitted by [`TypeBuilder::build`]. Returns one message per
/// mismatching type.
///
/// Every declared type is built and emitted as an external global in a scratch
/// module, which is then lowered and translated to LLVM IR so the real target
/// data layout answers the question.
pub fn find_type_layout_mismatches(
    native_context: &NativeContext,
    program: &Program,
) -> Result<Vec<String>> {
    const GLOBAL_PREFIX: &str = "cairo_native__layout_check__";

    initialize_llvm_targets();

    let context = native_context.context();
    let registry = ProgramRegistry::<CoreType, CoreLibfunc>::new(program)?;
    let location = Location::unknown(context);
    let mut module = Module::new(location);
    let mut metadata = MetadataStorage::new();

    let mut expected = Vec::with_capacity(program.type_declarations.len());
    for declaration in &program.type_declarations {
        let concrete_type = registry.get_type(&declaration.id)?;
        if matches!(concrete_type, CoreTypeConcrete::Const(_)) {
            // Const types only exist at compile time and have no MLIR type.
            continue;
        }

        let ty = match registry.build_type(context, &module, &mut metadata, &declaration.id) {
            Ok(ty) => ty,
            // Circuit gates and anything that contains them are compile-time only.
            Err(Error::SierraAssert(SierraAssertError::BadTypeInit(_))) => continue,
            Err(e) => return Err(e),
        };
        let layout = concrete_type.layout(&registry)?.pad_to_align();

        module.body().append_operation(
            ods::llvm::mlir_global(
                context,
                Region::new(),
                TypeAttribute::new(ty),
                StringAttribute::new(context, &format!("{GLOBAL_PREFIX}{}", expected.len())),
                Attribute::parse(context, "#llvm.linkage<external>")
                    .to_native_assert_error("linkage attribute should parse")?,
                location,
            )
            .into(),
        );
        expected.push((&declaration.long_id, layout));
    }

    run_pass_manager(context, &mut module)?;

    let data_layout_rep = CString::new(get_data_layout_rep()?)
        .to_native_assert_error("data layout string should not contain nul bytes")?;

    let mut mismatches = Vec::new();
    unsafe {
        let llvm_context = LLVMContextCreate();
        let llvm_module: LLVMModuleRef =
            mlirTranslateModuleToLLVMIR(module.as_operation().to_raw(), llvm_context as *mut _)
                as *mut _;
        LLVMSetDataLayout(llvm_module, data_layout_rep.as_ptr());
        let data_layout = LLVMGetModuleDataLayout(llvm_module);

        for (index, (type_id, layout)) in expected.iter().enumerate() {
            let name = CString::new(format!("{GLOBAL_PREFIX}{index}"))
                .to_native_assert_error("global name should not contain nul bytes")?;
            let global = LLVMGetNamedGlobal(llvm_module, name.as_ptr());
            assert!(
                !global.is_null(),
                "global for {type_id} was dropped during lowering"
            );
            let ty = LLVMGlobalGetValueType(global);
            let llvm_size = LLVMABISizeOfType(data_layout, ty) as usize;
            let llvm_align = LLVMABIAlignmentOfType(data_layout, ty) as usize;
            if llvm_size != layout.size() || llvm_align != layout.align() {
                mismatches.push(format!(
                    "{type_id}: layout() reports size {} align {}, but LLVM reports size \
                     {llvm_size} align {llvm_align}",
                    layout.size(),
                    layout.align(),
                ));
            }
        }

        LLVMDisposeModule(llvm_module);
        LLVMContextDispose(llvm_context);
    }

    Ok(mismatches)
}