Fix some issues with references and templates

crates/bevy_api_gen/Cargo.toml

@@ -40,6 +40,7 @@ prettyplease = "0.2"
 convert_case = "0.6"
 syn = { version = "2", features = ["parsing"], no-default-features = true }
 clap-verbosity-flag = "2.2"
+itertools = "0.12"
 
 [build-dependencies]
 toml = "0.8"

crates/bevy_api_gen/src/args.rs

@@ -31,6 +31,29 @@ pub struct Args {
 
     #[command(flatten)]
     pub verbose: Verbosity,
+
+    /// Features to enable when running cargo, requires workspace_root to be provided to work properly
+    #[arg(
+        global = true,
+        short,
+        long,
+        default_value = "",
+        use_value_delimiter = true,
+        value_delimiter = ','
+    )]
+    pub features: Vec<String>,
+
+    /// Disable default features when running cargo, requires workspace_root to be provided to work properly
+    #[arg(global = true, long, default_value = "false")]
+    pub no_default_features: bool,
+
+    /// If provided will use the workspace root to calculate effective dependencies and only generate code for currently active features
+    #[arg(global = true, long, default_value = "bevy")]
+    pub workspace_root: Option<String>,
+
+    /// additional template context in the form of json, provided to the templates under an 'args' key
+    #[arg(global = true, long)]
+    pub template_args: Option<String>,
 }
 
 #[derive(clap::Args, Debug, Clone, Default, Serialize, Deserialize)]
@@ -56,6 +79,26 @@ pub struct Verbosity {
     pub quiet: u8,
 }
 
+impl Verbosity {
+    pub fn get_log_level(&self) -> log::Level {
+        match (self.verbose as isize) - (self.quiet as isize) {
+            0 => log::Level::Info,
+            1 => log::Level::Debug,
+            x if x >= 2 => log::Level::Trace,
+            _ => log::Level::Trace,
+        }
+    }
+
+    pub fn get_rustlog_value(&self) -> &str {
+        match (self.verbose as isize) - (self.quiet as isize) {
+            0 => "info",
+            1 => "debug",
+            x if x >= 2 => "trace",
+            _ => "error",
+        }
+    }
+}
+
 #[derive(Subcommand, Deserialize, Serialize, strum::EnumIs)]
 pub enum Command {
     /// Prints built-in templates to stdout
@@ -106,10 +149,6 @@ pub enum Command {
         /// The data returned is the same as the one provided to the templates.
         #[arg(long, action)]
         template_data_only: bool,
-
-        /// additional template context in the form of json, provided to the templates under an 'args' key
-        #[arg(long)]
-        template_args: Option<String>,
     },
     // /// Final step, once you generate all the crate files you would like to have in your module, you can run this command to
     // /// generate a `mod.rs` file using the `collect.rs` template, which will be provided with all the generated filenames and can 'collect' all the items as it wishes
@@ -136,13 +175,25 @@ pub(crate) fn compute_default_dir() -> String {
 /// Utility for storing and retrieving workspace meta information in env vars
 #[derive(Default, Clone)]
 pub struct WorkspaceMeta {
+    /// the crates in the workspace
     pub crates: Vec<String>,
     pub plugin_target_dir: Utf8PathBuf,
+    pub include_crates: Option<Vec<String>>,
 }
 
 impl WorkspaceMeta {
     const CRATES_ENV_NAME: &'static str = "WORKSPACE_CRATES_META";
     const PLUGIN_DIR_NAME: &'static str = "WORKSPACE_PLUGIN_DIR_META";
+    const INCLUDE_CRATES_ENV_NAME: &'static str = "WORKSPACE_OPT_INCLUDE_CRATES_META";
+
+    /// Returns true if the given crate is in the workspace and if the plugin will run on it
+    pub fn is_workspace_and_included_crate(&self, crate_name: &str) -> bool {
+        self.include_crates
+            .as_ref()
+            .map(|include_crates| include_crates.contains(&crate_name.to_owned()))
+            .unwrap_or(true)
+            && self.crates.contains(&crate_name.to_owned())
+    }
 
     /// Will populate the meta from the environment variables, if empty will use defaults
     pub fn from_env() -> Self {
@@ -155,10 +206,17 @@ impl WorkspaceMeta {
             plugin_target_dir: std::env::var(Self::PLUGIN_DIR_NAME)
                 .unwrap_or_default()
                 .into(),
+            include_crates: std::env::var(Self::INCLUDE_CRATES_ENV_NAME)
+                .ok()
+                .map(|s| s.split(',').map(|s| s.to_owned()).collect()),
         }
     }
 
     pub fn set_env(&self) {
         std::env::set_var(Self::CRATES_ENV_NAME, self.crates.join(","));
+        std::env::set_var(Self::PLUGIN_DIR_NAME, &self.plugin_target_dir);
+        if let Some(include_crates) = &self.include_crates {
+            std::env::set_var(Self::INCLUDE_CRATES_ENV_NAME, include_crates.join(","));
+        }
     }
 }

crates/bevy_api_gen/src/bin/driver.rs

@@ -2,5 +2,6 @@ use bevy_api_gen::*;
 
 fn main() {
     env_logger::init();
+
     rustc_plugin::driver_main(BevyAnalyzer);
 }

crates/bevy_api_gen/src/bin/main.rs

@@ -2,7 +2,7 @@ use std::{
     collections::HashMap,
     env,
     fs::{create_dir_all, File},
-    io::{Read, Write},
+    io::Write,
     path::Path,
     process::{Command, Stdio},
 };
@@ -17,15 +17,37 @@ use tera::Context;
 const BOOTSTRAP_DEPS: [&str; 2] = ["mlua", "bevy_reflect"];
 
 fn main() {
+    // parse this here to early exit on wrong args
+    let args = Args::parse_from(env::args().skip(1));
+
+    if env::var("RUST_LOG").is_err() {
+        env::set_var("RUST_LOG", args.verbose.get_rustlog_value());
+    }
     env_logger::init();
-    debug!("CLI entrypoint");
-    debug!("Creating bootstrapping crate");
 
     let metadata = cargo_metadata::MetadataCommand::new()
         .no_deps()
         .other_options(["--all-features".to_string(), "--offline".to_string()])
         .exec()
         .unwrap();
+    let crates = metadata
+        .workspace_packages()
+        .iter()
+        .map(|p| p.name.to_owned())
+        .collect::<Vec<_>>();
+    let include_crates = match (&args.workspace_root, args.cmd.is_generate()) {
+        (Some(root), true) => {
+            let feature_graph = FeatureGraph::from_metadata(&metadata, root);
+            let dependencies = feature_graph
+                .dependencies_for_features(args.features.as_ref(), !args.no_default_features)
+                .into_iter()
+                .map(|s| s.to_owned())
+                .collect::<Vec<String>>();
+
+            Some(dependencies)
+        }
+        _ => None,
+    };
 
     let plugin_subdir = format!("plugin-{}", env!("RUSTC_CHANNEL"));
     let plugin_target_dir = metadata.target_directory.join(plugin_subdir);
@@ -34,30 +56,12 @@ fn main() {
     // know when to panic if a crate is not found
     // it's also useful to pass around the output directory for our Args default values to be able to compute them
     let workspace_meta = WorkspaceMeta {
-        crates: metadata
-            .workspace_packages()
-            .iter()
-            .map(|p| p.name.to_owned())
-            .collect::<Vec<_>>(),
+        crates,
         plugin_target_dir: plugin_target_dir.clone(),
+        include_crates,
     };
     workspace_meta.set_env();
 
-    // parse this here to early exit on wrong args
-    let args = Args::parse_from(env::args().skip(1));
-
-    if env::var("RUST_LOG").is_err() {
-        env::set_var(
-            "RUST_LOG",
-            match (args.verbose.verbose as isize) - (args.verbose.quiet as isize) {
-                0 => "bevy_api_gen=info",
-                1 => "bevy_api_gen=debug",
-                x if x >= 2 => "bevy_api_gen=trace",
-                _ => "bevy_api_gen=error",
-            },
-        );
-    }
-
     match args.cmd {
         bevy_api_gen::Command::Print { template } => {
             println!(
@@ -113,11 +117,15 @@ fn main() {
                     .collect(),
                 api_name,
             };
-            let context =
+            let mut context =
                 Context::from_serialize(context).expect("Could not create template context");
-            let file = File::create(output.join("mod.rs")).unwrap();
-            tera.render_to(&TemplateKind::SharedModule.to_string(), &context, file)
+
+            extend_context_with_args(args.template_args.as_deref(), &mut context);
+
+            let mut file = File::create(output.join("mod.rs")).unwrap();
+            tera.render_to(&TemplateKind::SharedModule.to_string(), &context, &mut file)
                 .expect("Failed to render mod.rs");
+            file.flush().unwrap();
             log::info!("Succesfully generated mod.rs");
             return;
         }
@@ -131,18 +139,7 @@ fn main() {
 
     write_bootstrap_files(temp_dir.path());
 
-    debug!("Building bootstrapping crate");
-
-    let mut cmd = Command::new("cargo")
-        .current_dir(&temp_dir)
-        .stdout(Stdio::piped())
-        .args(["build", "--message-format=json"])
-        .spawn()
-        .unwrap();
-
-    let reader = std::io::BufReader::new(cmd.stdout.take().unwrap());
-
-    let bootstrap_rlibs = build_bootstrap(reader, cmd, &plugin_target_dir.join("bootstrap"));
+    let bootstrap_rlibs = build_bootstrap(temp_dir.path(), &plugin_target_dir.join("bootstrap"));
 
     if bootstrap_rlibs.len() == BOOTSTRAP_DEPS.len() {
         let extern_args = bootstrap_rlibs
@@ -178,10 +175,11 @@ fn main() {
 /// Build bootstrap files if they don't exist
 /// use cached ones otherwise
 fn build_bootstrap(
-    reader: std::io::BufReader<std::process::ChildStdout>,
-    mut cmd: std::process::Child,
+    temp_dir: &Path,
     cache_dir: &Utf8Path,
 ) -> HashMap<String, cargo_metadata::camino::Utf8PathBuf> {
+    debug!("Building bootstrapping crate");
+
     // first check cache
     if cache_dir.exists() {
         let mut bootstrap_rlibs = HashMap::with_capacity(BOOTSTRAP_DEPS.len());
@@ -193,6 +191,16 @@ fn build_bootstrap(
         return bootstrap_rlibs;
     }
 
+    // run build command
+    let mut cmd = Command::new("cargo")
+        .current_dir(temp_dir)
+        .stdout(Stdio::piped())
+        .args(["build", "--message-format=json"])
+        .spawn()
+        .unwrap();
+
+    let reader = std::io::BufReader::new(cmd.stdout.take().unwrap());
+
     std::fs::create_dir_all(cache_dir).unwrap();
 
     let mut bootstrap_rlibs = HashMap::with_capacity(BOOTSTRAP_DEPS.len());

crates/bevy_api_gen/src/callback.rs

@@ -43,10 +43,12 @@ impl rustc_driver::Callbacks for BevyAnalyzerCallbacks {
         } = &self.args.cmd
         {
             templates_dir = templates.to_owned();
+            if let Some(meta_output) = meta_output {
+                meta_dirs.push(meta_output.to_owned())
+            };
             meta_dirs.push(output.to_owned());
             meta.iter()
                 .flatten()
-                .chain(meta_output.iter())
                 .for_each(|m| meta_dirs.push(m.to_owned()));
         };
 
@@ -63,7 +65,7 @@ impl rustc_driver::Callbacks for BevyAnalyzerCallbacks {
 
             let mut ctxt = crate::BevyCtxt::new(
                 tcx,
-                meta_dirs,
+                &meta_dirs,
                 WorkspaceMeta::from_env(),
                 include_private,
                 Some(Box::new(move |import_path| {

crates/bevy_api_gen/src/context.rs

@@ -24,7 +24,7 @@ impl<'tcx> BevyCtxt<'tcx> {
     /// Creates a new context with the provided TyCtxt and meta directories
     pub(crate) fn new(
         tcx: TyCtxt<'tcx>,
-        meta_dirs: Vec<Utf8PathBuf>,
+        meta_dirs: &[Utf8PathBuf],
         workspace_meta: crate::WorkspaceMeta,
         include_private_paths: bool,
         import_path_processor: Option<Box<dyn Fn(&str) -> String>>,
@@ -33,7 +33,7 @@ impl<'tcx> BevyCtxt<'tcx> {
             tcx,
             reflect_types: Default::default(),
             cached_traits: Default::default(),
-            meta_loader: MetaLoader::new(meta_dirs, workspace_meta),
+            meta_loader: MetaLoader::new(meta_dirs.to_vec(), workspace_meta),
             template_context: Default::default(),
             path_finder: ImportPathFinder::new(tcx, include_private_paths, import_path_processor),
         }
@@ -55,8 +55,8 @@ impl<'tcx> BevyCtxt<'tcx> {
 
 #[derive(Clone, Default, Debug)]
 pub(crate) struct ReflectType<'tcx> {
-    /// Trait implementations for the reflect type (from a selection)
-    pub(crate) trait_impls: Option<Vec<DefId>>,
+    /// Map from traits to their implementations for the reflect type (from a selection)
+    pub(crate) trait_impls: Option<HashMap<DefId, DefId>>,
     /// Information about the ADT structure, fields, and variants
     pub(crate) variant_data: Option<AdtDef<'tcx>>,
     /// Functions passing criteria to be proxied
@@ -84,8 +84,16 @@ impl ReflectType<'_> {
 pub(crate) const DEF_PATHS_FROM_LUA: [&str; 2] = ["value::FromLuaMulti", "mlua::FromLuaMulti"];
 pub(crate) const DEF_PATHS_INTO_LUA: [&str; 2] = ["value::IntoLuaMulti", "mlua::IntoLuaMulti"];
 pub(crate) const DEF_PATHS_REFLECT: [&str; 2] = ["bevy_reflect::Reflect", "reflect::Reflect"];
-pub(crate) const FN_SOURCE_TRAITS: [&str; 12] = [
+pub(crate) const DEF_PATHS_GET_TYPE_REGISTRATION: [&str; 2] = [
+    "bevy_reflect::GetTypeRegistration",
+    "reflect::GetTypeRegistration",
+];
+
+/// A collection of traits which we search for in the codebase, some of these are necessary to figure out if a type
+/// is Clone or Debug for the purposes of the macro code generation
+pub(crate) const FN_SOURCE_TRAITS: [&str; 13] = [
     // PRINTING
+    "std::fmt::Debug",
     "std::string::ToString",
     // OWNERSHIP
     "std::clone::Clone",
@@ -108,6 +116,7 @@ pub(crate) struct CachedTraits {
     pub(crate) mlua_from_lua_multi: Option<DefId>,
     pub(crate) mlua_into_lua_multi: Option<DefId>,
     pub(crate) bevy_reflect_reflect: Option<DefId>,
+    pub(crate) bevy_reflect_get_type_registration: Option<DefId>,
     /// Traits whose methods can be included in the generated code
     pub(crate) fn_source_traits: HashMap<String, DefId>,
 }
@@ -118,7 +127,7 @@ impl CachedTraits {
     }
 
     pub(crate) fn has_all_bevy_traits(&self) -> bool {
-        self.bevy_reflect_reflect.is_some()
+        self.bevy_reflect_reflect.is_some() && self.bevy_reflect_get_type_registration.is_some()
     }
 
     pub(crate) fn has_all_fn_source_traits(&self) -> bool {