cairo-lang-language-server 2.9.4

use std::collections::VecDeque;
use std::sync::Arc;

use cairo_lang_defs::db::DefsGroup;
use cairo_lang_defs::plugin::MacroPluginMetadata;
use cairo_lang_diagnostics::DiagnosticsBuilder;
use cairo_lang_filesystem::db::FilesGroup;
use cairo_lang_filesystem::ids::{FileId, FileKind, FileLongId, VirtualFile};
use cairo_lang_filesystem::span::{TextSpan, TextWidth};
use cairo_lang_formatter::FormatterConfig;
use cairo_lang_parser::db::ParserGroup;
use cairo_lang_parser::parser::Parser;
use cairo_lang_parser::utils::SimpleParserDatabase;
use cairo_lang_syntax::node::ast::{ExprInlineMacro, ModuleItemList};
use cairo_lang_syntax::node::kind::SyntaxKind;
use cairo_lang_syntax::node::{SyntaxNode, TypedSyntaxNode};
use cairo_lang_utils::Intern;
use indoc::formatdoc;
use lsp_types::TextDocumentPositionParams;

use crate::lang::db::{AnalysisDatabase, LsSemanticGroup, LsSyntaxGroup};
use crate::lang::lsp::{LsProtoGroup, ToCairo};

/// Tries to expand macro, returns it as string.
pub fn expand_macro(db: &AnalysisDatabase, params: &TextDocumentPositionParams) -> Option<String> {
    let file_id = db.file_for_url(&params.text_document.uri)?;
    let node = db.find_syntax_node_at_position(file_id, params.position.to_cairo())?;

    let module_id = db.find_module_file_containing_node(&node)?.0;
    let crate_id = module_id.owning_crate(db);
    let cfg_set = db
        .crate_config(crate_id)
        .and_then(|cfg| cfg.settings.cfg_set.map(Arc::new))
        .unwrap_or(db.cfg_set());
    let edition = db
        .crate_config(module_id.owning_crate(db))
        .map(|cfg| cfg.settings.edition)
        .unwrap_or_default();

    let metadata = MacroPluginMetadata {
        cfg_set: &cfg_set,
        declared_derives: &db.declared_derives(),
        allowed_features: &Default::default(),
        edition,
    };

    let module_file = db.module_main_file(module_id).ok()?;

    let item_ast_node =
        db.first_ancestor_of_kind_respective_child(node.clone(), SyntaxKind::ModuleItemList);
    let macro_ast_node = db.first_ancestor_of_kind(node.clone(), SyntaxKind::ExprInlineMacro);

    let (node_to_expand, top_level_macro_kind) = match (item_ast_node, macro_ast_node) {
        (Some(item_ast_node), Some(macro_ast_node)) => {
            if node_depth(item_ast_node.clone()) > node_depth(macro_ast_node.clone()) {
                (item_ast_node, TopLevelMacroKind::Attribute)
            } else {
                (macro_ast_node, TopLevelMacroKind::Inline)
            }
        }
        (Some(item_ast_node), None) => (item_ast_node, TopLevelMacroKind::Attribute),
        (None, Some(macro_ast_node)) => (macro_ast_node, TopLevelMacroKind::Inline),
        (None, None) => return None,
    };

    let files = match top_level_macro_kind {
        TopLevelMacroKind::Inline => VecDeque::from([module_file]),
        // If this is attribute or derive macro, it can return many files.
        TopLevelMacroKind::Attribute => {
            expanded_macro_files(db, module_file, node_to_expand.clone(), &metadata)?
        }
    };

    // Expand inline macros in generated files.
    expand_inline_macros(db, node_to_expand, files, &metadata, top_level_macro_kind)
}

#[derive(Copy, Clone)]
enum TopLevelMacroKind {
    Inline,
    Attribute,
}

/// Expands macro plugins and returns all file generated by plugins (origin file included).
fn expanded_macro_files(
    db: &dyn DefsGroup,
    module_file: FileId,
    item_ast_node: SyntaxNode,
    metadata: &MacroPluginMetadata<'_>,
) -> Option<VecDeque<FileId>> {
    let syntax_db = db.upcast();

    let item = ModuleItemList::from_syntax_node(syntax_db, item_ast_node.parent()?)
        .elements(syntax_db)
        .into_iter()
        .find(|e| e.as_syntax_node() == item_ast_node)?;

    let mut module_queue = VecDeque::from([(module_file, vec![item])]);
    let mut files = VecDeque::new();

    while let Some((module_file, item_asts)) = module_queue.pop_front() {
        files.push_back(module_file);

        for item_ast in item_asts {
            for plugin in db.macro_plugins() {
                let result = plugin.generate_code(db.upcast(), item_ast.clone(), metadata);

                if let Some(generated) = result.code {
                    let new_file = FileLongId::Virtual(VirtualFile {
                        parent: Some(module_file),
                        name: generated.name,
                        content: generated.content.into(),
                        code_mappings: Default::default(),
                        kind: FileKind::Module,
                    })
                    .intern(db);

                    module_queue.push_back((
                        new_file,
                        db.file_module_syntax(new_file)
                            .ok()?
                            .items(syntax_db)
                            .elements(db.upcast()),
                    ));
                }
                if result.remove_original_item {
                    break;
                }
            }
        }
    }

    Some(files)
}

/// Finds the depth of the node in tree.
fn node_depth(node: SyntaxNode) -> usize {
    std::iter::successors(Some(node), SyntaxNode::parent).count()
}

/// Expands inline macros for each file.
fn expand_inline_macros(
    db: &AnalysisDatabase,
    node_to_expand: SyntaxNode,
    mut files: VecDeque<FileId>,
    metadata: &MacroPluginMetadata<'_>,
    top_level_macro_kind: TopLevelMacroKind,
) -> Option<String> {
    let mut output = String::new();

    expand_inline_macros_in_single_file(
        db,
        metadata,
        files.pop_front().unwrap(),
        &mut files,
        &mut output,
        FileProcessorConfig::main_file(db, node_to_expand, top_level_macro_kind),
    )?;

    while let Some(file) = files.pop_front() {
        expand_inline_macros_in_single_file(
            db,
            metadata,
            file,
            &mut files,
            &mut output,
            FileProcessorConfig::generated_file(db, file, db.file_content(file)?.to_string())?,
        )?;
    }

    if output.is_empty() { None } else { Some(format_output(&output, top_level_macro_kind)) }
}

/// Formats output string.
fn format_output(output: &str, top_level_macro_kind: TopLevelMacroKind) -> String {
    let db = &SimpleParserDatabase::default();
    let virtual_file = FileLongId::Virtual(VirtualFile {
        parent: Default::default(),
        name: Default::default(),
        content: Default::default(),
        code_mappings: Default::default(),
        kind: FileKind::Module,
    })
    .intern(db);

    let syntax_root = match top_level_macro_kind {
        TopLevelMacroKind::Inline => {
            Parser::parse_file_expr(db, &mut DiagnosticsBuilder::default(), virtual_file, output)
                .as_syntax_node()
        }
        TopLevelMacroKind::Attribute => {
            Parser::parse_file(db, &mut DiagnosticsBuilder::default(), virtual_file, output)
                .as_syntax_node()
        }
    };

    cairo_lang_formatter::get_formatted_file(db, &syntax_root, FormatterConfig::default())
        .trim_end()
        .trim_end_matches("\n;")
        .to_owned()
}

struct FileProcessorConfig {
    content: String,
    offset_correction: TextWidth,
    macros: Vec<SyntaxNode>,
}

impl FileProcessorConfig {
    // In original user file we want to expand only selected/hovered ModuleItem.
    fn main_file(
        db: &AnalysisDatabase,
        node_to_expand: SyntaxNode,
        top_level_macro_kind: TopLevelMacroKind,
    ) -> Self {
        Self {
            content: node_to_expand.get_text(db),
            offset_correction: node_to_expand.offset() - Default::default(),
            macros: match top_level_macro_kind {
                TopLevelMacroKind::Inline => vec![node_to_expand],
                TopLevelMacroKind::Attribute => node_to_expand
                    .descendants(db)
                    .filter(|node| node.kind(db) == SyntaxKind::ExprInlineMacro)
                    .collect(),
            },
        }
    }

    // In generated file all ModuleItem should be expanded.
    fn generated_file(db: &AnalysisDatabase, file: FileId, file_content: String) -> Option<Self> {
        Some(Self {
            content: file_content,
            offset_correction: Default::default(),
            macros: db
                .file_syntax(file)
                .ok()?
                .descendants(db)
                .filter(|node| node.kind(db) == SyntaxKind::ExprInlineMacro)
                .collect(),
        })
    }
}

/// Inlines macros in single file, pushes new inlined file if any macro was inlined, otherwise
/// pushes output string.
fn expand_inline_macros_in_single_file(
    db: &AnalysisDatabase,
    metadata: &MacroPluginMetadata<'_>,
    file: FileId,
    files: &mut VecDeque<FileId>,
    output: &mut String,
    mut config: FileProcessorConfig,
) -> Option<()> {
    let plugins = db.inline_macro_plugins();

    if config.macros.is_empty() {
        append_file_with_header(db, file, &config.content, output);
    } else {
        // Iterate in reversed order so positions are not affected by inlining.
        for node in config.macros.into_iter().rev() {
            let inline_macro = ExprInlineMacro::from_syntax_node(db, node.clone());
            let code = plugins
                .get(&inline_macro.path(db).as_syntax_node().get_text_without_trivia(db))?
                .generate_code(db, &inline_macro, metadata)
                .code?
                .content;

            let offset = node.offset().sub_width(config.offset_correction);

            config.content.replace_range(
                TextSpan { start: offset, end: offset.add_width(node.width(db)) }.to_str_range(),
                &code,
            );
        }
        let new_file = FileLongId::Virtual(VirtualFile {
            parent: Some(file),
            name: file.file_name(db).into(),
            content: config.content.into(),
            code_mappings: Default::default(),
            kind: file.kind(db),
        })
        .intern(db);

        files.push_back(new_file);
    };

    Some(())
}

/// Extends output with single file.
fn append_file_with_header(
    db: &AnalysisDatabase,
    file: FileId,
    file_content: &str,
    output: &mut String,
) {
    let file_name = file.file_name(db);
    let file_underscore = "-".repeat(file_name.chars().count());

    output.push_str(&formatdoc!(
        "
        // {file_name}
        // {file_underscore}


        "
    ));

    output.push_str(file_content.trim_end_matches('\n'));
    output.push_str("\n\n");
}