From 8b5055d5106da402f9a132f0ed21571ef98b8ac2 Mon Sep 17 00:00:00 2001
From: Matan Kushner <hello@matchai.me>
Date: Thu, 9 May 2019 23:51:50 -0400
Subject: [PATCH] Parallelize prompt modules (#46)

---
 Cargo.lock                |  1 +
 Cargo.toml                |  1 +
 benches/my_benchmark.rs   | 27 ++++++++++++++++++++++++++-
 src/context.rs            |  8 +++++---
 src/modules/directory.rs  |  3 +--
 src/modules/git_branch.rs |  8 +++-----
 src/print.rs              |  3 ++-
 7 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index f27c70f0..0fa8f511 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -689,6 +689,7 @@ dependencies = [
  "criterion 0.2.11 (registry+https://github.com/rust-lang/crates.io-index)",
  "dirs 1.0.5 (registry+https://github.com/rust-lang/crates.io-index)",
  "git2 0.8.0 (registry+https://github.com/rust-lang/crates.io-index)",
+ "rayon 1.0.3 (registry+https://github.com/rust-lang/crates.io-index)",
  "serde_json 1.0.39 (registry+https://github.com/rust-lang/crates.io-index)",
  "tempfile 3.0.7 (registry+https://github.com/rust-lang/crates.io-index)",
  "toml 0.5.0 (registry+https://github.com/rust-lang/crates.io-index)",
diff --git a/Cargo.toml b/Cargo.toml
index ece4734e..016c1f3f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -11,6 +11,7 @@ dirs = "1.0.5"
 git2 = "0.8.0"
 toml = "0.5.0"
 serde_json = "1.0.39"
+rayon = "1.0.3"
 
 [dev-dependencies]
 tempfile = "3.0.7"
diff --git a/benches/my_benchmark.rs b/benches/my_benchmark.rs
index e16c8cab..be6898bb 100644
--- a/benches/my_benchmark.rs
+++ b/benches/my_benchmark.rs
@@ -6,6 +6,8 @@ use criterion::Criterion;
 use clap::{App, Arg};
 use starship::context::Context;
 use starship::modules;
+use std::fs;
+use tempfile::TempDir;
 
 fn char_segment(c: &mut Criterion) {
     let args = App::new("starship")
@@ -40,5 +42,28 @@ fn line_break_segment(c: &mut Criterion) {
     });
 }
 
-criterion_group!(benches, dir_segment, char_segment, line_break_segment);
+fn git_branch_segment(c: &mut Criterion) {
+    let tmp_dir = TempDir::new().unwrap();
+    let repo_dir = tmp_dir.path().join("rocket-controls");
+    fs::create_dir(&repo_dir).unwrap();
+
+    git2::Repository::init(&repo_dir).unwrap();
+
+    let args = App::new("starship")
+        .arg(Arg::with_name("status_code"))
+        .get_matches_from(vec!["starship", "0"]);
+    let context = Context::new_with_dir(args, "~");
+
+    c.bench_function("git_branch segment", move |b| {
+        b.iter(|| modules::handle("git_branch", &context))
+    });
+}
+
+criterion_group!(
+    benches,
+    char_segment,
+    dir_segment,
+    line_break_segment,
+    git_branch_segment
+);
 criterion_main!(benches);
diff --git a/src/context.rs b/src/context.rs
index 6baa1953..035eec90 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -8,7 +8,7 @@ pub struct Context<'a> {
     pub current_dir: PathBuf,
     pub dir_files: Vec<PathBuf>,
     pub arguments: ArgMatches<'a>,
-    pub repository: Option<Repository>,
+    pub repo_root: Option<PathBuf>,
 }
 
 impl<'a> Context<'a> {
@@ -36,13 +36,15 @@ impl<'a> Context<'a> {
             .map(|entry| entry.path())
             .collect::<Vec<PathBuf>>();
 
-        let repository: Option<Repository> = Repository::discover(&current_dir).ok();
+        let repo_root: Option<PathBuf> = Repository::discover(&current_dir)
+            .ok()
+            .and_then(|repo| repo.workdir().map(|repo| repo.to_path_buf()));
 
         Context {
             arguments,
             current_dir,
             dir_files,
-            repository,
+            repo_root,
         }
     }
 
diff --git a/src/modules/directory.rs b/src/modules/directory.rs
index a85287a6..e0f4a9e1 100644
--- a/src/modules/directory.rs
+++ b/src/modules/directory.rs
@@ -23,9 +23,8 @@ pub fn segment(context: &Context) -> Option<Module> {
     let current_dir = &context.current_dir;
 
     let dir_string;
-    if let Some(repo) = &context.repository {
+    if let Some(repo_root) = &context.repo_root {
         // Contract the path to the git repo root
-        let repo_root = repo.workdir().unwrap();
         let repo_folder_name = repo_root.file_name().unwrap().to_str().unwrap();
 
         dir_string = contract_path(&current_dir, repo_root, repo_folder_name);
diff --git a/src/modules/git_branch.rs b/src/modules/git_branch.rs
index a89e77f3..22b934da 100644
--- a/src/modules/git_branch.rs
+++ b/src/modules/git_branch.rs
@@ -7,12 +7,10 @@ use super::{Context, Module};
 ///
 /// Will display the branch name if the current directory is a git repo
 pub fn segment(context: &Context) -> Option<Module> {
-    if context.repository.is_none() {
-        return None;
-    }
+    let repo_root = context.repo_root.as_ref()?;
+    let repository = Repository::open(repo_root).ok()?;
 
-    let repository = context.repository.as_ref().unwrap();
-    match get_current_branch(repository) {
+    match get_current_branch(&repository) {
         Ok(branch_name) => {
             const GIT_BRANCH_CHAR: &str = " ";
             let segment_color = Color::Purple.bold();
diff --git a/src/print.rs b/src/print.rs
index 780d9099..3020f9de 100644
--- a/src/print.rs
+++ b/src/print.rs
@@ -1,4 +1,5 @@
 use clap::ArgMatches;
+use rayon::prelude::*;
 use std::io::{self, Write};
 
 use crate::context::Context;
@@ -29,7 +30,7 @@ pub fn prompt(args: ArgMatches) {
     writeln!(handle).unwrap();
 
     let modules = prompt_order
-        .iter()
+        .par_iter()
         .map(|module| modules::handle(module, &context)) // Compute modules
         .flatten()
         .collect::<Vec<Module>>(); // Remove segments set to `None`