From 4f5fa68d19d75c3cf7b62a32cfce344bf210a189 Mon Sep 17 00:00:00 2001
From: dcvz <david@dcvz.io>
Date: Thu, 30 Nov 2023 22:14:19 +0100
Subject: [PATCH] Iterate workspaces

---
 xtask/Cargo.toml       |  1 +
 xtask/src/runchecks.rs | 73 ++++++++++++++++++++++++++++++++----------
 2 files changed, 57 insertions(+), 17 deletions(-)

diff --git a/xtask/Cargo.toml b/xtask/Cargo.toml
index d92e7bb24b..5de7713dc8 100644
--- a/xtask/Cargo.toml
+++ b/xtask/Cargo.toml
@@ -9,3 +9,4 @@ license = "MIT OR Apache-2.0"
 [dependencies]
 anyhow = "1.0.75"
 clap = { version = "4.4.8", features = ["derive"] }
+serde_json = { version = "1" }
diff --git a/xtask/src/runchecks.rs b/xtask/src/runchecks.rs
index c18bcafcf3..be5c4a63b6 100644
--- a/xtask/src/runchecks.rs
+++ b/xtask/src/runchecks.rs
@@ -9,6 +9,8 @@ use std::process::{Child, Command, Stdio};
 use std::str;
 use std::time::Instant;
 
+use serde_json::Value;
+
 // Targets constants
 const WASM32_TARGET: &str = "wasm32-unknown-unknown";
 const ARM_TARGET: &str = "thumbv7m-none-eabi";
@@ -43,6 +45,38 @@ fn run_command(command: &str, args: &[&str], command_error: &str, child_error: &
     handle_child_process(command, child_error);
 }
 
+// Get the workspaces available
+fn get_workspaces() -> Vec<String> {
+    // Run `cargo metadata` command to get project metadata
+    let output = Command::new("cargo")
+        .arg("metadata")
+        .output()
+        .expect("Failed to execute command");
+
+    // Parse the JSON output
+    let metadata: Value = serde_json::from_slice(&output.stdout).expect("Failed to parse JSON");
+
+    // Extract workspaces from the metadata, excluding examples/ and xtask
+    let workspaces = metadata["workspace_members"]
+        .as_array()
+        .expect("Expected an array of workspace members")
+        .iter()
+        .filter_map(|member| {
+            let mut parts = member.as_str()?.split_whitespace();
+            let workspace_name = parts.next()?.to_string();
+            let workspace_path = parts.last()?.to_string();
+
+            if !workspace_path.contains("examples/") && workspace_name != "xtask" {
+                Some(workspace_name)
+            } else {
+                None
+            }
+        })
+        .collect();
+
+    workspaces
+}
+
 // Start section print in CI
 fn start_group(title: String) {
     if std::env::var("CI").is_ok() {
@@ -121,17 +155,18 @@ fn cargo_test(params: Params) {
 
 // Run cargo fmt command
 fn cargo_fmt() {
-    // Run cargo fmt
+    start_group("cargo: fmt".to_string());
     run_cargo(
         "fmt",
         ["--check", "--all", "--", "--color=always"].into(),
         "Failed to run cargo fmt",
     );
+    end_group();
 }
 
 // Run cargo clippy command
 fn cargo_clippy() {
-    if std::env::var("CI_RUN").is_ok() {
+    if std::env::var("CI").is_ok() {
         return;
     }
     // Run cargo clippy
@@ -220,8 +255,6 @@ fn run_grcov() {
 
 // Run no_std checks
 fn no_std_checks() {
-    println!("Checks for no_std environment...\n\n");
-
     // Install wasm32 target
     rustup("target", WASM32_TARGET);
 
@@ -243,8 +276,6 @@ fn no_std_checks() {
 
 // Test burn-core with tch and wgpu backend
 fn burn_core_std() {
-    println!("\n\nRun checks for burn-core crate with tch and wgpu backend");
-
     // Run cargo test --features test-tch
     start_group("Test: burn-core (tch)".to_string());
     cargo_test(["-p", "burn-core", "--features", "test-tch"].into());
@@ -292,8 +323,6 @@ fn std_checks() {
     let is_coverage = std::env::var("COVERAGE").is_ok();
     let disable_wgpu = std::env::var("DISABLE_WGPU").is_ok();
 
-    println!("Running std checks");
-
     // Check format
     cargo_fmt();
 
@@ -301,13 +330,17 @@ fn std_checks() {
     cargo_clippy();
 
     // Build each workspace
-    start_group("Build: workspaces".to_string());
-    if disable_wgpu {
-        cargo_build(["--workspace", "--exclude=xtask", "--exclude=burn-wgpu"].into());
-    } else {
-        cargo_build(["--workspace", "--exclude=xtask"].into());
+    let workspaces = get_workspaces();
+    println!("Workspaces: {:?}", workspaces);
+    for workspace in workspaces.clone() {
+        if disable_wgpu && workspace == "burn-wgpu" {
+            continue;
+        }
+
+        start_group(format!("Build: {}", workspace));
+        cargo_build(Params::from(["-p", &workspace]));
+        end_group()
     }
-    end_group();
 
     // Produce documentation for each workspace
     start_group("Docs: workspaces".to_string());
@@ -320,9 +353,15 @@ fn std_checks() {
     }
 
     // Test each workspace
-    start_group("Test: workspaces".to_string());
-    cargo_test(["--workspace"].into());
-    end_group();
+    for workspace in workspaces {
+        if disable_wgpu && workspace == "burn-wgpu" {
+            continue;
+        }
+
+        start_group(format!("Test: {}", workspace));
+        cargo_test(Params::from(["-p", &workspace]));
+        end_group();
+    }
 
     // Test burn-candle with accelerate (macOS only)
     #[cfg(target_os = "macos")]