use std::{
fs,
path::{Path, PathBuf},
process::Command,
};
use anyhow::Ok;
use regex::Regex;
use crate::{
core::manifest::Manifest,
util::{to_pascal_case, to_snake_case, Config, SubstrateResult},
};
pub struct AddOptions {
pub package_name: String,
pub package_path: PathBuf,
pub crate_spec: String,
pub features: Vec<String>,
pub default_features: Option<bool>,
pub source: CrateSource,
}
pub enum CrateSource {
DefaultRegistry,
Git(String, String),
Path(String),
CustomRegistry(String),
}
pub fn add_pallet_std_to_manifest(
cwd: &Path,
runtime_path: &Path,
crate_spec: &str,
) -> SubstrateResult<()> {
let feature = format!("{}/std", crate_spec);
let mut runtime_manifest = Manifest::new(cwd.join(runtime_path).join("Cargo.toml"));
let mut runtime_document = runtime_manifest.read_document()?;
let feature_array = runtime_document["features"].as_table_mut().unwrap()["std"]
.as_array_mut()
.unwrap();
if !feature_array.iter().any(|f| f.as_str().unwrap() == feature) {
feature_array.push::<String>(feature);
}
runtime_manifest.write_document(runtime_document)?;
Ok(())
}
pub fn add_pallet_to_runtime(
cwd: &Path,
runtime_path: &Path,
crate_spec: &str,
) -> SubstrateResult<Option<usize>> {
let runtime_lib_path = cwd.join(runtime_path).join("src/lib.rs");
let mod_name = to_snake_case(crate_spec);
let pallet_trait_existing = Regex::new(
format!(
r"(?xm)
^impl\s+{}::Config\s+for\s+Runtime\s+\{{
[^\}}]+
\}}
",
mod_name
)
.as_ref(),
)?;
let construct_runtime = Regex::new(
r"construct_runtime!\(\s*(?P<visibility>pub\s+)?(?P<variant>enum|struct)\s+Runtime\s*(where[^\{]+)?\{(?P<pallets>[\s\S]+)\}\s*\);",
)?;
let mut pallet_trait_impl = format!("impl {}::Config for Runtime {{ \n", mod_name);
pallet_trait_impl.push_str(&format!(" /* {} Trait config goes here */ \n", mod_name));
pallet_trait_impl.push('}');
let pallet_config = format!(
r"
{}: {},",
to_pascal_case(&mod_name),
mod_name
);
let original = fs::read_to_string(&runtime_lib_path)?;
let mut buffer = original.clone();
let mut line_number: Option<usize> = None;
if pallet_trait_existing.is_match(&original) {
buffer = pallet_trait_existing
.replace(&original, |caps: ®ex::Captures| {
line_number = Some(original[..caps.get(0).unwrap().start()].lines().count() + 1);
pallet_trait_impl.to_owned()
})
.to_string();
} else {
let mat = construct_runtime
.find(&original)
.ok_or_else(|| anyhow::anyhow!("couldn't find construct_runtime call"))?;
line_number = Some(original[..mat.start()].lines().count() + 1);
buffer.insert_str(mat.start(), format!("{}\n\n", pallet_trait_impl).as_str());
};
let modified = buffer.clone();
let caps = construct_runtime
.captures(&modified)
.ok_or_else(|| anyhow::anyhow!("couldn't find construct_runtime call"))?;
let pallets = caps.name("pallets").ok_or_else(|| {
anyhow::anyhow!("couldn't find runtime pallets config inside construct_runtime",)
})?;
let existing_pallets = modified.get(pallets.start()..pallets.end()).unwrap();
let line_number_with_mod_name = existing_pallets
.lines()
.position(|line| line.contains(&mod_name));
if let Some(line_number) = line_number_with_mod_name {
let line_to_replace = existing_pallets.lines().nth(line_number).unwrap();
let new_existing_pallets = existing_pallets.replace(line_to_replace, &pallet_config[1..]);
let new_buffer = modified.replacen(existing_pallets, &new_existing_pallets, 1);
fs::write(runtime_lib_path, new_buffer)?;
} else {
buffer.insert_str(pallets.end() - 2, &pallet_config);
fs::write(runtime_lib_path, buffer)?;
}
Ok(line_number)
}
pub fn add_pallet(opts: &AddOptions, config: &Config) -> SubstrateResult<()> {
let crate_source_arg = match &opts.source {
CrateSource::DefaultRegistry => vec![],
CrateSource::Git(url, branch) => {
let mut args = vec!["--git", url];
if !branch.is_empty() {
args.extend(["--branch", branch]);
}
args
},
CrateSource::Path(path) => vec!["--path", path],
CrateSource::CustomRegistry(registry) => vec!["--registry", registry],
};
let status = Command::new("cargo")
.arg("add")
.arg("-p")
.arg(&opts.package_name)
.arg(&opts.crate_spec)
.arg("--features")
.arg(&opts.features.join(","))
.arg("--no-default-features")
.args(crate_source_arg)
.status()?;
if !status.success() {
return Err(anyhow::anyhow!(
"Failed to install pallet: `{}`",
opts.crate_spec,
));
}
add_pallet_std_to_manifest(config.cwd(), &opts.package_path, &opts.crate_spec)?;
let trait_line_number =
add_pallet_to_runtime(config.cwd(), &opts.package_path, &opts.crate_spec)?;
println!(
"\nPallet `{}` has been successfully added to the runtime!",
opts.crate_spec
);
if let Some(line_number) = trait_line_number {
println!(
"Don't forget to implement the `Config` trait in `runtime/src/lib.rs`, line: {}",
line_number
);
}
Ok(())
}