feat: Add timings subcommand (#1629)

* feat: Add computational duration to all computed modules

This also means that in case we do some computations and these end up empty, we submit an empty module

* feat: Add timings subcommand

This outputs the timings of all computed modules, sorted by the duration it took to compute the module.

Useful for debugging why the prompt takes so long.

* feat: Add timings to explain output

* fix: Ensure that even empty custom modules get timings

* format main.rs

* feat: Only show interesting timings

* fix(tests): Change tests to look for empty string instead of None

* Use proper wording in timings help

* Revert "fix(tests): Change tests to look for empty string instead of None"

This reverts commit aca5bd1b03c48e1dee1b7ca91d66e2bda2d5a97c.

* fix(tests): Returning None in case the module produced an empty string

* fix: Ensure that linebreaks (and space) make a module not-empty

* Make cargo clippy happy

* Make Module.duration a proper Duration

* Only return a module if we would report it

* Change to cleaner way to return None for empty modules

* Avoid unnecessary module creation

* Simplify a string comparison

* Add timings to trace

Co-authored-by: Thomas O'Donnell <andytom@users.noreply.github.com>

Co-authored-by: Thomas O'Donnell <andytom@users.noreply.github.com>

7 files changed, 316 insertions, 173 deletions

diff --git a/src/main.rs b/src/main.rs
index 27876b9c8..3afcf6b3d 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -71,87 +71,89 @@ fn main() {
         .long("print-full-init")
         .help("Print the main initialization script (as opposed to the init stub)");
 
-    let mut app =
-        App::new("starship")
-            .about("The cross-shell prompt for astronauts. ☄🌌️")
-            // pull the version number from Cargo.toml
-            .version(crate_version!())
-            // pull the authors from Cargo.toml
-            .author(crate_authors!())
-            .after_help("https://github.com/starship/starship")
-            .setting(AppSettings::SubcommandRequiredElseHelp)
-            .subcommand(
-                SubCommand::with_name("init")
-                    .about("Prints the shell function used to execute starship")
-                    .arg(&shell_arg)
-                    .arg(&init_scripts_arg),
-            )
-            .subcommand(
-                SubCommand::with_name("prompt")
-                    .about("Prints the full starship prompt")
-                    .arg(&status_code_arg)
-                    .arg(&path_arg)
-                    .arg(&cmd_duration_arg)
-                    .arg(&keymap_arg)
-                    .arg(&jobs_arg),
-            )
-            .subcommand(
-                SubCommand::with_name("module")
-                    .about("Prints a specific prompt module")
-                    .arg(
-                        Arg::with_name("name")
-                            .help("The name of the module to be printed")
-                            .required(true)
-                            .required_unless("list"),
-                    )
-                    .arg(
-                        Arg::with_name("list")
-                            .short("l")
-                            .long("list")
-                            .help("List out all supported modules"),
-                    )
-                    .arg(&status_code_arg)
-                    .arg(&path_arg)
-                    .arg(&cmd_duration_arg)
-                    .arg(&keymap_arg)
-                    .arg(&jobs_arg),
-            )
-            .subcommand(
-                SubCommand::with_name("config")
-                    .alias("configure")
-                    .about("Edit the starship configuration")
-                    .arg(
-                        Arg::with_name("name")
-                            .help("Configuration key to edit")
-                            .required(false)
-                            .requires("value"),
-                    )
-                    .arg(Arg::with_name("value").help("Value to place into that key")),
-            )
-            .subcommand(SubCommand::with_name("bug-report").about(
+    let mut app = App::new("starship")
+        .about("The cross-shell prompt for astronauts. ☄🌌️")
+        // pull the version number from Cargo.toml
+        .version(crate_version!())
+        // pull the authors from Cargo.toml
+        .author(crate_authors!())
+        .after_help("https://github.com/starship/starship")
+        .setting(AppSettings::SubcommandRequiredElseHelp)
+        .subcommand(
+            SubCommand::with_name("init")
+                .about("Prints the shell function used to execute starship")
+                .arg(&shell_arg)
+                .arg(&init_scripts_arg),
+        )
+        .subcommand(
+            SubCommand::with_name("prompt")
+                .about("Prints the full starship prompt")
+                .arg(&status_code_arg)
+                .arg(&path_arg)
+                .arg(&cmd_duration_arg)
+                .arg(&keymap_arg)
+                .arg(&jobs_arg),
+        )
+        .subcommand(
+            SubCommand::with_name("module")
+                .about("Prints a specific prompt module")
+                .arg(
+                    Arg::with_name("name")
+                        .help("The name of the module to be printed")
+                        .required(true)
+                        .required_unless("list"),
+                )
+                .arg(
+                    Arg::with_name("list")
+                        .short("l")
+                        .long("list")
+                        .help("List out all supported modules"),
+                )
+                .arg(&status_code_arg)
+                .arg(&path_arg)
+                .arg(&cmd_duration_arg)
+                .arg(&keymap_arg)
+                .arg(&jobs_arg),
+        )
+        .subcommand(
+            SubCommand::with_name("config")
+                .alias("configure")
+                .about("Edit the starship configuration")
+                .arg(
+                    Arg::with_name("name")
+                        .help("Configuration key to edit")
+                        .required(false)
+                        .requires("value"),
+                )
+                .arg(Arg::with_name("value").help("Value to place into that key")),
+        )
+        .subcommand(
+            SubCommand::with_name("bug-report").about(
                 "Create a pre-populated GitHub issue with information about your configuration",
-            ))
-            .subcommand(
-                SubCommand::with_name("time")
-                    .about("Prints time in milliseconds")
-                    .settings(&[AppSettings::Hidden]),
-            )
-            .subcommand(
-                SubCommand::with_name("explain").about("Explains the currently showing modules"),
-            )
-            .subcommand(
-                SubCommand::with_name("completions")
-                    .about("Generate starship shell completions for your shell to stdout")
-                    .arg(
-                        Arg::with_name("shell")
-                            .takes_value(true)
-                            .possible_values(&Shell::variants())
-                            .help("the shell to generate completions for")
-                            .value_name("SHELL")
-                            .required(true)
-                            .env("STARSHIP_SHELL"),
-                    ),
-            );
+            ),
+        )
+        .subcommand(
+            SubCommand::with_name("time")
+                .about("Prints time in milliseconds")
+                .settings(&[AppSettings::Hidden]),
+        )
+        .subcommand(
+            SubCommand::with_name("explain").about("Explains the currently showing modules"),
+        )
+        .subcommand(SubCommand::with_name("timings").about("Prints timings of all active modules"))
+        .subcommand(
+            SubCommand::with_name("completions")
+                .about("Generate starship shell completions for your shell to stdout")
+                .arg(
+                    Arg::with_name("shell")
+                        .takes_value(true)
+                        .possible_values(&Shell::variants())
+                        .help("the shell to generate completions for")
+                        .value_name("SHELL")
+                        .required(true)
+                        .env("STARSHIP_SHELL"),
+                ),
+        );
 
     let matches = app.clone().get_matches();
 
@@ -197,6 +199,7 @@ fn main() {
             }
         }
         ("explain", Some(sub_m)) => print::explain(sub_m.clone()),
+        ("timings", Some(sub_m)) => print::timings(sub_m.clone()),
         ("completions", Some(sub_m)) => {
             let shell: Shell = sub_m
                 .value_of("shell")
diff --git a/src/module.rs b/src/module.rs
index 5ffb0cb69..768071c9c 100644
--- a/src/module.rs
+++ b/src/module.rs
@@ -3,6 +3,7 @@ use crate::segment::Segment;
 use crate::utils::wrap_colorseq_for_shell;
 use ansi_term::{ANSIString, ANSIStrings};
 use std::fmt;
+use std::time::Duration;
 
 // List of all modules
 // Keep these ordered alphabetically.
@@ -73,6 +74,9 @@ pub struct Module<'a> {
 
     /// The collection of segments that compose this module.
     pub segments: Vec<Segment>,
+
+    /// the time it took to compute this module
+    pub duration: Duration,
 }
 
 impl<'a> Module<'a> {
@@ -83,6 +87,7 @@ impl<'a> Module<'a> {
             name: name.to_string(),
             description: desc.to_string(),
             segments: Vec::new(),
+            duration: Duration::default(),
         }
     }
 
@@ -105,7 +110,8 @@ impl<'a> Module<'a> {
     pub fn is_empty(&self) -> bool {
         self.segments
             .iter()
-            .all(|segment| segment.value.trim().is_empty())
+            // no trim: if we add spaces/linebreaks it's not "empty" as we change the final output
+            .all(|segment| segment.value.is_empty())
     }
 
     /// Get values of the module's segments
@@ -167,6 +173,7 @@ mod tests {
             name: name.to_string(),
             description: desc.to_string(),
             segments: Vec::new(),
+            duration: Duration::default(),
         };
 
         assert!(module.is_empty());
@@ -181,8 +188,39 @@ mod tests {
             name: name.to_string(),
             description: desc.to_string(),
             segments: vec![Segment::new(None, "")],
+            duration: Duration::default(),
         };
 
         assert!(module.is_empty());
     }
+
+    #[test]
+    fn test_module_is_not_empty_with_linebreak_only() {
+        let name = "unit_test";
+        let desc = "This is a unit test";
+        let module = Module {
+            config: None,
+            name: name.to_string(),
+            description: desc.to_string(),
+            segments: vec![Segment::new(None, "\n")],
+            duration: Duration::default(),
+        };
+
+        assert!(!module.is_empty());
+    }
+
+    #[test]
+    fn test_module_is_not_empty_with_space_only() {
+        let name = "unit_test";
+        let desc = "This is a unit test";
+        let module = Module {
+            config: None,
+            name: name.to_string(),
+            description: desc.to_string(),
+            segments: vec![Segment::new(None, " ")],
+            duration: Duration::default(),
+        };
+
+        assert!(!module.is_empty());
+    }
 }
diff --git a/src/modules/custom.rs b/src/modules/custom.rs
index bb28b49af..d72a874b4 100644
--- a/src/modules/custom.rs
+++ b/src/modules/custom.rs
@@ -1,5 +1,6 @@
 use std::io::Write;
 use std::process::{Command, Output, Stdio};
+use std::time::Instant;
 
 use super::{Context, Module, RootModuleConfig};
 
@@ -13,6 +14,7 @@ use crate::{configs::custom::CustomConfig, formatter::StringFormatter};
 ///
 /// Finally, the content of the module itself is also set by a command.
 pub fn module<'a>(name: &str, context: &'a Context) -> Option<Module<'a>> {
+    let start: Instant = Instant::now();
     let toml_config = context.config.get_custom_module_config(name).expect(
         "modules::custom::module should only be called after ensuring that the module exists",
     );
@@ -47,37 +49,36 @@ pub fn module<'a>(name: &str, context: &'a Context) -> Option<Module<'a>> {
     let output = exec_command(config.command, &config.shell.0)?;
 
     let trimmed = output.trim();
-    if trimmed.is_empty() {
-        return None;
+    if !trimmed.is_empty() {
+        let parsed = StringFormatter::new(config.format).and_then(|formatter| {
+            formatter
+                .map_meta(|var, _| match var {
+                    "symbol" => Some(config.symbol),
+                    _ => None,
+                })
+                .map_style(|variable| match variable {
+                    "style" => Some(Ok(config.style)),
+                    _ => None,
+                })
+                .map(|variable| match variable {
+                    // This may result in multiple calls to `get_module_version` when a user have
+                    // multiple `$version` variables defined in `format`.
+                    "output" => Some(Ok(trimmed)),
+                    _ => None,
+                })
+                .parse(None)
+        });
+
+        match parsed {
+            Ok(segments) => module.set_segments(segments),
+            Err(error) => {
+                log::warn!("Error in module `custom.{}`:\n{}", name, error);
+            }
+        };
     }
-
-    let parsed = StringFormatter::new(config.format).and_then(|formatter| {
-        formatter
-            .map_meta(|var, _| match var {
-                "symbol" => Some(config.symbol),
-                _ => None,
-            })
-            .map_style(|variable| match variable {
-                "style" => Some(Ok(config.style)),
-                _ => None,
-            })
-            .map(|variable| match variable {
-                // This may result in multiple calls to `get_module_version` when a user have
-                // multiple `$version` variables defined in `format`.
-                "output" => Some(Ok(trimmed)),
-                _ => None,
-            })
-            .parse(None)
-    });
-
-    module.set_segments(match parsed {
-        Ok(segments) => segments,
-        Err(error) => {
-            log::warn!("Error in module `custom.{}`:\n{}", name, error);
-            return None;
-        }
-    });
-
+    let elapsed = start.elapsed();
+    log::trace!("Took {:?} to compute custom module {:?}", elapsed, name);
+    module.duration = elapsed;
     Some(module)
 }
 
diff --git a/src/modules/line_break.rs b/src/modules/line_break.rs
index e485458d4..944a1c222 100644
--- a/src/modules/line_break.rs
+++ b/src/modules/line_break.rs
@@ -11,3 +11,19 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
 
     Some(module)
 }
+
+#[cfg(test)]
+mod test {
+    use std::io;
+
+    use crate::test::ModuleRenderer;
+
+    #[test]
+    fn produces_result() -> io::Result<()> {
+        let expected = Some(String::from("\n"));
+        let actual = ModuleRenderer::new("line_break").collect();
+        assert_eq!(expected, actual);
+
+        Ok(())
+    }
+}
diff --git a/src/modules/mod.rs b/src/modules/mod.rs
index a9af8cadb..9098a746a 100644
--- a/src/modules/mod.rs
+++ b/src/modules/mod.rs
@@ -55,64 +55,82 @@ mod battery;
 use crate::config::RootModuleConfig;
 use crate::context::{Context, Shell};
 use crate::module::Module;
+use std::time::Instant;
 
 pub fn handle<'a>(module: &str, context: &'a Context) -> Option<Module<'a>> {
-    match module {
-        // Keep these ordered alphabetically.
-        // Default ordering is handled in configs/starship_root.rs
-        "aws" => aws::module(context),
-        #[cfg(feature = "battery")]
-        "battery" => battery::module(context),
-        "character" => character::module(context),
-        "cmake" => cmake::module(context),
-        "cmd_duration" => cmd_duration::module(context),
-        "conda" => conda::module(context),
-        "dart" => dart::module(context),
-        "directory" => directory::module(context),
-        "docker_context" => docker_context::module(context),
-        "dotnet" => dotnet::module(context),
-        "elixir" => elixir::module(context),
-        "elm" => elm::module(context),
-        "erlang" => erlang::module(context),
-        "env_var" => env_var::module(context),
-        "gcloud" => gcloud::module(context),
-        "git_branch" => git_branch::module(context),
-        "git_commit" => git_commit::module(context),
-        "git_state" => git_state::module(context),
-        "git_status" => git_status::module(context),
-        "golang" => golang::module(context),
-        "helm" => helm::module(context),
-        "hg_branch" => hg_branch::module(context),
-        "hostname" => hostname::module(context),
-        "java" => java::module(context),
-        "jobs" => jobs::module(context),
-        "julia" => julia::module(context),
-        "kubernetes" => kubernetes::module(context),
-        "line_break" => line_break::module(context),
-        "memory_usage" => memory_usage::module(context),
-        "nim" => nim::module(context),
-        "nix_shell" => nix_shell::module(context),
-        "nodejs" => nodejs::module(context),
-        "ocaml" => ocaml::module(context),
-        "package" => package::module(context),
-        "perl" => perl::module(context),
-        "php" => php::module(context),
-        "purescript" => purescript::module(context),
-        "python" => python::module(context),
-        "ruby" => ruby::module(context),
-        "rust" => rust::module(context),
-        "shlvl" => shlvl::module(context),
-        "singularity" => singularity::module(context),
-        "swift" => swift::module(context),
-        "terraform" => terraform::module(context),
-        "time" => time::module(context),
-        "crystal" => crystal::module(context),
-        "username" => username::module(context),
-        "zig" => zig::module(context),
-        _ => {
-            eprintln!("Error: Unknown module {}. Use starship module --list to list out all supported modules.", module);
-            None
+    let start: Instant = Instant::now();
+
+    let mut m: Option<Module> = {
+        match module {
+            // Keep these ordered alphabetically.
+            // Default ordering is handled in configs/starship_root.rs
+            "aws" => aws::module(context),
+            #[cfg(feature = "battery")]
+            "battery" => battery::module(context),
+            "character" => character::module(context),
+            "cmake" => cmake::module(context),
+            "cmd_duration" => cmd_duration::module(context),
+            "conda" => conda::module(context),
+            "dart" => dart::module(context),
+            "directory" => directory::module(context),
+            "docker_context" => docker_context::module(context),
+            "dotnet" => dotnet::module(context),
+            "elixir" => elixir::module(context),
+            "elm" => elm::module(context),
+            "erlang" => erlang::module(context),
+            "env_var" => env_var::module(context),
+            "gcloud" => gcloud::module(context),
+            "git_branch" => git_branch::module(context),
+            "git_commit" => git_commit::module(context),
+            "git_state" => git_state::module(context),
+            "git_status" => git_status::module(context),
+            "golang" => golang::module(context),
+            "helm" => helm::module(context),
+            "hg_branch" => hg_branch::module(context),
+            "hostname" => hostname::module(context),
+            "java" => java::module(context),
+            "jobs" => jobs::module(context),
+            "julia" => julia::module(context),
+            "kubernetes" => kubernetes::module(context),
+            "line_break" => line_break::module(context),
+            "memory_usage" => memory_usage::module(context),
+            "nim" => nim::module(context),
+            "nix_shell" => nix_shell::module(context),
+            "nodejs" => nodejs::module(context),
+            "ocaml" => ocaml::module(context),
+            "package" => package::module(context),
+            "perl" => perl::module(context),
+            "php" => php::module(context),
+            "purescript" => purescript::module(context),
+            "python" => python::module(context),
+            "ruby" => ruby::module(context),
+            "rust" => rust::module(context),
+            "shlvl" => shlvl::module(context),
+            "singularity" => singularity::module(context),
+            "swift" => swift::module(context),
+            "terraform" => terraform::module(context),
+            "time" => time::module(context),
+            "crystal" => crystal::module(context),
+            "username" => username::module(context),
+            "zig" => zig::module(context),
+            _ => {
+                eprintln!("Error: Unknown module {}. Use starship module --list to list out all supported modules.", module);
+                None
+            }
         }
+    };
+
+    let elapsed = start.elapsed();
+    log::trace!("Took {:?} to compute module {:?}", elapsed, module);
+    if elapsed.as_millis() < 1 {
+        // If we take less than 1ms to compute a None, then we will not return a module at all
+        // if we have a module: default duration is 0 so no need to change it
+        m
+    } else {
+        // if we took more than 1ms we want to report that and so--in case we have None currently--
+        // need to create an empty module just to hold the duration for that case
+        m.get_or_insert_with(|| context.new_module(module)).duration = elapsed;
+        m
     }
 }
 
diff --git a/src/print.rs b/src/print.rs
index ce45e3d04..84c0b8f13 100644
--- a/src/print.rs
+++ b/src/print.rs
@@ -4,6 +4,7 @@ use rayon::prelude::*;
 use std::collections::BTreeSet;
 use std::fmt::{self, Debug, Write as FmtWrite};
 use std::io::{self, Write};
+use std::time::Duration;
 use unicode_segmentation::UnicodeSegmentation;
 use unicode_width::UnicodeWidthChar;
 
@@ -90,6 +91,52 @@ pub fn get_module(module_name: &str, context: Context) -> Option<String> {
     modules::handle(module_name, &context).map(|m| m.to_string())
 }
 
+pub fn timings(args: ArgMatches) {
+    let context = Context::new(args);
+
+    struct ModuleTiming {
+        name: String,
+        name_len: usize,
+        value: String,
+        duration: Duration,
+        duration_len: usize,
+    }
+
+    let mut modules = compute_modules(&context)
+        .iter()
+        .filter(|module| !module.is_empty() || module.duration.as_millis() > 0)
+        .map(|module| ModuleTiming {
+            name: String::from(module.get_name().as_str()),
+            name_len: better_width(module.get_name().as_str()),
+            value: ansi_term::ANSIStrings(&module.ansi_strings())
+                .to_string()
+                .replace('\n', "\\n"),
+            duration: module.duration,
+            duration_len: better_width(format_duration(&module.duration).as_str()),
+        })
+        .collect::<Vec<ModuleTiming>>();
+
+    modules.sort_by(|a, b| b.duration.cmp(&a.duration));
+
+    let max_name_width = modules.iter().map(|i| i.name_len).max().unwrap_or(0);
+    let max_duration_width = modules.iter().map(|i| i.duration_len).max().unwrap_or(0);
+
+    println!("\n Here are the timings of modules in your prompt (>=1ms or output):");
+
+    // for now we do not expect a wrap around at the end... famous last words
+    // Overall a line looks like this: " {module name}  -  {duration}  -  {module value}".
+    for timing in &modules {
+        println!(
+            " {}{}  -  {}{}  -   {}",
+            timing.name,
+            " ".repeat(max_name_width - (timing.name_len)),
+            " ".repeat(max_duration_width - (timing.duration_len)),
+            format_duration(&timing.duration),
+            timing.value
+        );
+    }
+}
+
 pub fn explain(args: ArgMatches) {
     let context = Context::new(args);
 
@@ -97,6 +144,7 @@ pub fn explain(args: ArgMatches) {
         value: String,
         value_len: usize,
         desc: String,
+        duration: String,
     }
 
     let dont_print = vec!["line_break"];
@@ -104,21 +152,25 @@ pub fn explain(args: ArgMatches) {
     let modules = compute_modules(&context)
         .into_iter()
         .filter(|module| !dont_print.contains(&module.get_name().as_str()))
+        // this contains empty modules which should not print
+        .filter(|module| !module.is_empty())
         .map(|module| {
             let value = module.get_segments().join("");
             ModuleInfo {
                 value: ansi_term::ANSIStrings(&module.ansi_strings()).to_string(),
-                value_len: better_width(value.as_str()),
+                value_len: better_width(value.as_str())
+                    + better_width(format_duration(&module.duration).as_str()),
                 desc: module.get_description().to_owned(),
+                duration: format_duration(&module.duration),
             }
         })
         .collect::<Vec<ModuleInfo>>();
 
     let max_module_width = modules.iter().map(|i| i.value_len).max().unwrap_or(0);
 
-    // In addition to the module width itself there are also 6 padding characters in each line.
-    // Overall a line looks like this: " {module name}  -  {description}".
-    const PADDING_WIDTH: usize = 6;
+    // In addition to the module width itself there are also 9 padding characters in each line.
+    // Overall a line looks like this: " {module value} ({xxxms})  -  {description}".
+    const PADDING_WIDTH: usize = 9;
 
     let desc_width = term_size::dimensions()
         .map(|(w, _)| w)
@@ -133,9 +185,10 @@ pub fn explain(args: ArgMatches) {
             let mut escaping = false;
             // Print info
             print!(
-                " {}{}  -  ",
+                " {} ({}){}  -  ",
                 info.value,
-                " ".repeat(max_module_width - info.value_len)
+                info.duration,
+                " ".repeat(max_module_width - (info.value_len))
             );
             for g in info.desc.graphemes(true) {
                 // Handle ANSI escape sequnces
@@ -296,6 +349,15 @@ fn better_width(s: &str) -> usize {
     s.graphemes(true).map(grapheme_width).sum()
 }
 
+pub fn format_duration(duration: &Duration) -> String {
+    let milis = duration.as_millis();
+    if milis == 0 {
+        "<1ms".to_string()
+    } else {
+        format!("{:?}ms", &milis)
+    }
+}
+
 // Assume that graphemes have width of the first character in the grapheme
 fn grapheme_width(g: &str) -> usize {
     g.chars().next().and_then(|i| i.width()).unwrap_or(0)
diff --git a/src/test/mod.rs b/src/test/mod.rs
index 71402a0ac..ae19f6f5a 100644
--- a/src/test/mod.rs
+++ b/src/test/mod.rs
@@ -84,7 +84,12 @@ impl<'a> ModuleRenderer<'a> {
 
     /// Renders the module returning its output
     pub fn collect(self) -> Option<String> {
-        crate::print::get_module(self.name, self.context)
+        let ret = crate::print::get_module(self.name, self.context);
+        // all tests rely on the fact that an empty module produces None as output as the
+        // convention was that there would be no module but None. This is nowadays not anymore
+        // the case (to get durations for all modules). So here we make it so, that an empty
+        // module returns None in the tests...
+        ret.filter(|s| s != "")
     }
 }