LS: Proc macro plugins

commit-id:655a951f

crates/cairo-lang-language-server/src/lang/db/mod.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use cairo_lang_defs::db::{DefsDatabase, DefsGroup, try_ext_as_virtual_impl};
 use cairo_lang_doc::db::DocDatabase;
 use cairo_lang_filesystem::cfg::{Cfg, CfgSet};
@@ -13,13 +15,13 @@ use cairo_lang_semantic::inline_macros::get_default_plugin_suite;
 use cairo_lang_semantic::plugin::PluginSuite;
 use cairo_lang_starknet::starknet_plugin_suite;
 use cairo_lang_syntax::node::db::{SyntaxDatabase, SyntaxGroup};
-use cairo_lang_test_plugin::test_plugin_suite;
+use cairo_lang_test_plugin::test_assert_suite;
 use cairo_lang_utils::Upcast;
 
 pub use self::semantic::*;
 pub use self::swapper::*;
 pub use self::syntax::*;
-use super::proc_macros::db::ProcMacroDatabase;
+use super::proc_macros::db::{ProcMacroDatabase, init_proc_macro_cache_group};
 use crate::Tricks;
 
 mod semantic;
@@ -48,17 +50,24 @@ impl AnalysisDatabase {
 
         init_files_group(&mut db);
         init_lowering_group(&mut db, InliningStrategy::Default);
+        // proc-macro-server can be restarted many times but we want to keep these data across
+        // multiple server starts, so init it once per database, not per server.
+        init_proc_macro_cache_group(&mut db);
 
         db.set_cfg_set(Self::initial_cfg_set().into());
 
-        let plugin_suite =
-            [get_default_plugin_suite(), starknet_plugin_suite(), test_plugin_suite()]
-                .into_iter()
-                .chain(tricks.extra_plugin_suites.iter().flat_map(|f| f()))
-                .fold(PluginSuite::default(), |mut acc, suite| {
-                    acc.add(suite);
-                    acc
-                });
+        let plugin_suite = [
+            get_default_plugin_suite(),
+            starknet_plugin_suite(),
+            // test_plugin_suite() TODO(#6551)
+            test_assert_suite(),
+        ]
+        .into_iter()
+        .chain(tricks.extra_plugin_suites.iter().flat_map(|f| f()))
+        .fold(PluginSuite::default(), |mut acc, suite| {
+            acc.add(suite);
+            acc
+        });
         db.apply_plugin_suite(plugin_suite);
 
         db
@@ -85,6 +94,51 @@ impl AnalysisDatabase {
         self.set_inline_macro_plugins(plugin_suite.inline_macro_plugins.into());
         self.set_analyzer_plugins(plugin_suite.analyzer_plugins);
     }
+
+    /// Modifies plugins in the db:
+    ///  - removes plugins from `plugins_to_remove`,
+    ///  - adds plugins from `plugins_to_add`,
+    ///  - retains other plugins present in the db.
+    ///
+    /// Plugins are considered equal if they use the same [`Arc`].
+    pub(crate) fn replace_plugin_suite(
+        &mut self,
+        plugins_to_remove: Option<PluginSuite>,
+        plugins_to_add: PluginSuite,
+    ) {
+        let mut macro_plugins = self.macro_plugins();
+        let mut analyzer_plugins = self.analyzer_plugins();
+        let mut inline_macro_plugins = Arc::unwrap_or_clone(self.inline_macro_plugins());
+
+        if let Some(plugins_to_remove) = plugins_to_remove {
+            macro_plugins.retain(|plugin| {
+                plugins_to_remove
+                    .plugins
+                    .iter()
+                    .all(|previous_plugin| !Arc::ptr_eq(previous_plugin, plugin))
+            });
+            analyzer_plugins.retain(|plugin| {
+                plugins_to_remove
+                    .analyzer_plugins
+                    .iter()
+                    .all(|previous_plugin| !Arc::ptr_eq(previous_plugin, plugin))
+            });
+            inline_macro_plugins.retain(|_, plugin| {
+                plugins_to_remove
+                    .inline_macro_plugins
+                    .iter()
+                    .all(|(_, previous_plugin)| !Arc::ptr_eq(previous_plugin, plugin))
+            });
+        }
+
+        macro_plugins.extend(plugins_to_add.plugins);
+        analyzer_plugins.extend(plugins_to_add.analyzer_plugins);
+        inline_macro_plugins.extend(plugins_to_add.inline_macro_plugins);
+
+        self.set_macro_plugins(macro_plugins);
+        self.set_analyzer_plugins(analyzer_plugins);
+        self.set_inline_macro_plugins(Arc::new(inline_macro_plugins));
+    }
 }
 
 impl salsa::Database for AnalysisDatabase {}

crates/cairo-lang-language-server/src/lang/db/swapper.rs

@@ -3,8 +3,10 @@ use std::panic::{AssertUnwindSafe, catch_unwind};
 use std::sync::Arc;
 use std::time::{Duration, SystemTime};
 
+use cairo_lang_defs::db::DefsGroup;
 use cairo_lang_filesystem::db::FilesGroup;
 use cairo_lang_filesystem::ids::FileId;
+use cairo_lang_semantic::db::SemanticGroup;
 use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
 use cairo_lang_utils::{Intern, LookupIntern};
 use lsp_types::Url;
@@ -13,6 +15,7 @@ use tracing::{error, warn};
 use crate::config::Config;
 use crate::lang::db::AnalysisDatabase;
 use crate::lang::lsp::LsProtoGroup;
+use crate::lang::proc_macros::db::ProcMacroGroup;
 use crate::server::client::Notifier;
 use crate::toolchain::scarb::ScarbToolchain;
 use crate::{Backend, Tricks, env_config};
@@ -86,6 +89,7 @@ impl AnalysisDatabaseSwapper {
     ) {
         let Ok(new_db) = catch_unwind(AssertUnwindSafe(|| {
             let mut new_db = AnalysisDatabase::new(tricks);
+            self.migrate_proc_macro_state(&mut new_db, db);
             self.migrate_file_overrides(&mut new_db, db, open_files);
             self.detect_crates_for_open_files(&mut new_db, open_files, config, notifier);
             new_db
@@ -99,6 +103,21 @@ impl AnalysisDatabaseSwapper {
         self.last_replace = SystemTime::now();
     }
 
+    /// Copies current proc macro state into new db.
+    fn migrate_proc_macro_state(&self, new_db: &mut AnalysisDatabase, old_db: &AnalysisDatabase) {
+        new_db.set_macro_plugins(old_db.macro_plugins());
+        new_db.set_inline_macro_plugins(old_db.inline_macro_plugins());
+        new_db.set_analyzer_plugins(old_db.analyzer_plugins());
+
+        new_db.set_proc_macro_client_status(old_db.proc_macro_client_status());
+
+        // TODO(#6646) probably this should not be part of migration as it will be ever growing.
+        // But diagnostics going crazy every 5 minutes are no better.
+        new_db.set_attribute_macro_resolution(old_db.attribute_macro_resolution());
+        new_db.set_derive_macro_resolution(old_db.derive_macro_resolution());
+        new_db.set_inline_macro_resolution(old_db.inline_macro_resolution());
+    }
+
     /// Makes sure that all open files exist in the new db, with their current changes.
     fn migrate_file_overrides(
         &self,

crates/cairo-lang-language-server/src/lang/proc_macros/controller.rs

@@ -2,6 +2,7 @@ use std::num::NonZeroU32;
 use std::sync::Arc;
 use std::time::Duration;
 
+use cairo_lang_semantic::plugin::PluginSuite;
 use crossbeam::channel::{Receiver, Sender};
 use governor::clock::QuantaClock;
 use governor::state::{InMemoryState, NotKeyed};
@@ -13,11 +14,13 @@ use super::client::connection::ProcMacroServerConnection;
 use super::client::status::ClientStatus;
 use crate::lang::db::AnalysisDatabase;
 use crate::lang::proc_macros::db::ProcMacroGroup;
+use crate::lang::proc_macros::plugins::proc_macro_plugin_suite;
 use crate::toolchain::scarb::ScarbToolchain;
 
 /// Manages lifecycle of proc-macro-server client.
 pub struct ProcMacroClientController {
     scarb: ScarbToolchain,
+    plugin_suite: Option<PluginSuite>,
     initialization_retries: RateLimiter<NotKeyed, InMemoryState, QuantaClock>,
     channels: Option<ProcMacroChannelsSenders>,
 }
@@ -36,6 +39,7 @@ impl ProcMacroClientController {
     pub fn new(scarb: ScarbToolchain) -> Self {
         Self {
             scarb,
+            plugin_suite: Default::default(),
             initialization_retries: RateLimiter::direct(
                 Quota::with_period(
                     Duration::from_secs(180 / 5), // Across 3 minutes (180 seconds) / 5 retries.
@@ -111,13 +115,19 @@ impl ProcMacroClientController {
     pub fn on_response(&mut self, db: &mut AnalysisDatabase) {
         match db.proc_macro_client_status() {
             ClientStatus::Initializing(client) => {
-                let Ok(_defined_macros) = client.finish_initialize() else {
+                let Ok(defined_macros) = client.finish_initialize() else {
                     self.handle_error(db);
 
                     return;
                 };
 
-                // TODO setup db plugins.
+                let new_plugin_suite = proc_macro_plugin_suite(defined_macros);
+
+                // Store current plugins for identity comparison, so we can remove them if we
+                // restart proc-macro-server.
+                let previous_plugin_suite = self.plugin_suite.replace(new_plugin_suite.clone());
+
+                db.replace_plugin_suite(previous_plugin_suite, new_plugin_suite);
 
                 db.set_proc_macro_client_status(ClientStatus::Ready(client));
             }

crates/cairo-lang-language-server/src/lang/proc_macros/db.rs

@@ -1,7 +1,27 @@
+use rustc_hash::FxHashMap;
+use scarb_proc_macro_server_types::methods::ProcMacroResult;
+use scarb_proc_macro_server_types::methods::expand::{
+    ExpandAttributeParams, ExpandDeriveParams, ExpandInlineMacroParams,
+};
+
 use super::client::ClientStatus;
 
 #[salsa::query_group(ProcMacroDatabase)]
 pub trait ProcMacroGroup {
+    #[salsa::input]
+    fn attribute_macro_resolution(&self) -> FxHashMap<ExpandAttributeParams, ProcMacroResult>;
+    #[salsa::input]
+    fn derive_macro_resolution(&self) -> FxHashMap<ExpandDeriveParams, ProcMacroResult>;
+    #[salsa::input]
+    fn inline_macro_resolution(&self) -> FxHashMap<ExpandInlineMacroParams, ProcMacroResult>;
+
     #[salsa::input]
     fn proc_macro_client_status(&self) -> ClientStatus;
 }
+
+pub fn init_proc_macro_cache_group(db: &mut dyn ProcMacroGroup) {
+    db.set_attribute_macro_resolution(Default::default());
+    db.set_derive_macro_resolution(Default::default());
+    db.set_inline_macro_resolution(Default::default());
+    db.set_proc_macro_client_status(Default::default());
+}

crates/cairo-lang-language-server/src/lang/proc_macros/mod.rs

@@ -1,3 +1,4 @@
 mod client;
 pub mod controller;
 pub mod db;
+mod plugins;

crates/cairo-lang-language-server/src/lang/proc_macros/plugins/mod.rs

@@ -0,0 +1,63 @@
+use std::sync::Arc;
+
+use cairo_lang_defs::plugin::{InlineMacroExprPlugin, MacroPlugin};
+use cairo_lang_semantic::plugin::PluginSuite;
+use scarb_proc_macro_server_types::methods::defined_macros::DefinedMacrosResponse;
+
+pub(crate) fn proc_macro_plugin_suite(defined_macros: DefinedMacrosResponse) -> PluginSuite {
+    let mut plugin_suite = PluginSuite::default();
+
+    plugin_suite.add_plugin_ex(Arc::new(ProcMacroPlugin {
+        defined_attributes: defined_macros.attributes,
+        defined_derives: defined_macros.derives,
+        defined_executable_attributes: defined_macros.executables,
+    }));
+
+    let inline_plugin = Arc::new(InlineProcMacroPlugin);
+
+    for inline_macro in defined_macros.inline_macros {
+        plugin_suite.add_inline_macro_plugin_ex(&inline_macro, inline_plugin.clone());
+    }
+
+    plugin_suite
+}
+
+#[derive(Debug)]
+struct ProcMacroPlugin {
+    defined_attributes: Vec<String>,
+    defined_derives: Vec<String>,
+    defined_executable_attributes: Vec<String>,
+}
+
+impl MacroPlugin for ProcMacroPlugin {
+    fn generate_code(
+        &self,
+        _db: &dyn cairo_lang_syntax::node::db::SyntaxGroup,
+        _item_ast: cairo_lang_syntax::node::ast::ModuleItem,
+        _metadata: &cairo_lang_defs::plugin::MacroPluginMetadata<'_>,
+    ) -> cairo_lang_defs::plugin::PluginResult {
+        todo!();
+    }
+
+    fn declared_attributes(&self) -> Vec<String> {
+        [&self.defined_attributes[..], &self.defined_executable_attributes[..]].concat()
+    }
+
+    fn declared_derives(&self) -> Vec<String> {
+        self.defined_derives.clone()
+    }
+}
+
+#[derive(Debug)]
+struct InlineProcMacroPlugin;
+
+impl InlineMacroExprPlugin for InlineProcMacroPlugin {
+    fn generate_code(
+        &self,
+        _db: &dyn cairo_lang_syntax::node::db::SyntaxGroup,
+        _item_ast: &cairo_lang_syntax::node::ast::ExprInlineMacro,
+        _metadata: &cairo_lang_defs::plugin::MacroPluginMetadata<'_>,
+    ) -> cairo_lang_defs::plugin::InlinePluginResult {
+        todo!();
+    }
+}