add git info to the prompt

4 files changed, 245 insertions, 10 deletions

diff --git a/src/shell/event.rs b/src/shell/event.rs
index 9885821..9157c5e 100644
--- a/src/shell/event.rs
+++ b/src/shell/event.rs
@@ -6,6 +6,7 @@ pub enum Event {
     PtyClose,
     ChildRunPipeline(usize, (usize, usize)),
     ChildSuspend(usize),
+    GitInfo(Option<super::git::Info>),
     ClockTimer,
 }
 
@@ -59,6 +60,7 @@ struct Pending {
     pty_close: bool,
     child_run_pipeline: std::collections::VecDeque<(usize, (usize, usize))>,
     child_suspend: std::collections::VecDeque<usize>,
+    git_info: Option<Option<super::git::Info>>,
     clock_timer: bool,
     done: bool,
 }
@@ -76,6 +78,7 @@ impl Pending {
             || self.pty_close
             || !self.child_run_pipeline.is_empty()
             || !self.child_suspend.is_empty()
+            || self.git_info.is_some()
             || self.clock_timer
     }
 
@@ -99,6 +102,9 @@ impl Pending {
         if let Some(idx) = self.child_suspend.pop_front() {
             return Some(Event::ChildSuspend(idx));
         }
+        if let Some(info) = self.git_info.take() {
+            return Some(Event::GitInfo(info));
+        }
         if self.clock_timer {
             self.clock_timer = false;
             return Some(Event::ClockTimer);
@@ -125,6 +131,7 @@ impl Pending {
             Some(Event::ChildSuspend(idx)) => {
                 self.child_suspend.push_back(idx);
             }
+            Some(Event::GitInfo(info)) => self.git_info = Some(info),
             Some(Event::ClockTimer) => self.clock_timer = true,
             None => self.done = true,
         }
diff --git a/src/shell/git.rs b/src/shell/git.rs
new file mode 100644
index 0000000..48e5eea
--- /dev/null
+++ b/src/shell/git.rs
@@ -0,0 +1,201 @@
+#[derive(Debug)]
+pub struct Info {
+    modified_files: bool,
+    staged_files: bool,
+    new_files: bool,
+    commits: bool,
+    active_operation: ActiveOperation,
+    branch: Option<String>,
+    remote_branch_diff: Option<(usize, usize)>,
+}
+
+const MODIFIED: git2::Status = git2::Status::WT_DELETED
+    .union(git2::Status::WT_MODIFIED)
+    .union(git2::Status::WT_RENAMED)
+    .union(git2::Status::WT_TYPECHANGE)
+    .union(git2::Status::CONFLICTED);
+const STAGED: git2::Status = git2::Status::INDEX_DELETED
+    .union(git2::Status::INDEX_MODIFIED)
+    .union(git2::Status::INDEX_NEW)
+    .union(git2::Status::INDEX_RENAMED)
+    .union(git2::Status::INDEX_TYPECHANGE);
+const NEW: git2::Status = git2::Status::WT_NEW;
+
+impl Info {
+    pub fn new(git: &git2::Repository) -> Self {
+        let mut status_options = git2::StatusOptions::new();
+        status_options.include_untracked(true);
+        status_options.update_index(true);
+
+        let statuses = git.statuses(Some(&mut status_options));
+
+        let mut modified_files = false;
+        let mut staged_files = false;
+        let mut new_files = false;
+        if let Ok(statuses) = statuses {
+            for file in statuses.iter() {
+                if file.status().intersects(MODIFIED) {
+                    modified_files = true;
+                }
+                if file.status().intersects(STAGED) {
+                    staged_files = true;
+                }
+                if file.status().intersects(NEW) {
+                    new_files = true;
+                }
+            }
+        }
+
+        let head = git.head();
+        let mut commits = false;
+        let mut branch = None;
+        let mut remote_branch_diff = None;
+
+        if let Ok(head) = head {
+            commits = true;
+            if head.is_branch() {
+                branch = head.shorthand().map(ToString::to_string);
+                remote_branch_diff =
+                    head.resolve()
+                        .ok()
+                        .map(|head| {
+                            (
+                                head.target(),
+                                head.shorthand().map(ToString::to_string),
+                            )
+                        })
+                        .and_then(|(head_id, name)| {
+                            head_id.and_then(|head_id| {
+                                name.and_then(|name| {
+                                    git.refname_to_id(&format!(
+                                        "refs/remotes/origin/{}",
+                                        name
+                                    ))
+                                    .ok()
+                                    .and_then(|remote_id| {
+                                        git.graph_ahead_behind(
+                                            head_id, remote_id,
+                                        )
+                                        .ok()
+                                    })
+                                })
+                            })
+                        });
+            } else {
+                branch =
+                    head.resolve().ok().and_then(|head| head.target()).map(
+                        |oid| {
+                            let mut sha: String = oid
+                                .as_bytes()
+                                .iter()
+                                .take(4)
+                                .map(|b| format!("{:02x}", b))
+                                .collect();
+                            sha.truncate(7);
+                            sha
+                        },
+                    );
+            }
+        }
+
+        let active_operation = match git.state() {
+            git2::RepositoryState::Merge => ActiveOperation::Merge,
+            git2::RepositoryState::Revert
+            | git2::RepositoryState::RevertSequence => {
+                ActiveOperation::Revert
+            }
+            git2::RepositoryState::CherryPick
+            | git2::RepositoryState::CherryPickSequence => {
+                ActiveOperation::CherryPick
+            }
+            git2::RepositoryState::Bisect => ActiveOperation::Bisect,
+            git2::RepositoryState::Rebase
+            | git2::RepositoryState::RebaseInteractive
+            | git2::RepositoryState::RebaseMerge => ActiveOperation::Rebase,
+            _ => ActiveOperation::None,
+        };
+
+        Self {
+            modified_files,
+            staged_files,
+            new_files,
+            commits,
+            active_operation,
+            branch,
+            remote_branch_diff,
+        }
+    }
+}
+
+impl std::fmt::Display for Info {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "g")?;
+
+        if self.modified_files {
+            write!(f, "*")?;
+        }
+        if self.staged_files {
+            write!(f, "+")?;
+        }
+        if self.new_files {
+            write!(f, "?")?;
+        }
+        if !self.commits {
+            write!(f, "!")?;
+            return Ok(());
+        }
+
+        let branch = self.branch.as_ref().map_or("???", |branch| {
+            if branch == "master" {
+                ""
+            } else {
+                branch
+            }
+        });
+        if !branch.is_empty() {
+            write!(f, ":")?;
+        }
+        write!(f, "{}", branch)?;
+
+        if let Some((local, remote)) = self.remote_branch_diff {
+            if local > 0 || remote > 0 {
+                write!(f, ":")?;
+            }
+            if local > 0 {
+                write!(f, "+{}", local)?;
+            }
+            if remote > 0 {
+                write!(f, "-{}", remote)?;
+            }
+        } else {
+            write!(f, ":-")?;
+        }
+
+        write!(f, "{}", self.active_operation)?;
+
+        Ok(())
+    }
+}
+
+#[derive(Debug, Copy, Clone)]
+pub enum ActiveOperation {
+    None,
+    Merge,
+    Revert,
+    CherryPick,
+    Bisect,
+    Rebase,
+}
+
+impl std::fmt::Display for ActiveOperation {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            ActiveOperation::None => Ok(()),
+            ActiveOperation::Merge => write!(f, "(m)"),
+            ActiveOperation::Revert => write!(f, "(v)"),
+            ActiveOperation::CherryPick => write!(f, "(c)"),
+            ActiveOperation::Bisect => write!(f, "(b)"),
+            ActiveOperation::Rebase => write!(f, "(r)"),
+        }
+    }
+}
diff --git a/src/shell/mod.rs b/src/shell/mod.rs
index f82909d..d2d9324 100644
--- a/src/shell/mod.rs
+++ b/src/shell/mod.rs
@@ -3,6 +3,7 @@ use crate::shell::prelude::*;
 use textmode::Textmode as _;
 
 mod event;
+mod git;
 mod history;
 mod prelude;
 mod readline;
@@ -79,7 +80,22 @@ pub async fn main() -> anyhow::Result<i32> {
         });
     }
 
-    let mut shell = Shell::new(crate::info::get_offset());
+    let (git_w, git_r) = async_std::channel::unbounded();
+    {
+        let event_w = event_w.clone();
+        async_std::task::spawn(async move {
+            while let Ok(mut dir) = git_r.recv().await {
+                while let Ok(newer_dir) = git_r.try_recv() {
+                    dir = newer_dir;
+                }
+                let repo = git2::Repository::discover(dir).ok();
+                let info = repo.map(|repo| git::Info::new(&repo));
+                event_w.send(Event::GitInfo(info)).await.unwrap();
+            }
+        });
+    }
+
+    let mut shell = Shell::new(crate::info::get_offset(), git_w).await;
     let event_reader = event::Reader::new(event_r);
     while let Some(event) = event_reader.recv().await {
         match shell.handle_event(event, &event_w).await {
@@ -128,7 +144,8 @@ pub struct Shell {
     readline: readline::Readline,
     history: history::History,
     env: Env,
-    git: Option<git2::Repository>,
+    git: Option<git::Info>,
+    git_w: async_std::channel::Sender<std::path::PathBuf>,
     focus: Focus,
     scene: Scene,
     escape: bool,
@@ -137,20 +154,24 @@ pub struct Shell {
 }
 
 impl Shell {
-    pub fn new(offset: time::UtcOffset) -> Self {
+    pub async fn new(
+        offset: time::UtcOffset,
+        git_w: async_std::channel::Sender<std::path::PathBuf>,
+    ) -> Self {
         let env = Env::new();
         let mut self_ = Self {
             readline: readline::Readline::new(),
             history: history::History::new(),
             env,
             git: None,
+            git_w,
             focus: Focus::Readline,
             scene: Scene::Readline,
             escape: false,
             hide_readline: false,
             offset,
         };
-        self_.update_git();
+        self_.update_git().await;
         self_
     }
 
@@ -297,7 +318,7 @@ impl Shell {
                             let idx = self.env.idx();
                             self.env = entry.env().clone();
                             self.env.set_idx(idx);
-                            self.update_git();
+                            self.update_git().await;
                         }
                         self.set_focus(
                             if self.hide_readline {
@@ -319,6 +340,9 @@ impl Shell {
                     self.set_focus(Focus::Readline, None).await;
                 }
             }
+            Event::GitInfo(info) => {
+                self.git = info;
+            }
             Event::ClockTimer => {}
         };
         Some(Action::Refresh)
@@ -583,7 +607,10 @@ impl Shell {
         self.focus_idx().map_or(Focus::Readline, Focus::History)
     }
 
-    fn update_git(&mut self) {
-        self.git = git2::Repository::discover(self.env.current_dir()).ok();
+    async fn update_git(&mut self) {
+        self.git_w
+            .send(self.env.current_dir().to_path_buf())
+            .await
+            .unwrap();
     }
 }
diff --git a/src/shell/readline.rs b/src/shell/readline.rs
index a1f90bb..ffe92fe 100644
--- a/src/shell/readline.rs
+++ b/src/shell/readline.rs
@@ -23,7 +23,7 @@ impl Readline {
         &self,
         out: &mut impl textmode::Textmode,
         env: &Env,
-        git: Option<&git2::Repository>,
+        git: Option<&super::git::Info>,
         focus: bool,
         offset: time::UtcOffset,
     ) -> anyhow::Result<()> {
@@ -58,8 +58,8 @@ impl Readline {
         }
         out.write_str(" (");
         out.write_str(&crate::format::path(pwd));
-        if git.is_some() {
-            out.write_str("|g");
+        if let Some(info) = git {
+            out.write_str(&format!("|{}", info));
         }
         out.write_str(")");
         out.move_to(self.size.0 - 2, self.size.1 - 4 - idlen - timelen);