Implement output of colored messages with optional check context

10 files changed, 152 insertions, 55 deletions

diff --git a/src/tools/features-status-dump/src/main.rs b/src/tools/features-status-dump/src/main.rs
index 1ce98d1506d..a4f88362ab8 100644
--- a/src/tools/features-status-dump/src/main.rs
+++ b/src/tools/features-status-dump/src/main.rs
@@ -5,6 +5,7 @@ use std::path::PathBuf;
 
 use anyhow::{Context, Result};
 use clap::Parser;
+use tidy::diagnostics::RunningCheck;
 use tidy::features::{Feature, collect_lang_features, collect_lib_features};
 
 #[derive(Debug, Parser)]
@@ -29,7 +30,7 @@ struct FeaturesStatus {
 fn main() -> Result<()> {
     let Cli { compiler_path, library_path, output_path } = Cli::parse();
 
-    let lang_features_status = collect_lang_features(&compiler_path, &mut false);
+    let lang_features_status = collect_lang_features(&compiler_path, &mut RunningCheck::new_noop());
     let lib_features_status = collect_lib_features(&library_path)
         .into_iter()
         .filter(|&(ref name, _)| !lang_features_status.contains_key(name))
diff --git a/src/tools/tidy/src/alphabetical/tests.rs b/src/tools/tidy/src/alphabetical/tests.rs
index 4d05bc33ced..b181ab8f744 100644
--- a/src/tools/tidy/src/alphabetical/tests.rs
+++ b/src/tools/tidy/src/alphabetical/tests.rs
@@ -1,19 +1,22 @@
-use std::io::Write;
-use std::str::from_utf8;
+use std::path::Path;
 
-use super::*;
+use crate::alphabetical::check_lines;
+use crate::diagnostics::DiagCtx;
 
 #[track_caller]
 fn test(lines: &str, name: &str, expected_msg: &str, expected_bad: bool) {
-    let mut actual_msg = Vec::new();
-    let mut actual_bad = false;
-    let mut err = |args: &_| {
-        write!(&mut actual_msg, "{args}")?;
-        Ok(())
-    };
-    check_lines(&name, lines.lines().enumerate(), &mut err, &mut actual_bad);
-    assert_eq!(expected_msg, from_utf8(&actual_msg).unwrap());
-    assert_eq!(expected_bad, actual_bad);
+    let diag_ctx = DiagCtx::new(Path::new("/"), false);
+    let mut check = diag_ctx.start_check("alphabetical-test");
+    check_lines(&name, lines.lines().enumerate(), &mut check);
+
+    assert_eq!(expected_bad, check.is_bad());
+    let errors = check.get_errors();
+    if expected_bad {
+        assert_eq!(errors.len(), 1);
+        assert_eq!(expected_msg, errors[0]);
+    } else {
+        assert!(errors.is_empty());
+    }
 }
 
 #[track_caller]
diff --git a/src/tools/tidy/src/deps.rs b/src/tools/tidy/src/deps.rs
index a9d07c75315..e275d3042cb 100644
--- a/src/tools/tidy/src/deps.rs
+++ b/src/tools/tidy/src/deps.rs
@@ -9,7 +9,7 @@ use build_helper::ci::CiEnv;
 use cargo_metadata::semver::Version;
 use cargo_metadata::{Metadata, Package, PackageId};
 
-use crate::diagnostics::{CheckId, DiagCtx, RunningCheck};
+use crate::diagnostics::{DiagCtx, RunningCheck};
 
 #[path = "../../../bootstrap/src/utils/proc_macro_deps.rs"]
 mod proc_macro_deps;
@@ -664,7 +664,7 @@ const PERMITTED_CRANELIFT_DEPENDENCIES: &[&str] = &[
 /// `root` is path to the directory with the root `Cargo.toml` (for the workspace). `cargo` is path
 /// to the cargo executable.
 pub fn check(root: &Path, cargo: &Path, bless: bool, diag_ctx: DiagCtx) {
-    let mut check = diag_ctx.start_check(CheckId::new("deps").path(root));
+    let mut check = diag_ctx.start_check("deps");
 
     let mut checked_runtime_licenses = false;
 
diff --git a/src/tools/tidy/src/diagnostics.rs b/src/tools/tidy/src/diagnostics.rs
index d7889f925ea..6e95f97d010 100644
--- a/src/tools/tidy/src/diagnostics.rs
+++ b/src/tools/tidy/src/diagnostics.rs
@@ -1,9 +1,9 @@
 use std::collections::HashSet;
-use std::fmt::Display;
+use std::fmt::{Display, Formatter};
 use std::path::{Path, PathBuf};
 use std::sync::{Arc, Mutex};
 
-use termcolor::WriteColor;
+use termcolor::{Color, WriteColor};
 
 /// Collects diagnostics from all tidy steps, and contains shared information
 /// that determines how should message and logs be presented.
@@ -14,26 +14,40 @@ use termcolor::WriteColor;
 pub struct DiagCtx(Arc<Mutex<DiagCtxInner>>);
 
 impl DiagCtx {
-    pub fn new(verbose: bool) -> Self {
+    pub fn new(root_path: &Path, verbose: bool) -> Self {
         Self(Arc::new(Mutex::new(DiagCtxInner {
             running_checks: Default::default(),
             finished_checks: Default::default(),
+            root_path: root_path.to_path_buf(),
             verbose,
         })))
     }
 
     pub fn start_check<Id: Into<CheckId>>(&self, id: Id) -> RunningCheck {
-        let id = id.into();
+        let mut id = id.into();
 
         let mut ctx = self.0.lock().unwrap();
+
+        // Shorten path for shorter diagnostics
+        id.path = match id.path {
+            Some(path) => Some(path.strip_prefix(&ctx.root_path).unwrap_or(&path).to_path_buf()),
+            None => None,
+        };
+
         ctx.start_check(id.clone());
-        RunningCheck { id, bad: false, ctx: self.0.clone() }
+        RunningCheck {
+            id,
+            bad: false,
+            ctx: self.0.clone(),
+            #[cfg(test)]
+            errors: vec![],
+        }
     }
 
-    pub fn into_conclusion(self) -> bool {
-        let ctx = self.0.lock().unwrap();
+    pub fn into_failed_checks(self) -> Vec<FinishedCheck> {
+        let ctx = Arc::into_inner(self.0).unwrap().into_inner().unwrap();
         assert!(ctx.running_checks.is_empty(), "Some checks are still running");
-        ctx.finished_checks.iter().any(|c| c.bad)
+        ctx.finished_checks.into_iter().filter(|c| c.bad).collect()
     }
 }
 
@@ -41,6 +55,7 @@ struct DiagCtxInner {
     running_checks: HashSet<CheckId>,
     finished_checks: HashSet<FinishedCheck>,
     verbose: bool,
+    root_path: PathBuf,
 }
 
 impl DiagCtxInner {
@@ -48,6 +63,7 @@ impl DiagCtxInner {
         if self.has_check_id(&id) {
             panic!("Starting a check named `{id:?}` for the second time");
         }
+
         self.running_checks.insert(id);
     }
 
@@ -57,6 +73,13 @@ impl DiagCtxInner {
             "Finishing check `{:?}` that was not started",
             check.id
         );
+
+        if check.bad {
+            output_message("FAIL", Some(&check.id), Some(COLOR_ERROR));
+        } else if self.verbose {
+            output_message("OK", Some(&check.id), Some(COLOR_SUCCESS));
+        }
+
         self.finished_checks.insert(check);
     }
 
@@ -71,8 +94,8 @@ impl DiagCtxInner {
 /// Identifies a single step
 #[derive(PartialEq, Eq, Hash, Clone, Debug)]
 pub struct CheckId {
-    name: String,
-    path: Option<PathBuf>,
+    pub name: String,
+    pub path: Option<PathBuf>,
 }
 
 impl CheckId {
@@ -91,40 +114,70 @@ impl From<&'static str> for CheckId {
     }
 }
 
+impl Display for CheckId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.name)?;
+        if let Some(path) = &self.path {
+            write!(f, " ({})", path.display())?;
+        }
+        Ok(())
+    }
+}
+
 #[derive(PartialEq, Eq, Hash, Debug)]
-struct FinishedCheck {
+pub struct FinishedCheck {
     id: CheckId,
     bad: bool,
 }
 
+impl FinishedCheck {
+    pub fn id(&self) -> &CheckId {
+        &self.id
+    }
+}
+
 /// Represents a single tidy check, identified by its `name`, running.
 pub struct RunningCheck {
     id: CheckId,
     bad: bool,
     ctx: Arc<Mutex<DiagCtxInner>>,
+    #[cfg(test)]
+    errors: Vec<String>,
 }
 
 impl RunningCheck {
+    /// Creates a new instance of a running check without going through the diag
+    /// context.
+    /// Useful if you want to run some functions from tidy without configuring
+    /// diagnostics.
+    pub fn new_noop() -> Self {
+        let ctx = DiagCtx::new(Path::new(""), false);
+        ctx.start_check("noop")
+    }
+
     /// Immediately output an error and mark the check as failed.
-    pub fn error<T: Display>(&mut self, t: T) {
+    pub fn error<T: Display>(&mut self, msg: T) {
         self.mark_as_bad();
-        tidy_error(&t.to_string()).expect("failed to output error");
+        let msg = msg.to_string();
+        output_message(&msg, Some(&self.id), Some(COLOR_ERROR));
+        #[cfg(test)]
+        self.errors.push(msg);
     }
 
     /// Immediately output a warning.
-    pub fn warning<T: Display>(&mut self, t: T) {
-        eprintln!("WARNING: {t}");
+    pub fn warning<T: Display>(&mut self, msg: T) {
+        output_message(&msg.to_string(), Some(&self.id), Some(COLOR_WARNING));
     }
 
     /// Output an informational message
-    pub fn message<T: Display>(&mut self, t: T) {
-        eprintln!("{t}");
+    pub fn message<T: Display>(&mut self, msg: T) {
+        output_message(&msg.to_string(), Some(&self.id), None);
     }
 
     /// Output a message only if verbose output is enabled.
-    pub fn verbose_msg<T: Display>(&mut self, t: T) {
+    pub fn verbose_msg<T: Display>(&mut self, msg: T) {
         if self.is_verbose_enabled() {
-            self.message(t);
+            self.message(msg);
         }
     }
 
@@ -138,6 +191,11 @@ impl RunningCheck {
         self.ctx.lock().unwrap().verbose
     }
 
+    #[cfg(test)]
+    pub fn get_errors(&self) -> Vec<String> {
+        self.errors.clone()
+    }
+
     fn mark_as_bad(&mut self) {
         self.bad = true;
     }
@@ -149,17 +207,37 @@ impl Drop for RunningCheck {
     }
 }
 
-fn tidy_error(args: &str) -> std::io::Result<()> {
+pub const COLOR_SUCCESS: Color = Color::Green;
+pub const COLOR_ERROR: Color = Color::Red;
+pub const COLOR_WARNING: Color = Color::Yellow;
+
+/// Output a message to stderr.
+/// The message can be optionally scoped to a certain check, and it can also have a certain color.
+pub fn output_message(msg: &str, id: Option<&CheckId>, color: Option<Color>) {
     use std::io::Write;
 
-    use termcolor::{Color, ColorChoice, ColorSpec, StandardStream};
+    use termcolor::{ColorChoice, ColorSpec, StandardStream};
 
-    let mut stderr = StandardStream::stdout(ColorChoice::Auto);
-    stderr.set_color(ColorSpec::new().set_fg(Some(Color::Red)))?;
+    let mut stderr = StandardStream::stderr(ColorChoice::Auto);
+    if let Some(color) = &color {
+        stderr.set_color(ColorSpec::new().set_fg(Some(*color))).unwrap();
+    }
 
-    write!(&mut stderr, "tidy error")?;
-    stderr.set_color(&ColorSpec::new())?;
+    match id {
+        Some(id) => {
+            write!(&mut stderr, "tidy [{}", id.name).unwrap();
+            if let Some(path) = &id.path {
+                write!(&mut stderr, " ({})", path.display()).unwrap();
+            }
+            write!(&mut stderr, "]").unwrap();
+        }
+        None => {
+            write!(&mut stderr, "tidy").unwrap();
+        }
+    }
+    if color.is_some() {
+        stderr.set_color(&ColorSpec::new()).unwrap();
+    }
 
-    writeln!(&mut stderr, ": {args}")?;
-    Ok(())
+    writeln!(&mut stderr, ": {msg}").unwrap();
 }
diff --git a/src/tools/tidy/src/extdeps.rs b/src/tools/tidy/src/extdeps.rs
index 21612980007..f75de13b45c 100644
--- a/src/tools/tidy/src/extdeps.rs
+++ b/src/tools/tidy/src/extdeps.rs
@@ -4,7 +4,7 @@ use std::fs;
 use std::path::Path;
 
 use crate::deps::WorkspaceInfo;
-use crate::diagnostics::{CheckId, DiagCtx};
+use crate::diagnostics::DiagCtx;
 
 /// List of allowed sources for packages.
 const ALLOWED_SOURCES: &[&str] = &[
@@ -16,7 +16,7 @@ const ALLOWED_SOURCES: &[&str] = &[
 /// Checks for external package sources. `root` is the path to the directory that contains the
 /// workspace `Cargo.toml`.
 pub fn check(root: &Path, diag_ctx: DiagCtx) {
-    let mut check = diag_ctx.start_check(CheckId::new("extdeps").path(root));
+    let mut check = diag_ctx.start_check("extdeps");
 
     for &WorkspaceInfo { path, submodules, .. } in crate::deps::WORKSPACES {
         if crate::deps::has_missing_submodule(root, submodules) {
diff --git a/src/tools/tidy/src/filenames.rs b/src/tools/tidy/src/filenames.rs
index d52a8229738..835cbefbf69 100644
--- a/src/tools/tidy/src/filenames.rs
+++ b/src/tools/tidy/src/filenames.rs
@@ -10,10 +10,10 @@
 use std::path::Path;
 use std::process::Command;
 
-use crate::diagnostics::{CheckId, DiagCtx};
+use crate::diagnostics::DiagCtx;
 
 pub fn check(root_path: &Path, diag_ctx: DiagCtx) {
-    let mut check = diag_ctx.start_check(CheckId::new("filenames").path(root_path));
+    let mut check = diag_ctx.start_check("filenames");
     let stat_output = Command::new("git")
         .arg("-C")
         .arg(root_path)
diff --git a/src/tools/tidy/src/main.rs b/src/tools/tidy/src/main.rs
index 73ff183868a..93bc1611199 100644
--- a/src/tools/tidy/src/main.rs
+++ b/src/tools/tidy/src/main.rs
@@ -11,7 +11,7 @@ use std::str::FromStr;
 use std::thread::{self, ScopedJoinHandle, scope};
 use std::{env, process};
 
-use tidy::diagnostics::DiagCtx;
+use tidy::diagnostics::{COLOR_ERROR, COLOR_SUCCESS, DiagCtx, output_message};
 use tidy::*;
 
 fn main() {
@@ -50,7 +50,7 @@ fn main() {
     let extra_checks =
         cfg_args.iter().find(|s| s.starts_with("--extra-checks=")).map(String::as_str);
 
-    let diag_ctx = DiagCtx::new(verbose);
+    let diag_ctx = DiagCtx::new(&root_path, verbose);
     let ci_info = CiInfo::new(diag_ctx.clone());
 
     let drain_handles = |handles: &mut VecDeque<ScopedJoinHandle<'_, ()>>| {
@@ -175,8 +175,22 @@ fn main() {
         );
     });
 
-    if diag_ctx.into_conclusion() {
-        eprintln!("some tidy checks failed");
+    let failed_checks = diag_ctx.into_failed_checks();
+    if !failed_checks.is_empty() {
+        let mut failed: Vec<String> =
+            failed_checks.into_iter().map(|c| c.id().to_string()).collect();
+        failed.sort();
+        output_message(
+            &format!(
+                "The following check{} failed: {}",
+                if failed.len() > 1 { "s" } else { "" },
+                failed.join(", ")
+            ),
+            None,
+            Some(COLOR_ERROR),
+        );
         process::exit(1);
+    } else {
+        output_message("All tidy checks succeeded", None, Some(COLOR_SUCCESS));
     }
 }
diff --git a/src/tools/tidy/src/tests_placement.rs b/src/tools/tidy/src/tests_placement.rs
index e978e1c453c..8ba8cf552bd 100644
--- a/src/tools/tidy/src/tests_placement.rs
+++ b/src/tools/tidy/src/tests_placement.rs
@@ -1,12 +1,12 @@
 use std::path::Path;
 
-use crate::diagnostics::{CheckId, DiagCtx};
+use crate::diagnostics::DiagCtx;
 
 const FORBIDDEN_PATH: &str = "src/test";
 const ALLOWED_PATH: &str = "tests";
 
 pub fn check(root_path: &Path, diag_ctx: DiagCtx) {
-    let mut check = diag_ctx.start_check(CheckId::new("tests-placement").path(root_path));
+    let mut check = diag_ctx.start_check("tests_placement");
 
     if root_path.join(FORBIDDEN_PATH).exists() {
         check.error(format!(
diff --git a/src/tools/tidy/src/triagebot.rs b/src/tools/tidy/src/triagebot.rs
index 9eccef29f2e..41d61dcd141 100644
--- a/src/tools/tidy/src/triagebot.rs
+++ b/src/tools/tidy/src/triagebot.rs
@@ -4,10 +4,10 @@ use std::path::Path;
 
 use toml::Value;
 
-use crate::diagnostics::{CheckId, DiagCtx};
+use crate::diagnostics::DiagCtx;
 
 pub fn check(path: &Path, diag_ctx: DiagCtx) {
-    let mut check = diag_ctx.start_check(CheckId::new("triagebot").path(path));
+    let mut check = diag_ctx.start_check("triagebot");
     let triagebot_path = path.join("triagebot.toml");
 
     // This check is mostly to catch broken path filters *within* `triagebot.toml`, and not enforce
diff --git a/src/tools/unstable-book-gen/src/main.rs b/src/tools/unstable-book-gen/src/main.rs
index a7c6173d88c..16550f83003 100644
--- a/src/tools/unstable-book-gen/src/main.rs
+++ b/src/tools/unstable-book-gen/src/main.rs
@@ -5,6 +5,7 @@ use std::env;
 use std::fs::{self, write};
 use std::path::Path;
 
+use tidy::diagnostics::RunningCheck;
 use tidy::features::{Features, collect_env_vars, collect_lang_features, collect_lib_features};
 use tidy::t;
 use tidy::unstable_book::{
@@ -122,7 +123,7 @@ fn main() {
     let src_path = Path::new(&src_path_str);
     let dest_path = Path::new(&dest_path_str);
 
-    let lang_features = collect_lang_features(compiler_path, &mut false);
+    let lang_features = collect_lang_features(compiler_path, &mut RunningCheck::new_noop());
     let lib_features = collect_lib_features(library_path)
         .into_iter()
         .filter(|&(ref name, _)| !lang_features.contains_key(name))