1 files changed, 183 insertions, 0 deletions

diff --git a/src/main.rs b/src/main.rs
new file mode 100644
index 0000000..0b1ec1b
--- /dev/null
+++ b/src/main.rs
@@ -0,0 +1,183 @@
+use std::{env, fmt, path::Path};
+
+use git2::{Oid, Repository, Status};
+use tico::tico;
+
+fn main() {
+    let args = env::args().collect::<Vec<_>>();
+    match args
+        .iter()
+        .map(String::as_str)
+        .collect::<Vec<&str>>()
+        .as_slice()
+    {
+        [_, "cwd", target] => print!("{}", cwd(target)),
+        [_, "vcs", target] => {
+            if let Some(status) = vcs(target) {
+                print!("{}", status)
+            }
+        }
+        _ => (),
+    }
+}
+
+fn cwd(target: &str) -> String {
+    let home = env::var("HOME").unwrap();
+
+    let home_dir_ext = format!("{}{}", home, "/");
+    if target == home.as_str() || target.starts_with(&home_dir_ext) {
+        let replaced = target.replacen(home.as_str(), "~", 1);
+        tico(&replaced)
+    } else {
+        tico(&target)
+    }
+}
+
+struct VcsStatus {
+    branch: Branch,
+    dist: Dist,
+    status: StatusSummary,
+}
+
+impl fmt::Display for VcsStatus {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "{}{}{}#[fg=colour7]",
+            self.branch, self.dist, self.status
+        )
+    }
+}
+
+enum Branch {
+    Id(Oid),
+    Ref(String),
+    Unknown,
+}
+
+impl fmt::Display for Branch {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Branch::Id(id) => write!(f, "#[fg=colour3]{:.7} ", id),
+            Branch::Ref(s) => write!(f, "#[fg=colour8]{} ", s),
+            Branch::Unknown => write!(f, ""),
+        }
+    }
+}
+
+enum Dist {
+    Ahead,
+    Behind,
+    Both,
+    Neither,
+}
+
+impl fmt::Display for Dist {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "#[fg=colour8]{}",
+            match self {
+                Self::Ahead => "↑ ",
+                Self::Behind => "↓ ",
+                Self::Both => "↑↓ ",
+                Self::Neither => "",
+            }
+        )
+    }
+}
+
+enum StatusSummary {
+    WtModified,
+    IdxModified,
+    Conflict,
+    Clean,
+}
+
+impl fmt::Display for StatusSummary {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "{}",
+            match self {
+                Self::WtModified => "#[fg=colour1]×",
+                Self::IdxModified => "#[fg=colour3]±",
+                Self::Conflict => "#[fg=colour5]!",
+                Self::Clean => "#[fg=colour2]·",
+            }
+        )
+    }
+}
+
+fn vcs(target: &str) -> Option<VcsStatus> {
+    let repo = match Path::new(target)
+        .ancestors()
+        .map(Repository::open)
+        .find_map(|r| r.ok())
+    {
+        Some(r) => r,
+        None => return None,
+    };
+
+    let dist = match get_ahead_behind(&repo) {
+        Some((ahead, behind)) if ahead > 0 && behind > 0 => Dist::Both,
+        Some((ahead, _)) if ahead > 0 => Dist::Ahead,
+        Some((_, behind)) if behind > 0 => Dist::Behind,
+        _ => Dist::Neither,
+    };
+
+    let branch = match repo.head() {
+        Ok(reference) if reference.is_branch() => {
+            Branch::Ref(reference.shorthand().unwrap().to_string())
+        }
+        Ok(reference) => Branch::Id(reference.peel_to_commit().unwrap().id()),
+        _ => Branch::Unknown,
+    };
+
+    let status = repo_status(&repo);
+
+    Some(VcsStatus {
+        branch,
+        dist,
+        status,
+    })
+}
+
+fn repo_status(repo: &Repository) -> StatusSummary {
+    for file in repo.statuses(None).unwrap().iter() {
+        match file.status() {
+            // STATE: conflicted
+            Status::CONFLICTED => return StatusSummary::Conflict,
+            // STATE: unstaged (working tree modified)
+            Status::WT_NEW
+            | Status::WT_MODIFIED
+            | Status::WT_DELETED
+            | Status::WT_TYPECHANGE
+            | Status::WT_RENAMED => return StatusSummary::WtModified,
+            // STATE: staged (changes added to index)
+            Status::INDEX_NEW
+            | Status::INDEX_MODIFIED
+            | Status::INDEX_DELETED
+            | Status::INDEX_TYPECHANGE
+            | Status::INDEX_RENAMED => return StatusSummary::IdxModified,
+            // STATE: committed (changes have been saved in the repo)
+            _ => return StatusSummary::Clean,
+        }
+    }
+    StatusSummary::Clean
+}
+
+fn get_ahead_behind(r: &Repository) -> Option<(usize, usize)> {
+    let head = (r.head().ok())?;
+    if !head.is_branch() {
+        return None;
+    }
+
+    let head_name = head.shorthand()?;
+    let head_branch = r.find_branch(head_name, git2::BranchType::Local).ok()?;
+    let upstream = head_branch.upstream().ok()?;
+    let head_oid = head.target()?;
+    let upstream_oid = upstream.get().target()?;
+
+    r.graph_ahead_behind(head_oid, upstream_oid).ok()
+}