From 602257f11428834f07733b688e07fd3e98096272 Mon Sep 17 00:00:00 2001
From: Stephan Dilly <dilly.stephan@gmail.com>
Date: Mon, 15 Jun 2020 01:32:39 +0200
Subject: [PATCH] simplification and fix staging hunks in untracked file

---
 src/app.rs             | 15 --------------
 src/components/diff.rs | 45 ++++++++++++++++++++++++++++++++----------
 src/queue.rs           |  2 --
 3 files changed, 35 insertions(+), 27 deletions(-)

diff --git a/src/app.rs b/src/app.rs
index de933729..15f6ee07 100644
--- a/src/app.rs
+++ b/src/app.rs
@@ -343,21 +343,6 @@ impl App {
                 self.reset.open(action)?;
                 flags.insert(NeedsUpdate::COMMANDS);
             }
-            InternalEvent::StageHunk(hash) => {
-                if let Some((path, is_stage)) =
-                    self.status_tab.selected_path()
-                {
-                    if is_stage {
-                        if sync::unstage_hunk(CWD, path, hash)? {
-                            flags.insert(NeedsUpdate::ALL);
-                        }
-                    } else if sync::stage_hunk(CWD, path, hash)
-                        .is_ok()
-                    {
-                        flags.insert(NeedsUpdate::ALL);
-                    }
-                }
-            }
             InternalEvent::ShowErrorMsg(msg) => {
                 self.msg.show_msg(msg.as_str())?;
                 flags
diff --git a/src/components/diff.rs b/src/components/diff.rs
index de6fe656..c8395f03 100644
--- a/src/components/diff.rs
+++ b/src/components/diff.rs
@@ -2,13 +2,13 @@ use super::{CommandBlocking, DrawableComponent, ScrollType};
 use crate::{
     components::{CommandInfo, Component},
     keys,
-    queue::{Action, InternalEvent, Queue, ResetItem},
+    queue::{Action, InternalEvent, NeedsUpdate, Queue, ResetItem},
     strings,
     ui::{calc_scroll_top, style::Theme},
 };
-use asyncgit::{hash, DiffLine, DiffLineType, FileDiff};
+use asyncgit::{hash, sync, DiffLine, DiffLineType, FileDiff, CWD};
 use crossterm::event::Event;
-use std::{borrow::Cow, cmp};
+use std::{borrow::Cow, cmp, path::Path};
 use strings::commands;
 use tui::{
     backend::Backend,
@@ -271,19 +271,40 @@ impl DiffComponent {
         false
     }
 
-    fn add_hunk(&self) -> Result<()> {
+    fn unstage_hunk(&mut self) -> Result<()> {
         if let Some(hunk) = self.selected_hunk {
             let hash = self.diff.hunks[hunk].header_hash;
-            self.queue
-                .as_ref()
-                .expect("try using queue in immutable diff")
-                .borrow_mut()
-                .push_back(InternalEvent::StageHunk(hash));
+            sync::unstage_hunk(CWD, self.current.path.clone(), hash)?;
+            self.queue_update();
         }
 
         Ok(())
     }
 
+    fn stage_hunk(&mut self) -> Result<()> {
+        if let Some(hunk) = self.selected_hunk {
+            let path = self.current.path.clone();
+            if self.diff.untracked {
+                sync::stage_add_file(CWD, Path::new(&path))?;
+            } else {
+                let hash = self.diff.hunks[hunk].header_hash;
+                sync::stage_hunk(CWD, path, hash)?;
+            }
+
+            self.queue_update();
+        }
+
+        Ok(())
+    }
+
+    fn queue_update(&mut self) {
+        self.queue
+            .as_ref()
+            .expect("try using queue in immutable diff")
+            .borrow_mut()
+            .push_back(InternalEvent::Update(NeedsUpdate::ALL));
+    }
+
     fn reset_hunk(&self) -> Result<()> {
         if let Some(hunk) = self.selected_hunk {
             let hash = self.diff.hunks[hunk].header_hash;
@@ -434,7 +455,11 @@ impl Component for DiffComponent {
                         Ok(true)
                     }
                     keys::ENTER if !self.is_immutable() => {
-                        self.add_hunk()?;
+                        if self.current.is_stage {
+                            self.unstage_hunk()?;
+                        } else {
+                            self.stage_hunk()?;
+                        }
                         Ok(true)
                     }
                     keys::DIFF_RESET_HUNK
diff --git a/src/queue.rs b/src/queue.rs
index e760e6b6..f682d080 100644
--- a/src/queue.rs
+++ b/src/queue.rs
@@ -37,8 +37,6 @@ pub enum InternalEvent {
     ///
     ConfirmedAction(Action),
     ///
-    StageHunk(u64),
-    ///
     ShowErrorMsg(String),
     ///
     Update(NeedsUpdate),