cairo-lang-semantic 2.21.0

use std::sync::{LazyLock, Mutex};

use cairo_lang_defs::db::{DefsGroup, init_defs_group, init_external_files};
use cairo_lang_defs::ids::{FreeFunctionId, FunctionWithBodyId, ModuleId};
use cairo_lang_defs::plugin::{MacroPlugin, MacroPluginMetadata, PluginResult};
use cairo_lang_diagnostics::{Diagnostics, DiagnosticsBuilder};
use cairo_lang_filesystem::db::{
    CrateSettings, Edition, ExperimentalFeaturesConfig, init_dev_corelib, init_files_group,
};
use cairo_lang_filesystem::detect::detect_corelib;
use cairo_lang_filesystem::ids::{
    BlobId, CrateId, CrateLongId, FileKind, FileLongId, SmolStrId, VirtualFile,
};
use cairo_lang_parser::db::ParserGroup;
use cairo_lang_syntax::node::ast;
use cairo_lang_syntax::node::helpers::QueryAttrs;
use cairo_lang_test_utils::parse_test_file::TestRunnerResult;
use cairo_lang_test_utils::{get_direct_or_file_content, verify_diagnostics_expectation};
use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
use cairo_lang_utils::{Intern, extract_matches};
use salsa::Database;

use crate::db::{PluginSuiteInput, SemanticGroup, init_semantic_group};
use crate::inline_macros::get_default_plugin_suite;
use crate::items::function_with_body::FunctionWithBodySemantic;
use crate::plugin::PluginSuite;
use crate::{ConcreteFunctionWithBodyId, SemanticDiagnostic, semantic};

/// The attribute marking the function(s) a test block targets, within its `cairo_code` input.
pub const TARGET_FUNCTION_ATTR: &str = "target_function";

/// A test-only no-op macro plugin, whose sole purpose is to declare the `#[target_function]`
/// attribute, so that marking a function with it in test code is not reported as an
/// "Unsupported attribute." error.
///
/// It never generates code and never removes the original item, so it has no effect on any other
/// plugin in the suite.
#[derive(Debug, Default)]
pub struct TargetFunctionPlugin;

impl MacroPlugin for TargetFunctionPlugin {
    fn generate_code<'db>(
        &self,
        _db: &'db dyn Database,
        _item_ast: ast::ModuleItem<'db>,
        _metadata: &MacroPluginMetadata<'_>,
    ) -> PluginResult<'db> {
        PluginResult::default()
    }

    fn declared_attributes<'db>(&self, db: &'db dyn Database) -> Vec<SmolStrId<'db>> {
        vec![SmolStrId::from(db, TARGET_FUNCTION_ATTR)]
    }
}

/// Adds [TargetFunctionPlugin] to the given suite, and returns it.
///
/// Every testing database that runs golden test files must use this, otherwise `#[target_function]`
/// markers in the test code would produce "Unsupported attribute." diagnostics.
pub fn with_target_function_plugin(mut suite: PluginSuite) -> PluginSuite {
    suite.add_plugin::<TargetFunctionPlugin>();
    suite
}

/// Returns the free functions of the given module marked with `#[target_function]`, in the order
/// they appear in the module.
pub fn resolve_target_functions<'a>(
    db: &'a dyn Database,
    module_id: ModuleId<'a>,
) -> Vec<FreeFunctionId<'a>> {
    db.module_free_functions_ids(module_id)
        .expect("Failed to load module.")
        .iter()
        .copied()
        .filter(|free_function_id| {
            db.module_free_function_by_id(*free_function_id)
                .expect("Failed to load free function.")
                .has_attr(db, TARGET_FUNCTION_ATTR)
        })
        .collect()
}

#[salsa::db]
#[derive(Clone)]
pub struct SemanticDatabaseForTesting {
    storage: salsa::Storage<SemanticDatabaseForTesting>,
}

#[salsa::db]
impl Database for SemanticDatabaseForTesting {}

impl SemanticDatabaseForTesting {
    pub fn new_empty() -> Self {
        let suite = get_default_plugin_suite();
        SemanticDatabaseForTesting::with_plugin_suite(suite)
    }

    pub fn with_plugin_suite(suite: PluginSuite) -> Self {
        let mut res = SemanticDatabaseForTesting { storage: Default::default() };
        init_external_files(&mut res);
        init_files_group(&mut res);
        init_defs_group(&mut res);
        init_semantic_group(&mut res);

        res.set_default_plugins_from_suite(with_target_function_plugin(suite));

        let corelib_path = detect_corelib().expect("Corelib not found in default location.");
        init_dev_corelib(&mut res, corelib_path);
        res
    }
}
pub static SHARED_DB: LazyLock<Mutex<SemanticDatabaseForTesting>> =
    LazyLock::new(|| Mutex::new(SemanticDatabaseForTesting::new_empty()));
impl Default for SemanticDatabaseForTesting {
    fn default() -> Self {
        SHARED_DB.lock().unwrap().clone()
    }
}

pub struct WithStringDiagnostics<T> {
    value: T,
    diagnostics: String,
}
impl<T> WithStringDiagnostics<T> {
    /// Verifies that there are no diagnostics (fails otherwise), and returns the inner value.
    pub fn unwrap(self) -> T {
        assert!(self.diagnostics.is_empty(), "Unexpected diagnostics:\n{}", self.diagnostics);
        self.value
    }

    /// Returns the inner value and the diagnostics (as a string).
    pub fn split(self) -> (T, String) {
        (self.value, self.diagnostics)
    }

    /// Returns the diagnostics (as a string).
    pub fn get_diagnostics(self) -> String {
        self.diagnostics
    }
}

/// Helper struct for the return value of [setup_test_module].
pub struct TestModule<'a> {
    pub crate_id: CrateId<'a>,
    pub module_id: ModuleId<'a>,
}

/// Sets up a crate with given content, and returns its crate id.
pub fn setup_test_crate_ex<'a>(
    db: &'a dyn Database,
    content: &str,
    crate_settings: Option<&str>,
    cache_file: Option<BlobId<'a>>,
) -> CrateId<'a> {
    let file_long_id = FileLongId::Virtual(VirtualFile {
        parent: None,
        name: SmolStrId::from(db, "lib.cairo"),
        content: SmolStrId::from(db, content),
        code_mappings: [].into(),
        kind: FileKind::Module,
        original_item_removed: false,
    });

    let settings: CrateSettings = if let Some(crate_settings) = crate_settings {
        toml::from_str(crate_settings).expect("Invalid config.")
    } else {
        CrateSettings {
            name: None,
            edition: Edition::default(),
            version: None,
            dependencies: Default::default(),
            experimental_features: ExperimentalFeaturesConfig {
                negative_impls: true,
                associated_item_constraints: true,
                coupons: true,
                user_defined_inline_macros: true,
                repr_ptrs: true,
            },
            cfg_set: Default::default(),
        }
    };

    CrateLongId::Virtual {
        name: SmolStrId::from(db, "test"),
        file_id: file_long_id.intern(db),
        settings: toml::to_string_pretty(&settings).unwrap(),
        cache_file,
    }
    .intern(db)
}

/// See [setup_test_crate_ex].
pub fn setup_test_crate<'a>(db: &'a dyn Database, content: &str) -> CrateId<'a> {
    setup_test_crate_ex(db, content, None, None)
}

/// Sets up a module with given content, and returns its module id.
pub fn setup_test_module_ex<'a>(
    db: &'a dyn Database,
    content: &str,
    crate_settings: Option<&str>,
    cached_crate: Option<BlobId<'a>>,
) -> WithStringDiagnostics<TestModule<'a>> {
    let crate_id = setup_test_crate_ex(db, content, crate_settings, cached_crate);
    let module_id = ModuleId::CrateRoot(crate_id);
    let file_id = db.module_main_file(module_id).unwrap();

    let syntax_diagnostics = db.file_syntax_diagnostics(file_id).format(db);
    let semantic_diagnostics = get_recursive_module_semantic_diagnostics(db, module_id).format(db);

    WithStringDiagnostics {
        value: TestModule { crate_id, module_id },
        diagnostics: format!("{syntax_diagnostics}{semantic_diagnostics}"),
    }
}

/// See [setup_test_module_ex].
pub fn setup_test_module<'a>(
    db: &'a dyn Database,
    content: &str,
) -> WithStringDiagnostics<TestModule<'a>> {
    setup_test_module_ex(db, content, None, None)
}

/// Helper struct for the return value of [setup_test_function].
pub struct TestFunction<'a> {
    pub module_id: ModuleId<'a>,
    pub function_id: FunctionWithBodyId<'a>,
    pub concrete_function_id: ConcreteFunctionWithBodyId<'a>,
    pub signature: semantic::Signature<'a>,
    pub body: semantic::ExprId,
}

/// Returns the code of the module described by the given test inputs.
///
/// `inputs["cairo_code"]` is the whole module - possibly a `>>> file: <path>` reference.
pub fn test_module_code(inputs: &OrderedHashMap<String, String>) -> String {
    let (_path, content) = get_direct_or_file_content(&inputs["cairo_code"]);
    content
}

/// Returns the semantic model of the function a test targets.
///
/// The module code is taken as described in [test_module_code], and the target function within it
/// is the single function marked with `#[target_function]`.
///
/// `inputs["crate_settings"]` - optional extra settings for the crate.
pub fn setup_test_function<'a>(
    db: &'a dyn Database,
    inputs: &OrderedHashMap<String, String>,
) -> WithStringDiagnostics<TestFunction<'a>> {
    setup_test_function_from_content(
        db,
        &test_module_code(inputs),
        inputs.get("crate_settings").map(String::as_str),
        None,
    )
}

/// Returns the semantic model of a function within the module of the given `content`.
///
/// The function is the single one marked with the `#[target_function]` marker.
pub fn setup_test_function_from_content<'a>(
    db: &'a dyn Database,
    content: &str,
    crate_settings: Option<&str>,
    cache_crate: Option<BlobId<'a>>,
) -> WithStringDiagnostics<TestFunction<'a>> {
    let (test_module, diagnostics) =
        setup_test_module_ex(db, content, crate_settings, cache_crate).split();
    let target_functions = resolve_target_functions(db, test_module.module_id);
    let free_function_id = match target_functions.as_slice() {
        [free_function_id] => *free_function_id,
        [] => panic!(
            "No function marked with `#[{TARGET_FUNCTION_ATTR}]` was found. Mark the tested \
             function."
        ),
        _ => panic!(
            "Expected a single function marked with `#[{TARGET_FUNCTION_ATTR}]`, found {}.",
            target_functions.len()
        ),
    };
    let function_id = FunctionWithBodyId::Free(free_function_id);
    WithStringDiagnostics {
        value: TestFunction {
            module_id: test_module.module_id,
            function_id,
            concrete_function_id: ConcreteFunctionWithBodyId::from_no_generics_free(
                db,
                free_function_id,
            )
            .unwrap(),
            signature: db.function_with_body_signature(function_id).unwrap().clone(),
            body: db.function_body_expr(function_id).unwrap(),
        },
        diagnostics,
    }
}

/// Helper struct for the return value of [setup_test_expr].
pub struct TestExpr<'a> {
    pub module_id: ModuleId<'a>,
    pub function_id: FunctionWithBodyId<'a>,
    pub signature: semantic::Signature<'a>,
    pub body: semantic::ExprId,
    pub expr_id: semantic::ExprId,
}

/// Returns the semantic model of a given expression.
/// module_code - extra setup code in the module context.
/// function_body - extra setup code in the function context.
pub fn setup_test_expr<'a>(
    db: &'a dyn Database,
    expr_code: &str,
    module_code: &str,
    function_body: &str,
    crate_settings: Option<&str>,
) -> WithStringDiagnostics<TestExpr<'a>> {
    let ((test_function, statements, tail), diagnostics) =
        setup_test_code_block(db, expr_code, module_code, function_body, crate_settings, false)
            .split();
    assert!(statements.is_empty(), "expr_code is not a valid expression.");
    WithStringDiagnostics {
        value: TestExpr {
            module_id: test_function.module_id,
            function_id: test_function.function_id,
            signature: test_function.signature,
            body: test_function.body,
            expr_id: tail.unwrap(),
        },
        diagnostics,
    }
}

/// Helper struct for the return value of [setup_test_exprs].
pub struct TestExprs<'a> {
    pub function_id: FunctionWithBodyId<'a>,
    /// The tested expressions, each with the name it was bound to.
    pub named_exprs: Vec<(SmolStrId<'a>, semantic::ExprId)>,
}

/// Returns the semantic model of the given expressions, each provided as a
/// `let <name> = <expr>;` statement, paired with its name.
/// module_code - extra setup code in the module context.
/// function_body - extra setup code in the function context.
pub fn setup_test_exprs<'a>(
    db: &'a dyn Database,
    exprs_code: &str,
    module_code: &str,
    function_body: &str,
    crate_settings: Option<&str>,
) -> WithStringDiagnostics<TestExprs<'a>> {
    let ((test_function, statements, tail), diagnostics) =
        setup_test_code_block(db, exprs_code, module_code, function_body, crate_settings, true)
            .split();
    const EXPECTATION: &str = "exprs_code must consist of `let <name> = <expr>;` statements.";
    assert!(tail.is_none(), "{EXPECTATION}");
    let function_id = test_function.function_id;
    let named_exprs = statements
        .iter()
        .map(|statement| {
            let statement_let = extract_matches!(
                db.statement_semantic(function_id, *statement),
                semantic::Statement::Let,
                "{EXPECTATION}"
            );
            let pattern = extract_matches!(
                db.pattern_semantic(function_id, statement_let.pattern),
                semantic::Pattern::Variable,
                "{EXPECTATION}"
            );
            (pattern.name, statement_let.expr)
        })
        .collect();
    WithStringDiagnostics { value: TestExprs { function_id, named_exprs }, diagnostics }
}

/// Returns the test function created by wrapping the given `code` in a block within a test
/// function, together with the statements and tail of that block.
fn setup_test_code_block<'a>(
    db: &'a dyn Database,
    code: &str,
    module_code: &str,
    function_body: &str,
    crate_settings: Option<&str>,
    allow_unused_variables: bool,
) -> WithStringDiagnostics<(TestFunction<'a>, Vec<semantic::StatementId>, Option<semantic::ExprId>)>
{
    let allow_attr = if allow_unused_variables { "#[allow(unused_variables)] " } else { "" };
    // The attributes share the wrapper's first line so the spans of `code` do not shift.
    let function_code = format!(
        "{allow_attr}#[{TARGET_FUNCTION_ATTR}] fn test_func() {{ {function_body} {{\n{code}\n}}; \
         }}"
    );
    let content = if module_code.is_empty() {
        function_code
    } else {
        format!("{module_code}\n{function_code}")
    };
    let (test_function, diagnostics) =
        setup_test_function_from_content(db, &content, crate_settings, None).split();
    let semantic::ExprBlock { statements, .. } = extract_matches!(
        db.expr_semantic(test_function.function_id, test_function.body),
        semantic::Expr::Block
    );
    let statement_expr = extract_matches!(
        db.statement_semantic(test_function.function_id, *statements.last().unwrap()),
        semantic::Statement::Expr
    );
    let semantic::ExprBlock { statements, tail, .. } = extract_matches!(
        db.expr_semantic(test_function.function_id, statement_expr.expr),
        semantic::Expr::Block
    );
    WithStringDiagnostics { value: (test_function, statements, tail), diagnostics }
}

pub fn test_expr_diagnostics(
    inputs: &OrderedHashMap<String, String>,
    args: &OrderedHashMap<String, String>,
) -> TestRunnerResult {
    let db = &SemanticDatabaseForTesting::default();

    let diagnostics = setup_test_expr(
        db,
        &inputs["expr_code"],
        inputs.get("module_code").map_or("", String::as_str),
        &inputs["function_body"],
        inputs.get("crate_settings").map(String::as_str),
    )
    .get_diagnostics();
    let error = verify_diagnostics_expectation(args, &diagnostics);

    TestRunnerResult {
        outputs: OrderedHashMap::from([("expected_diagnostics".into(), diagnostics)]),
        error,
    }
}

pub fn test_function_diagnostics(
    inputs: &OrderedHashMap<String, String>,
    args: &OrderedHashMap<String, String>,
) -> TestRunnerResult {
    let db = &SemanticDatabaseForTesting::default();

    // The reported diagnostics are the recursive module diagnostics, so no function has to be
    // resolved at all - which also allows test blocks that have no function to target.
    let diagnostics = setup_test_module_ex(
        db,
        &test_module_code(inputs),
        inputs.get("crate_settings").map(String::as_str),
        None,
    )
    .get_diagnostics();
    let error = verify_diagnostics_expectation(args, &diagnostics);

    TestRunnerResult {
        outputs: OrderedHashMap::from([("expected_diagnostics".into(), diagnostics)]),
        error,
    }
}

/// Gets the diagnostics for all the modules (including nested) in the given crate.
pub fn get_crate_semantic_diagnostics<'a>(
    db: &'a dyn Database,
    crate_id: CrateId<'a>,
) -> Diagnostics<'a, SemanticDiagnostic<'a>> {
    let submodules = db.crate_modules(crate_id);
    let mut diagnostics = DiagnosticsBuilder::default();
    for submodule_id in submodules.iter() {
        diagnostics.extend(db.module_semantic_diagnostics(*submodule_id).unwrap());
    }
    diagnostics.build()
}

/// Gets the diagnostics for all the modules (including nested) in the given module.
fn get_recursive_module_semantic_diagnostics<'a>(
    db: &'a dyn Database,
    module_id: ModuleId<'a>,
) -> Diagnostics<'a, SemanticDiagnostic<'a>> {
    let mut diagnostics: DiagnosticsBuilder<'_, _> =
        db.module_semantic_diagnostics(module_id).unwrap().into();
    for submodule_id in db.module_submodules_ids(module_id).unwrap().iter() {
        if db.is_submodule_inline(*submodule_id) {
            diagnostics.extend(get_recursive_module_semantic_diagnostics(
                db,
                ModuleId::Submodule(*submodule_id),
            ));
        }
    }
    diagnostics.build()
}