Represent a package in Proc Macro Server more reliably

scarb/src/compiler/compilation_unit.rs

@@ -106,6 +106,18 @@ impl CompilationUnitComponent {
     }
 }
 
+impl From<&CompilationUnitComponent> for scarb_proc_macro_server_types::scope::Package {
+    fn from(value: &CompilationUnitComponent) -> Self {
+        Self {
+            // Name given here must be identital to the name present in metadata.
+            // This implementation detail is crucial for communication between PMS and LS.
+            // Always remember to verify this invariant when changing the internals.
+            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::Package;
 
 use crate::{
     compiler::{CairoCompilationUnit, CompilationUnitComponentId, CompilationUnitDependency},
@@ -58,7 +59,7 @@ 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
     /// **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_packages: HashMap<Package, Arc<ProcMacroHostPlugin>>,
 }
 
 impl WorkspaceProcMacros {
@@ -72,25 +73,32 @@ 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 = unit
+                    .components
+                    .iter()
+                    .find(|component| component.id == component_id)
+                    .expect("component should always exist");
+
+                let package: Package = component.into();
+
                 macros_for_components
-                    .entry(key)
+                    .entry(package)
                     .or_default()
                     .append(&mut macro_instances);
             }
         }
 
         let macros_for_packages = macros_for_components
             .into_iter()
-            .map(|(package_id, macro_instances)| {
+            .map(|(package, 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((package, plugin))
             })
             .collect::<Result<HashMap<_, _>>>()?;
 
@@ -100,8 +108,8 @@ impl WorkspaceProcMacros {
     }
 
     /// 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()
+    pub fn get(&self, package: &Package) -> Option<Arc<ProcMacroHostPlugin>> {
+        self.macros_for_packages.get(package).cloned()
     }
 }
 

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.package)
             .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.package)
                 .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.package)
             .with_context(|| format!("No macros found in scope: {context:?}"))?;
 
         let instance = plugin

scarb/tests/proc_macro_prebuilt.rs

@@ -219,14 +219,14 @@ fn load_prebuilt_proc_macros() {
     let mut proc_macro_client = ProcMacroClient::new_without_cargo(&project);
 
     let DefinedMacrosInfo {
-        package_id: compilation_unit_main_component_id,
+        package: compilation_unit_main_component_id,
         ..
     } = proc_macro_client.defined_macros_for_package("test_package");
 
     let response = proc_macro_client
         .request_and_wait::<ExpandInline>(ExpandInlineMacroParams {
             context: ProcMacroScope {
-                package_id: compilation_unit_main_component_id,
+                package: compilation_unit_main_component_id,
             },
             name: "some".to_string(),
             args: TokenStream::new("42".to_string()),

scarb/tests/proc_macro_server.rs

@@ -80,12 +80,12 @@ fn expand_attribute() {
 
     let mut proc_macro_client = ProcMacroClient::new(&project);
 
-    let DefinedMacrosInfo { package_id, .. } =
+    let DefinedMacrosInfo { package, .. } =
         proc_macro_client.defined_macros_for_package("test_package");
 
     let response = proc_macro_client
         .request_and_wait::<ExpandAttribute>(ExpandAttributeParams {
-            context: ProcMacroScope { package_id },
+            context: ProcMacroScope { package },
             attr: "rename_to_very_new_name".to_string(),
             args: TokenStream::empty(),
             item: TokenStream::new("fn some_test_fn(){}".to_string()),
@@ -119,14 +119,14 @@ fn expand_derive() {
 
     let mut proc_macro_client = ProcMacroClient::new(&project);
 
-    let DefinedMacrosInfo { package_id, .. } =
+    let DefinedMacrosInfo { package, .. } =
         proc_macro_client.defined_macros_for_package("test_package");
 
     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 { package },
             derives: vec!["some_derive".to_string()],
             item,
         })
@@ -166,12 +166,12 @@ fn expand_inline() {
 
     let mut proc_macro_client = ProcMacroClient::new(&project);
 
-    let DefinedMacrosInfo { package_id, .. } =
+    let DefinedMacrosInfo { package, .. } =
         proc_macro_client.defined_macros_for_package("test_package");
 
     let response = proc_macro_client
         .request_and_wait::<ExpandInline>(ExpandInlineMacroParams {
-            context: ProcMacroScope { package_id },
+            context: ProcMacroScope { package },
             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,5 +1,7 @@
 use std::collections::HashMap;
 
+use crate::scope::Package;
+
 use super::Method;
 use serde::{Deserialize, Serialize};
 
@@ -16,7 +18,7 @@ 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>,
+    pub macros_by_package_id: HashMap<Package, PackageDefinedMacrosInfo>,
 }
 
 /// Response structure containing lists of all defined macros available for one package.

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` and `discriminator` must refer to the same `CompilationUnitComponent`.
+/// 2. `name` must be identical to the one from `scarb-metadata`.
+///    At the moment, the package name is obtained by calling `CompilationUnitComponent::cairo_package_name`.
+#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize, Hash)]
+pub struct Package {
+    /// Name of the `CompilationUnitComponent` associated with the package.
+    pub name: String,
+    /// A `CompilationUnitComponent` discriminator which represents
+    /// the package uniquely in the workspace.
+    /// `None` only for corelib.
+    pub discriminator: Option<String>,
+}
+
+impl Package {
+    /// Builds a new [`Package`] from `name` and `discriminator`
+    /// without checking them for consistency.
+    ///
+    /// # 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 package in which context the action occurs.
+    pub package: Package,
 }

utils/scarb-test-support/src/proc_macro_server.rs

@@ -7,6 +7,7 @@ use scarb_proc_macro_server_types::methods::defined_macros::DefinedMacros;
 use scarb_proc_macro_server_types::methods::defined_macros::DefinedMacrosParams;
 use scarb_proc_macro_server_types::methods::defined_macros::PackageDefinedMacrosInfo;
 use scarb_proc_macro_server_types::methods::Method;
+use scarb_proc_macro_server_types::scope::Package;
 use std::collections::HashMap;
 use std::io::BufRead;
 use std::io::BufReader;
@@ -69,11 +70,10 @@ pub struct ProcMacroClient {
     responses: HashMap<RequestId, RpcResponse>,
 }
 
-/// A helper structure containing macros available for a package
-/// identified by the `package_id` which is a serialized `PackageId`.
+/// A helper structure containing macros available for a package.
 pub struct DefinedMacrosInfo {
-    /// An ID of the package recognized by PMS.
-    pub package_id: String,
+    /// A package recognized by PMS.
+    pub package: Package,
     /// A proper part of the response, related to the main component of the main CU.
     pub defined_macros: PackageDefinedMacrosInfo,
 }
@@ -191,18 +191,18 @@ impl ProcMacroClient {
         let mut response = response.macros_by_package_id;
 
         // Usually, we can't discover the ID of the mock package used in test, so we extract it from the PMS response.
-        let package_id = response
+        let package = response
             .keys()
-            .find(|cu_id| cu_id.starts_with(package_name))
+            .find(|package| package.name == package_name)
             .expect("Response from Proc Macro Server should contain the main compilation unit.")
             .to_owned();
 
-        let defined_macros = response.remove(&package_id).expect(
+        let defined_macros = response.remove(&package).expect(
             "Response from Proc Macro Server should contain the main compilation unit component.",
         );
 
         DefinedMacrosInfo {
-            package_id,
+            package,
             defined_macros,
         }
     }