#[drink::contract_bundle_provider] macro (#73)

Cargo.toml

@@ -30,6 +30,7 @@ clap = { version = "4.3.4" }
 contract-build = { version = "3.0.1" }
 contract-metadata = { version = "3.2.0" }
 contract-transcode = { version = "3.2.0" }
+convert_case = { version = "0.6.0" }
 crossterm = { version = "0.26.0" }
 parity-scale-codec = { version = "3.4" }
 parity-scale-codec-derive = { version = "3.4" }

drink/src/session/bundle.rs → drink/src/bundle.rs

@@ -5,7 +5,7 @@ use std::{path::PathBuf, rc::Rc};
 use contract_metadata::ContractMetadata;
 use contract_transcode::ContractMessageTranscoder;
 
-use super::error::SessionError;
+use crate::{DrinkResult, Error};
 
 /// A struct representing the result of parsing a `.contract` bundle file.
 ///
@@ -22,18 +22,19 @@ pub struct ContractBundle {
 }
 
 impl ContractBundle {
-    /// Load and parse the information in a `.contract` bundle under `path`, producing a `ContractBundle` struct.
-    pub fn load<P>(path: P) -> Result<Self, SessionError>
+    /// Load and parse the information in a `.contract` bundle under `path`, producing a
+    /// `ContractBundle` struct.
+    pub fn load<P>(path: P) -> DrinkResult<Self>
     where
         P: AsRef<std::path::Path>,
     {
         let metadata: ContractMetadata = ContractMetadata::load(&path).map_err(|e| {
-            SessionError::BundleLoadFailed(format!("Failed to load the contract file:\n{e:?}"))
+            Error::BundleLoadFailed(format!("Failed to load the contract file:\n{e:?}"))
         })?;
 
         let ink_metadata = serde_json::from_value(serde_json::Value::Object(metadata.abi))
             .map_err(|e| {
-                SessionError::BundleLoadFailed(format!(
+                Error::BundleLoadFailed(format!(
                     "Failed to parse metadata from the contract file:\n{e:?}"
                 ))
             })?;
@@ -43,7 +44,7 @@ impl ContractBundle {
         let wasm = metadata
             .source
             .wasm
-            .ok_or(SessionError::BundleLoadFailed(
+            .ok_or(Error::BundleLoadFailed(
                 "Failed to get the WASM blob from the contract file".to_string(),
             ))?
             .0;
@@ -68,7 +69,7 @@ impl ContractBundle {
 #[macro_export]
 macro_rules! local_contract_file {
     () => {
-        drink::session::ContractBundle::local(
+        drink::ContractBundle::local(
             env!("CARGO_MANIFEST_DIR"),
             env!("CARGO_CRATE_NAME").to_owned() + ".contract",
         )

drink/src/errors.rs

@@ -14,6 +14,9 @@ pub enum Error {
     /// Block couldn't have been finalized.
     #[error("Failed to finalize block: {0}")]
     BlockFinalize(String),
+    /// Bundle loading and parsing has failed
+    #[error("Loading the contract bundle has failed: {0}")]
+    BundleLoadFailed(String),
 }
 
 /// Every contract message wraps its return value in `Result<T, LangResult>`. This is the error

drink/src/lib.rs

@@ -3,6 +3,7 @@
 
 #![warn(missing_docs)]
 
+mod bundle;
 pub mod chain_api;
 pub mod contract_api;
 pub mod errors;
@@ -16,7 +17,8 @@ use std::{
     sync::{Arc, Mutex},
 };
 
-pub use drink_test_macro::test;
+pub use bundle::ContractBundle;
+pub use drink_test_macro::{contract_bundle_provider, test};
 pub use errors::Error;
 use frame_support::sp_runtime::{traits::One, BuildStorage};
 pub use frame_support::{

drink/src/session.rs

@@ -16,14 +16,15 @@ use crate::{
     EventRecordOf, Sandbox, DEFAULT_GAS_LIMIT,
 };
 
-mod bundle;
-pub use bundle::ContractBundle;
 pub mod error;
 mod transcoding;
 
 use error::SessionError;
 
-use crate::{errors::MessageResult, mock::MockingApi, session::transcoding::TranscoderRegistry};
+use crate::{
+    bundle::ContractBundle, errors::MessageResult, mock::MockingApi,
+    session::transcoding::TranscoderRegistry,
+};
 
 type Balance = u128;
 
@@ -98,9 +99,10 @@ pub const NO_ARGS: &[String] = &[];
 /// ```rust, no_run
 /// # use drink::{
 /// #   local_contract_file,
-/// #   session::{ContractBundle, Session},
+/// #   session::Session,
 /// #   session::NO_ARGS,
-/// #   runtime::MinimalRuntime
+/// #   runtime::MinimalRuntime,
+/// #   ContractBundle,
 /// # };
 ///
 /// # fn main() -> Result<(), drink::session::error::SessionError> {

drink/src/session/error.rs

@@ -36,7 +36,4 @@ pub enum SessionError {
     /// There is no registered transcoder to encode/decode messages for the called contract.
     #[error("Missing transcoder")]
     NoTranscoder,
-    /// Bundle loading and parsing has failed
-    #[error("Loading the contract bundle has failed: {0}")]
-    BundleLoadFailed(String),
 }

drink/test-macro/Cargo.toml

@@ -16,6 +16,7 @@ proc-macro = true
 cargo_metadata = { workspace = true }
 contract-build = { workspace = true }
 contract-metadata = { workspace = true }
+convert_case = { workspace = true }
 proc-macro2 = { workspace = true }
 syn = { workspace = true, features = ["full"] }
 quote = { workspace = true }

drink/test-macro/src/bundle_provision.rs

@@ -0,0 +1,85 @@
+use std::{collections::HashMap, path::PathBuf};
+
+use convert_case::{Case, Casing};
+use proc_macro2::{Ident, Span, TokenStream as TokenStream2};
+use quote::quote;
+use syn::ItemEnum;
+
+pub struct BundleProviderGenerator {
+    root_contract_name: Option<String>,
+    bundles: HashMap<String, PathBuf>,
+}
+
+impl BundleProviderGenerator {
+    pub fn new<I: Iterator<Item = (String, PathBuf)>>(
+        bundles: I,
+        root_contract_name: Option<String>,
+    ) -> Self {
+        let root_contract_name = root_contract_name.map(|name| name.to_case(Case::Pascal));
+        let bundles = HashMap::from_iter(bundles.map(|(name, path)| {
+            let name = name.to_case(Case::Pascal);
+            (name, path)
+        }));
+
+        if let Some(root_contract_name) = &root_contract_name {
+            assert!(
+                bundles.contains_key(root_contract_name),
+                "Root contract must be part of the bundles"
+            );
+        }
+
+        Self {
+            root_contract_name,
+            bundles,
+        }
+    }
+
+    pub fn generate_bundle_provision(&self, enum_item: ItemEnum) -> TokenStream2 {
+        let enum_name = &enum_item.ident;
+        let enum_vis = &enum_item.vis;
+        let enum_attrs = &enum_item.attrs;
+
+        let local = match &self.root_contract_name {
+            None => quote! {},
+            Some(root_name) => {
+                let local_bundle = self.bundles[root_name].to_str().expect("Invalid path");
+                quote! {
+                    pub fn local() -> ::drink::DrinkResult<::drink::ContractBundle> {
+                        ::drink::ContractBundle::load(#local_bundle)
+                    }
+                }
+            }
+        };
+
+        let (contract_names, matches): (Vec<_>, Vec<_>) = self
+            .bundles
+            .keys()
+            .map(|name| {
+                let name_ident = Ident::new(name, Span::call_site());
+                let path = self.bundles[name].to_str().expect("Invalid path");
+                let matcher = quote! {
+                    #enum_name::#name_ident => ::drink::ContractBundle::load(#path),
+                };
+                (name_ident, matcher)
+            })
+            .unzip();
+
+        quote! {
+            #(#enum_attrs)*
+            #[derive(Copy, Clone, PartialEq, Eq, Debug)]
+            #enum_vis enum #enum_name {
+                #(#contract_names,)*
+            }
+
+            impl #enum_name {
+                #local
+
+                pub fn bundle(self) -> ::drink::DrinkResult<::drink::ContractBundle> {
+                    match self {
+                        #(#matches)*
+                    }
+                }
+            }
+        }
+    }
+}

drink/test-macro/src/contract_building.rs

@@ -1,5 +1,5 @@
 use std::{
-    collections::HashSet,
+    collections::{hash_map::Entry, HashMap},
     path::PathBuf,
     sync::{Mutex, OnceLock},
 };
@@ -10,32 +10,40 @@ use contract_build::{
     OutputType, Target, UnstableFlags, Verbosity,
 };
 
+use crate::bundle_provision::BundleProviderGenerator;
+
 /// Contract package differentiator.
 const INK_AS_DEPENDENCY_FEATURE: &str = "ink-as-dependency";
 
 /// Stores the manifest paths of all contracts that have already been built.
 ///
 /// This prevents from building the same contract for every testcase separately.
-static CONTRACTS_BUILT: OnceLock<Mutex<HashSet<PathBuf>>> = OnceLock::new();
+static CONTRACTS_BUILT: OnceLock<Mutex<HashMap<PathBuf, (String, PathBuf)>>> = OnceLock::new();
 
 /// Build the current package with `cargo contract build --release` (if it is a contract package),
-/// as well as all its contract dependencies.
+/// as well as all its contract dependencies. Return a collection of paths to corresponding
+/// `.contract` files.
 ///
 /// A package is considered as a contract package, if it has the `ink-as-dependency` feature.
 ///
 /// A contract dependency, is a package defined in the `Cargo.toml` file with the
 /// `ink-as-dependency` feature enabled.
-pub fn build_contracts() {
+pub fn build_contracts() -> BundleProviderGenerator {
     let metadata = MetadataCommand::new()
         .exec()
         .expect("Error invoking `cargo metadata`");
 
-    for contract_crate in get_contract_crates(&metadata) {
-        build_contract_crate(contract_crate);
-    }
+    let (maybe_root, contract_deps) = get_contract_crates(&metadata);
+    let maybe_root = maybe_root.map(build_contract_crate);
+    let contract_deps = contract_deps.map(build_contract_crate);
+
+    BundleProviderGenerator::new(
+        maybe_root.clone().into_iter().chain(contract_deps),
+        maybe_root.map(|(name, _)| name),
+    )
 }
 
-fn get_contract_crates(metadata: &Metadata) -> Vec<&Package> {
+fn get_contract_crates(metadata: &Metadata) -> (Option<&Package>, impl Iterator<Item = &Package>) {
     let pkg_lookup = |id| {
         metadata
             .packages
@@ -65,43 +73,52 @@ fn get_contract_crates(metadata: &Metadata) -> Vec<&Package> {
         .expect("Error resolving root package");
     let root = pkg_lookup(root.clone());
 
-    root.features
-        .contains_key(INK_AS_DEPENDENCY_FEATURE)
-        .then_some(root)
-        .into_iter()
-        .chain(contract_deps)
-        .collect()
+    (
+        root.features
+            .contains_key(INK_AS_DEPENDENCY_FEATURE)
+            .then_some(root),
+        contract_deps,
+    )
 }
 
-fn build_contract_crate(pkg: &Package) {
+fn build_contract_crate(pkg: &Package) -> (String, PathBuf) {
     let manifest_path = get_manifest_path(pkg);
 
-    if !CONTRACTS_BUILT
-        .get_or_init(|| Mutex::new(HashSet::new()))
+    match CONTRACTS_BUILT
+        .get_or_init(|| Mutex::new(HashMap::new()))
         .lock()
         .expect("Error locking mutex")
-        .insert(manifest_path.clone().into())
+        .entry(manifest_path.clone().into())
     {
-        return;
-    }
+        Entry::Occupied(ready) => ready.get().clone(),
+        Entry::Vacant(todo) => {
+            let args = ExecuteArgs {
+                manifest_path,
+                verbosity: Verbosity::Default,
+                build_mode: BuildMode::Release,
+                features: Features::default(),
+                network: Network::Online,
+                build_artifact: BuildArtifacts::All,
+                unstable_flags: UnstableFlags::default(),
+                optimization_passes: Some(OptimizationPasses::default()),
+                keep_debug_symbols: false,
+                lint: false,
+                output_type: OutputType::HumanReadable,
+                skip_wasm_validation: false,
+                target: Target::Wasm,
+            };
 
-    let args = ExecuteArgs {
-        manifest_path,
-        verbosity: Verbosity::Default,
-        build_mode: BuildMode::Release,
-        features: Features::default(),
-        network: Network::Online,
-        build_artifact: BuildArtifacts::All,
-        unstable_flags: UnstableFlags::default(),
-        optimization_passes: Some(OptimizationPasses::default()),
-        keep_debug_symbols: false,
-        lint: false,
-        output_type: OutputType::HumanReadable,
-        skip_wasm_validation: false,
-        target: Target::Wasm,
-    };
+            let result = contract_build::execute(args).expect("Error building contract");
+            let bundle_path = result
+                .metadata_result
+                .expect("Metadata should have been generated")
+                .dest_bundle;
 
-    contract_build::execute(args).expect("Error building contract");
+            let new_entry = (pkg.name.clone(), bundle_path);
+            todo.insert(new_entry.clone());
+            new_entry
+        }
+    }
 }
 
 fn get_manifest_path(package: &Package) -> ManifestPath {