pub(crate) mod config;
mod dependency;
mod expand_type;
mod f;
mod ids;
mod module;
mod templates;
use std::cell::RefCell;
use std::rc::Rc;
use anyhow::Result;
pub use config::configure;
use expand_type::expand_type;
use ids::*;
use itertools::Itertools;
use module::Module;
use proc_macro2::{Ident, TokenStream};
use quote::quote;
use templates::TypeOptions;
use wick_config::{FetchOptions, WickConfiguration};
use wick_interface_types::{OperationSignature, OperationSignatures, TypeDefinition};
#[derive(Debug, Clone, Copy, PartialEq)]
pub(crate) enum Direction {
In,
Out,
}
fn put_impl_in_path(module: &Rc<RefCell<Module>>, mut path_parts_reverse: Vec<&str>, implementation: TokenStream) {
if let Some(next) = path_parts_reverse.pop() {
let module = module.borrow_mut().get_or_add(next);
put_impl_in_path(&module, path_parts_reverse, implementation);
} else {
module.borrow_mut().add(implementation);
}
}
fn gen_types<'a>(
module_name: &str,
config: &mut config::Config,
ty: impl Iterator<Item = &'a TypeDefinition>,
) -> TokenStream {
let types = ty
.map(|v| templates::type_def(config, v, TypeOptions::empty()))
.collect_vec();
let root = Module::new(module_name);
for (mod_parts, implementation) in types {
put_impl_in_path(&root, mod_parts, implementation);
}
let borrowed = root.borrow();
borrowed.codegen()
}
fn gen_wrapper_fns<'a>(
config: &mut config::Config,
component: &Ident,
op: impl Iterator<Item = &'a OperationSignature>,
) -> Vec<TokenStream> {
op.map(|op| templates::gen_wrapper_fn(config, component, op))
.collect_vec()
}
fn gen_trait_fns<'a>(
config: &mut config::Config,
op: impl Iterator<Item = &'a OperationSignature>,
) -> Vec<TokenStream> {
op.map(|op| {
let op_name = id(&snake(op.name()));
let op_config = templates::op_config(config, &generic_config_id(), op);
let op_output = templates::op_outgoing(config, "Outputs", op.outputs());
let op_input = templates::op_incoming(config, "Inputs", op.inputs());
let trait_sig = templates::trait_signature(config, op);
let desc = format!("Types associated with the `{}` operation", op.name());
quote! {
#[doc = #desc]
pub mod #op_name {
#[allow(unused)]
use super::*;
#op_config
#op_output
#op_input
#trait_sig
}
}
})
.collect_vec()
}
#[allow(clippy::needless_pass_by_value, clippy::too_many_lines)]
fn codegen(wick_config: WickConfiguration, gen_config: &mut config::Config) -> Result<String> {
let (ops, types, required, imported, root_config) = match &wick_config {
wick_config::WickConfiguration::Component(comp) => {
let types = comp
.types()?
.into_iter()
.sorted_by(|a, b| a.name().cmp(b.name()))
.collect();
let root_config = comp.config().to_owned();
let requires = comp.requires().clone().to_vec();
let ops = comp.component().operation_signatures();
let imports = comp.import().to_vec();
(ops, types, requires, imports, Some(root_config))
}
wick_config::WickConfiguration::Types(config) => (
config.operation_signatures(),
config.types().to_vec(),
Default::default(),
Default::default(),
None,
),
_ => panic!("Code generation only supports `wick/component` and `wick/types` configurations"),
};
let component_name = id("Component");
let wrapper_fns = gen_wrapper_fns(gen_config, &component_name, ops.iter());
let trait_defs = gen_trait_fns(gen_config, ops.iter());
let typedefs = gen_types("types", gen_config, types.iter());
let init = (!ops.is_empty())
.then(|| templates::gen_component_impls(gen_config, &component_name, ops.iter(), &required, &imported));
let root_config = templates::component_config(gen_config, root_config);
let imports = gen_config.deps.iter().map(|dep| quote! { #dep }).collect_vec();
let imports = quote! { #( #imports )* };
let components = gen_config.components.then(|| {
quote! {
#[derive(Default, Clone)]
#[doc = "The struct that the component implementation hinges around"]
pub struct #component_name;
impl #component_name {
#( #wrapper_fns )*
}
}
});
let expanded = quote! {
#imports
#[allow(unused)]
pub(crate) use wick_component::*;
#[allow(unused)]
pub(crate) use wick_component::WickStream;
pub use wick_component::flow_component::Context;
#init
#root_config
#[doc = "Additional generated types"]
#typedefs
#( #trait_defs )*
#components
};
let source = expanded.to_string();
match syn::parse_file(source.as_str()) {
Ok(reparsed) => {
let formatted = prettyplease::unparse(&reparsed);
Ok(formatted)
}
Err(e) => {
println!("Failed to parse generated code: {}", e);
Ok(source)
}
}
}
pub fn build(config: config::Config) -> Result<()> {
let rt = tokio::runtime::Runtime::new()?;
rt.block_on(async_build(config))
}
pub async fn async_build(mut config: config::Config) -> Result<()> {
let path = config.spec.as_path().to_str().unwrap();
let wick_config = wick_config::WickConfiguration::fetch_uninitialized_tree(path, FetchOptions::default())
.await?
.element
.into_inner();
let src = codegen(wick_config, &mut config)?;
tokio::fs::create_dir_all(&config.out_dir).await?;
let target = config.out_dir.join("mod.rs");
println!("Writing to {}", target.display());
tokio::fs::write(target, src).await?;
Ok(())
}
#[cfg(test)]
mod test {
use anyhow::Result;
use super::*;
use crate::generate::config::ConfigBuilder;
#[tokio::test]
async fn test_build() -> Result<()> {
let mut config = ConfigBuilder::new().spec("./tests/testdata/component.yaml").build()?;
let wick_config = WickConfiguration::fetch(&config.spec, Default::default())
.await
.unwrap()
.finish()?;
let src = codegen(wick_config, &mut config)?;
assert!(src.contains("pub struct Component"));
Ok(())
}
}