Represent a package in Proc Macro Server more reliably

software-mansion · Mar 6, 2025 · 2946b39 · 2946b39
1 parent bf30ab7
commit 2946b39
Show file tree

Hide file tree

Showing 10 changed files with 58 additions and 27 deletions.
diff --git a/scarb/src/compiler/plugin/collection.rs b/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,7 +73,7 @@ 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 key: Package = component_id.package_id.into();
                 macros_for_components
                     .entry(key)
                     .or_default()
@@ -82,15 +83,15 @@ impl WorkspaceProcMacros {
 
         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 +101,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()
     }
 }
 

diff --git a/scarb/src/core/package/id.rs b/scarb/src/core/package/id.rs
@@ -2,6 +2,7 @@ use std::fmt;
 use std::ops::Deref;
 
 use anyhow::Result;
+use scarb_proc_macro_server_types::scope::Package;
 use semver::Version;
 use serde::{Deserialize, Deserializer, Serialize, Serializer};
 use smol_str::SmolStr;
@@ -108,6 +109,15 @@ impl PackageId {
     }
 }
 
+impl From<PackageId> for Package {
+    fn from(value: PackageId) -> Self {
+        Package {
+            name: value.name.to_string(),
+            id: value.to_serialized_string(),
+        }
+    }
+}
+
 impl Deref for PackageId {
     type Target = PackageIdInner;
 

diff --git a/scarb/src/ops/proc_macro_server/methods/expand_attribute.rs b/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

diff --git a/scarb/src/ops/proc_macro_server/methods/expand_derive.rs b/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

diff --git a/scarb/src/ops/proc_macro_server/methods/expand_inline.rs b/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

diff --git a/scarb/tests/proc_macro_prebuilt.rs b/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()),

diff --git a/scarb/tests/proc_macro_server.rs b/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(),

diff --git a/utils/scarb-proc-macro-server-types/src/methods/defined_macros.rs b/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.

diff --git a/utils/scarb-proc-macro-server-types/src/scope.rs b/utils/scarb-proc-macro-server-types/src/scope.rs
@@ -1,8 +1,26 @@
 use serde::{Deserialize, Serialize};
 
+/// Representation of the Scarb package.
+#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize, Hash)]
+pub struct Package {
+    /// Name of the Compilation Unit associated with the package.
+    pub name: String,
+    /// A string-serialized `PackageId` from the Scarb workspace
+    pub id: String,
+}
+
+impl Package {
+    pub fn new(name: impl AsRef<str>, id: impl AsRef<str>) -> Self {
+        Self {
+            name: name.as_ref().to_string(),
+            id: id.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,
+    pub package: Package,
 }
diff --git a/utils/scarb-test-support/src/proc_macro_server.rs b/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,
         }
     }