Represent a package in Proc Macro Server more reliably

scarb/src/compiler/compilation_unit.rs

@@ -106,6 +106,20 @@ impl CompilationUnitComponent {
     }
 }
 
+impl From<&CompilationUnitComponent>
+    for scarb_proc_macro_server_types::scope::CompilationUnitComponent
+{
+    fn from(value: &CompilationUnitComponent) -> Self {
+        // `name` and `discriminator` used here must be identital to the scarb-metadata.
+        // This implementation detail is crucial for communication between PMS and LS.
+        // Always remember to verify this invariant when changing the internals.
+        Self {
+            name: value.cairo_package_name().to_string(),
+            discriminator: value.id.to_discriminator().map(Into::into),
+        }
+    }
+}
+
 /// The kind of the compilation unit dependency.
 #[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Debug)]
 pub enum CompilationUnitDependency {

scarb/src/compiler/plugin/collection.rs

@@ -3,6 +3,7 @@ use std::{collections::HashMap, sync::Arc};
 use anyhow::Result;
 use cairo_lang_semantic::{inline_macros::get_default_plugin_suite, plugin::PluginSuite};
 use itertools::Itertools;
+use scarb_proc_macro_server_types::scope::CompilationUnitComponent;
 
 use crate::{
     compiler::{CairoCompilationUnit, CompilationUnitComponentId, CompilationUnitDependency},
@@ -48,17 +49,12 @@ impl PluginsForComponents {
 // NOTE: Since this structure is used to handle JsonRPC requests, its keys have to be serialized to strings.
 //
 /// A container for Proc Macro Server to manage macros present in the analyzed workspace.
-///
-/// # Invariant
-/// Correct usage of this struct during proc macro server <-> LS communication
-/// relies on the implicit contract that keys of `macros_for_packages` are of form
-/// `PackageId.to_serialized_string()` which is always equal to
-/// `scarb_metadata::CompilationUnitComponentId.repr`.
 pub struct WorkspaceProcMacros {
-    /// A mapping of the form: `package (as a serialized [`PackageId`]) -> plugin`.
-    /// Contains IDs of all packages from the workspace, each mapped to a [`ProcMacroHostPlugin`] which contains
+    /// A mapping of the form: `cu_component (as a [`CompilationUnitComponent`]) -> plugin`.
+    /// Contains IDs of all components of all compilation units from the workspace,
+    /// each mapped to a [`ProcMacroHostPlugin`] which contains
     /// **all proc macro dependencies of the package** collected from **all compilation units it appears in**.
-    pub macros_for_packages: HashMap<String, Arc<ProcMacroHostPlugin>>,
+    pub macros_for_components: HashMap<CompilationUnitComponent, Arc<ProcMacroHostPlugin>>,
 }
 
 impl WorkspaceProcMacros {
@@ -71,37 +67,42 @@ impl WorkspaceProcMacros {
 
         for &unit in compilation_units {
             for (component_id, mut macro_instances) in collect_proc_macros(workspace, unit)? {
-                // Here we assume that CompilationUnitComponentId.repr == PackageId.to_serialized_string().
-                let key = component_id.package_id.to_serialized_string();
+                let component: CompilationUnitComponent = unit
+                    .components
+                    .iter()
+                    .find(|component| component.id == component_id)
+                    .expect("component should always exist")
+                    .into();
+
                 macros_for_components
-                    .entry(key)
+                    .entry(component)
                     .or_default()
                     .append(&mut macro_instances);
             }
         }
 
-        let macros_for_packages = macros_for_components
+        let macros_for_components = macros_for_components
             .into_iter()
-            .map(|(package_id, macro_instances)| {
+            .map(|(component, macro_instances)| {
                 let deduplicated_instances = macro_instances
                     .into_iter()
                     .unique_by(|instance| instance.package_id())
                     .collect();
 
                 let plugin = Arc::new(ProcMacroHostPlugin::try_new(deduplicated_instances)?);
 
-                Ok((package_id, plugin))
+                Ok((component, plugin))
             })
             .collect::<Result<HashMap<_, _>>>()?;
 
         Ok(Self {
-            macros_for_packages,
+            macros_for_components,
         })
     }
 
-    /// Returns a [`ProcMacroHostPlugin`] assigned to the package with `package_id`.
-    pub fn get(&self, package_id: &str) -> Option<Arc<ProcMacroHostPlugin>> {
-        self.macros_for_packages.get(package_id).cloned()
+    /// Returns a [`ProcMacroHostPlugin`] assigned to the [`CompilationUnitComponent`].
+    pub fn get(&self, component: &CompilationUnitComponent) -> Option<Arc<ProcMacroHostPlugin>> {
+        self.macros_for_components.get(component).cloned()
     }
 }
 

scarb/src/ops/proc_macro_server/methods/defined_macros.rs

@@ -1,10 +1,10 @@
-use std::{collections::HashMap, sync::Arc};
+use std::sync::Arc;
 
 use anyhow::Result;
 use cairo_lang_defs::plugin::MacroPlugin;
 use convert_case::{Case, Casing};
 use scarb_proc_macro_server_types::methods::defined_macros::{
-    DefinedMacros, DefinedMacrosResponse, PackageDefinedMacrosInfo,
+    CompilationUnitComponentMacros, DefinedMacros, DefinedMacrosResponse,
 };
 
 use super::Handler;
@@ -15,10 +15,10 @@ impl Handler for DefinedMacros {
         workspace_macros: Arc<WorkspaceProcMacros>,
         _params: Self::Params,
     ) -> Result<Self::Response> {
-        let macros_by_package_id = workspace_macros
-            .macros_for_packages
+        let macros_for_cu_components = workspace_macros
+            .macros_for_components
             .iter()
-            .map(|(package_id, plugin)| {
+            .map(|(component, plugin)| {
                 let attributes = plugin.declared_attributes_without_executables();
                 let inline_macros = plugin.declared_inline_macros();
                 let derives = plugin
@@ -28,20 +28,18 @@ impl Handler for DefinedMacros {
                     .collect();
                 let executables = plugin.executable_attributes();
 
-                (
-                    package_id.to_owned(),
-                    PackageDefinedMacrosInfo {
-                        attributes,
-                        inline_macros,
-                        derives,
-                        executables,
-                    },
-                )
+                CompilationUnitComponentMacros {
+                    component: component.to_owned(),
+                    attributes,
+                    inline_macros,
+                    derives,
+                    executables,
+                }
             })
-            .collect::<HashMap<_, _>>();
+            .collect();
 
         Ok(DefinedMacrosResponse {
-            macros_by_package_id,
+            macros_for_cu_components,
         })
     }
 }

scarb/src/ops/proc_macro_server/methods/expand_attribute.rs

@@ -19,7 +19,7 @@ impl Handler for ExpandAttribute {
         } = params;
 
         let plugin = workspace_macros
-            .get(&context.package_id)
+            .get(&context.component)
             .with_context(|| format!("No macros found in scope: {context:?}"))?;
 
         let instance = plugin

scarb/src/ops/proc_macro_server/methods/expand_derive.rs

@@ -29,7 +29,7 @@ impl Handler for ExpandDerive {
             let expansion = Expansion::new(derive.to_case(Case::Snake), ExpansionKind::Derive);
 
             let plugin = workspace_macros
-                .get(&context.package_id)
+                .get(&context.component)
                 .with_context(|| format!("No macros found in scope {context:?}"))?;
 
             let instance = plugin

scarb/src/ops/proc_macro_server/methods/expand_inline.rs

@@ -19,7 +19,7 @@ impl Handler for ExpandInline {
         } = params;
 
         let plugin = workspace_macros
-            .get(&context.package_id)
+            .get(&context.component)
             .with_context(|| format!("No macros found in scope: {context:?}"))?;
 
         let instance = plugin

scarb/tests/proc_macro_prebuilt.rs

@@ -8,7 +8,7 @@ use scarb_proc_macro_server_types::methods::expand::{ExpandInline, ExpandInlineM
 use scarb_proc_macro_server_types::scope::ProcMacroScope;
 use scarb_test_support::cairo_plugin_project_builder::CairoPluginProjectBuilder;
 use scarb_test_support::command::Scarb;
-use scarb_test_support::proc_macro_server::{DefinedMacrosInfo, ProcMacroClient};
+use scarb_test_support::proc_macro_server::ProcMacroClient;
 use scarb_test_support::project_builder::ProjectBuilder;
 use scarb_test_support::workspace_builder::WorkspaceBuilder;
 use snapbox::cmd::Command;
@@ -218,16 +218,13 @@ fn load_prebuilt_proc_macros() {
 
     let mut proc_macro_client = ProcMacroClient::new_without_cargo(&project);
 
-    let DefinedMacrosInfo {
-        package_id: compilation_unit_main_component_id,
-        ..
-    } = proc_macro_client.defined_macros_for_package("test_package");
+    let component = proc_macro_client
+        .defined_macros_for_package("test_package")
+        .component;
 
     let response = proc_macro_client
         .request_and_wait::<ExpandInline>(ExpandInlineMacroParams {
-            context: ProcMacroScope {
-                package_id: compilation_unit_main_component_id,
-            },
+            context: ProcMacroScope { component },
             name: "some".to_string(),
             args: TokenStream::new("42".to_string()),
         })

scarb/tests/proc_macro_server.rs

@@ -9,7 +9,6 @@ use scarb_proc_macro_server_types::methods::expand::ExpandInline;
 use scarb_proc_macro_server_types::methods::expand::ExpandInlineMacroParams;
 use scarb_proc_macro_server_types::scope::ProcMacroScope;
 use scarb_test_support::cairo_plugin_project_builder::CairoPluginProjectBuilder;
-use scarb_test_support::proc_macro_server::DefinedMacrosInfo;
 use scarb_test_support::proc_macro_server::ProcMacroClient;
 use scarb_test_support::proc_macro_server::SIMPLE_MACROS;
 use scarb_test_support::project_builder::ProjectBuilder;
@@ -34,8 +33,7 @@ fn defined_macros() {
 
     let mut proc_macro_client = ProcMacroClient::new(&project);
 
-    let DefinedMacrosInfo { defined_macros, .. } =
-        proc_macro_client.defined_macros_for_package("test_package");
+    let defined_macros = proc_macro_client.defined_macros_for_package("test_package");
 
     assert_eq!(&defined_macros.attributes, &["some".to_string()]);
     assert_eq!(&defined_macros.derives, &["some_derive".to_string()]);
@@ -80,12 +78,13 @@ fn expand_attribute() {
 
     let mut proc_macro_client = ProcMacroClient::new(&project);
 
-    let DefinedMacrosInfo { package_id, .. } =
-        proc_macro_client.defined_macros_for_package("test_package");
+    let component = proc_macro_client
+        .defined_macros_for_package("test_package")
+        .component;
 
     let response = proc_macro_client
         .request_and_wait::<ExpandAttribute>(ExpandAttributeParams {
-            context: ProcMacroScope { package_id },
+            context: ProcMacroScope { component },
             attr: "rename_to_very_new_name".to_string(),
             args: TokenStream::empty(),
             item: TokenStream::new("fn some_test_fn(){}".to_string()),
@@ -119,14 +118,15 @@ fn expand_derive() {
 
     let mut proc_macro_client = ProcMacroClient::new(&project);
 
-    let DefinedMacrosInfo { package_id, .. } =
-        proc_macro_client.defined_macros_for_package("test_package");
+    let component = proc_macro_client
+        .defined_macros_for_package("test_package")
+        .component;
 
     let item = TokenStream::new("fn some_test_fn(){}".to_string());
 
     let response = proc_macro_client
         .request_and_wait::<ExpandDerive>(ExpandDeriveParams {
-            context: ProcMacroScope { package_id },
+            context: ProcMacroScope { component },
             derives: vec!["some_derive".to_string()],
             item,
         })
@@ -166,12 +166,13 @@ fn expand_inline() {
 
     let mut proc_macro_client = ProcMacroClient::new(&project);
 
-    let DefinedMacrosInfo { package_id, .. } =
-        proc_macro_client.defined_macros_for_package("test_package");
+    let component = proc_macro_client
+        .defined_macros_for_package("test_package")
+        .component;
 
     let response = proc_macro_client
         .request_and_wait::<ExpandInline>(ExpandInlineMacroParams {
-            context: ProcMacroScope { package_id },
+            context: ProcMacroScope { component },
             name: "replace_all_15_with_25".to_string(),
             args: TokenStream::new(
                 "struct A { field: 15 , other_field: macro_call!(12)}".to_string(),

utils/scarb-proc-macro-server-types/src/methods/defined_macros.rs

@@ -1,28 +1,31 @@
-use std::collections::HashMap;
+use crate::scope::CompilationUnitComponent;
 
 use super::Method;
 use serde::{Deserialize, Serialize};
 
-/// Response structure containing a mapping from package IDs
-/// to the information about the macros they use.
+/// Response structure containing a listed information
+/// about the macros used by packages from the workspace.
 ///
 /// # Invariant
-/// Correct usage of this struct during proc macro server <-> LS communication
-/// relies on the implicit contract that keys of `macros_by_package_id` are of form
-/// `PackageId.to_serialized_string()` which is always equal to
-/// `scarb_metadata::CompilationUnitComponentId.repr`.
+/// Each [`CompilationUnitComponentMacros`] in `macros_for_packages` should have
+/// an exclusive `package` field which identifies it in the response.
+/// Effectively, it simulates a HashMap which cannot be used directly
+/// because of the JSON serialization.
 #[derive(Clone, Debug, Default, Serialize, Deserialize)]
 pub struct DefinedMacrosResponse {
-    /// A mapping of the form: `package (as a serialized `PackageId`) -> macros info`.
-    /// Contains serialized IDs of all packages from the workspace,
-    /// mapped to the [`PackageDefinedMacrosInfo`], describing macros available for them.
-    pub macros_by_package_id: HashMap<String, PackageDefinedMacrosInfo>,
+    /// A list of [`CompilationUnitComponentMacros`], describing macros
+    /// available for each package from the workspace.
+    pub macros_for_cu_components: Vec<CompilationUnitComponentMacros>,
 }
 
-/// Response structure containing lists of all defined macros available for one package.
-/// Details the types of macros that can be expanded, such as attributes, inline macros, and derives.
+/// Response structure containing lists of all defined macros available for one compilation unit component.
+/// Provides the types of macros that can be expanded, such as attributes, inline macros, and derives.
 #[derive(Clone, Debug, Default, Serialize, Deserialize)]
-pub struct PackageDefinedMacrosInfo {
+pub struct CompilationUnitComponentMacros {
+    /// A component for which the macros are defined.
+    /// It should identify [`CompilationUnitComponentMacros`]
+    /// uniquely in the [`DefinedMacrosResponse`].
+    pub component: CompilationUnitComponent,
     /// List of attribute macro names available.
     pub attributes: Vec<String>,
     /// List of inline macro names available.

utils/scarb-proc-macro-server-types/src/scope.rs

@@ -1,8 +1,40 @@
 use serde::{Deserialize, Serialize};
 
+/// Representation of the Scarb package.
+///
+/// # Invariants
+/// 1. (`name`, `discriminator`) pair must represent the package uniquely in the workspace.
+/// 2. `name` and `discriminator` must refer to the same `CompilationUnitComponent` and must be identical to those from `scarb-metadata`.
+///    At the moment, they are obtained using `CompilationUnitComponent::cairo_package_name`
+///    and `CompilationUnitComponentId::to_discriminator`, respectively.
+#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize, Hash)]
+pub struct CompilationUnitComponent {
+    /// Name of the `CompilationUnitComponent` associated with the package.
+    pub name: String,
+    /// A `CompilationUnitComponent` discriminator.
+    /// `None` only for corelib.
+    pub discriminator: Option<String>,
+}
+
+impl CompilationUnitComponent {
+    /// Builds a new [`CompilationUnitComponent`] from `name` and `discriminator`
+    /// without checking for consistency between them and with the metadata.
+    ///
+    /// # Safety
+    /// Communication between PMS and LS relies on the invariant that `name` and `discriminator`
+    /// refer to the same CU component and are consistent with `scarb-metadata`.
+    /// The caller must ensure correctness of the provided values.
+    pub fn new(name: impl AsRef<str>, discriminator: impl AsRef<str>) -> Self {
+        Self {
+            name: name.as_ref().to_string(),
+            discriminator: Some(discriminator.as_ref().to_string()),
+        }
+    }
+}
+
 /// A description of the location in the workspace where particular macro is available.
 #[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize, Hash)]
 pub struct ProcMacroScope {
-    /// Serialized `PackageId` of the package in which context the action occurs.
-    pub package_id: String,
+    /// A [`CompilationUnitComponent`] in which context the action occurs.
+    pub component: CompilationUnitComponent,
 }