Add collect aux data mechanism

commit-id:af02bb94

scarb/src/compiler/db.rs

@@ -11,30 +11,38 @@ use smol_str::SmolStr;
 use std::sync::Arc;
 use tracing::trace;
 
-use crate::compiler::plugin::proc_macro::ProcMacroHost;
+use crate::compiler::plugin::proc_macro::{ProcMacroHost, ProcMacroHostPlugin};
 use crate::compiler::{CairoCompilationUnit, CompilationUnitAttributes, CompilationUnitComponent};
 use crate::core::Workspace;
 use crate::DEFAULT_MODULE_MAIN_FILE;
 
+pub struct BuiltScarbDatabase {
+    pub db: RootDatabase,
+    pub proc_macro_host: Arc<ProcMacroHostPlugin>,
+}
+
 // TODO(mkaput): ScarbDatabase?
 pub(crate) fn build_scarb_root_database(
     unit: &CairoCompilationUnit,
     ws: &Workspace<'_>,
-) -> Result<RootDatabase> {
+) -> Result<BuiltScarbDatabase> {
     let mut b = RootDatabase::builder();
     b.with_project_config(build_project_config(unit)?);
     b.with_cfg(unit.cfg_set.clone());
-    load_plugins(unit, ws, &mut b)?;
+    let proc_macro_host = load_plugins(unit, ws, &mut b)?;
     let mut db = b.build()?;
     inject_virtual_wrapper_lib(&mut db, unit)?;
-    Ok(db)
+    Ok(BuiltScarbDatabase {
+        db,
+        proc_macro_host,
+    })
 }
 
 fn load_plugins(
     unit: &CairoCompilationUnit,
     ws: &Workspace<'_>,
     builder: &mut RootDatabaseBuilder,
-) -> Result<()> {
+) -> Result<Arc<ProcMacroHostPlugin>> {
     let mut proc_macros = ProcMacroHost::default();
     for plugin_info in &unit.cairo_plugins {
         if plugin_info.builtin {
@@ -46,8 +54,9 @@ fn load_plugins(
             proc_macros.register(plugin_info.package.clone(), ws.config())?;
         }
     }
-    builder.with_plugin_suite(proc_macros.into_plugin_suite());
-    Ok(())
+    let macro_host = Arc::new(proc_macros.into_plugin());
+    builder.with_plugin_suite(ProcMacroHostPlugin::build_plugin_suite(macro_host.clone()));
+    Ok(macro_host)
 }
 
 /// Generates a wrapper lib file for appropriate compilation units.

scarb/src/compiler/plugin/proc_macro/host.rs

@@ -1,6 +1,7 @@
 use crate::compiler::plugin::proc_macro::{FromItemAst, ProcMacroInstance};
 use crate::core::{Config, Package, PackageId};
 use anyhow::Result;
+use cairo_lang_defs::db::DefsGroup;
 use cairo_lang_defs::plugin::PluginDiagnostic;
 use cairo_lang_defs::plugin::{
     DynGeneratedFileAuxData, GeneratedFileAuxData, MacroPlugin, MacroPluginMetadata,
@@ -128,6 +129,35 @@ impl ProcMacroHostPlugin {
             .find(|m| m.declared_attributes().contains(&name))
             .map(|m| m.package_id())
     }
+
+    pub fn build_plugin_suite(macr_host: Arc<Self>) -> PluginSuite {
+        let mut suite = PluginSuite::default();
+        suite.add_plugin_ex(macr_host);
+        suite
+    }
+
+    #[tracing::instrument(level = "trace", skip_all)]
+    pub fn collect_aux_data(&self, db: &dyn DefsGroup) -> Result<()> {
+        for crate_id in db.crates() {
+            let crate_modules = db.crate_modules(crate_id);
+            for module in crate_modules.iter() {
+                let file_infos = db.module_generated_file_infos(*module);
+                if let Ok(file_infos) = file_infos {
+                    for file_info in file_infos.iter().flatten() {
+                        let aux_data = file_info
+                            .aux_data
+                            .as_ref()
+                            .and_then(|ad| ad.as_any().downcast_ref::<ProcMacroAuxData>());
+                        if let Some(_aux_data) = aux_data {
+                            // Do something with `aux_data`.
+                            // dbg!(aux_data);
+                        }
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
 }
 
 impl MacroPlugin for ProcMacroHostPlugin {
@@ -241,10 +271,7 @@ impl ProcMacroHost {
         Ok(())
     }
 
-    pub fn into_plugin_suite(self) -> PluginSuite {
-        let macro_host = ProcMacroHostPlugin::new(self.macros);
-        let mut suite = PluginSuite::default();
-        suite.add_plugin_ex(Arc::new(macro_host));
-        suite
+    pub fn into_plugin(self) -> ProcMacroHostPlugin {
+        ProcMacroHostPlugin::new(self.macros)
     }
 }

scarb/src/ops/compile.rs

@@ -1,13 +1,14 @@
 use anyhow::{anyhow, Result};
 use cairo_lang_compiler::db::RootDatabase;
 use cairo_lang_compiler::diagnostics::DiagnosticsError;
+use cairo_lang_utils::Upcast;
 use indoc::formatdoc;
 use itertools::Itertools;
 
 use scarb_ui::components::Status;
 use scarb_ui::HumanDuration;
 
-use crate::compiler::db::{build_scarb_root_database, has_starknet_plugin};
+use crate::compiler::db::{build_scarb_root_database, has_starknet_plugin, BuiltScarbDatabase};
 use crate::compiler::helpers::build_compiler_config;
 use crate::compiler::plugin::proc_macro;
 use crate::compiler::{CairoCompilationUnit, CompilationUnit, CompilationUnitAttributes};
@@ -104,9 +105,14 @@ fn compile_unit(unit: CompilationUnit, ws: &Workspace<'_>) -> Result<()> {
     let result = match unit {
         CompilationUnit::ProcMacro(unit) => proc_macro::compile_unit(unit, ws),
         CompilationUnit::Cairo(unit) => {
-            let mut db = build_scarb_root_database(&unit, ws)?;
+            let BuiltScarbDatabase {
+                mut db,
+                proc_macro_host,
+            } = build_scarb_root_database(&unit, ws)?;
             check_starknet_dependency(&unit, ws, &db, &package_name);
-            ws.config().compilers().compile(unit, &mut db, ws)
+            ws.config().compilers().compile(unit, &mut db, ws)?;
+            proc_macro_host.collect_aux_data(db.upcast())?;
+            Ok(())
         }
     };
 
@@ -129,7 +135,10 @@ fn check_unit(unit: CompilationUnit, ws: &Workspace<'_>) -> Result<()> {
     match unit {
         CompilationUnit::ProcMacro(unit) => proc_macro::check_unit(unit, ws)?,
         CompilationUnit::Cairo(unit) => {
-            let db = build_scarb_root_database(&unit, ws)?;
+            let BuiltScarbDatabase {
+                db,
+                proc_macro_host,
+            } = build_scarb_root_database(&unit, ws)?;
 
             check_starknet_dependency(&unit, ws, &db, &package_name);
 
@@ -146,6 +155,8 @@ fn check_unit(unit: CompilationUnit, ws: &Workspace<'_>) -> Result<()> {
 
                     anyhow!("could not check `{package_name}` due to previous error")
                 })?;
+
+            proc_macro_host.collect_aux_data(db.upcast())?;
         }
     }