Checkout unstaged changes in working directory

4 files changed, 80 insertions, 27 deletions

diff --git a/src/Command.hs b/src/Command.hs
index 599bd90..7ca257a 100644
--- a/src/Command.hs
+++ b/src/Command.hs
@@ -15,6 +15,7 @@ module Command (
     getTerminalOutput,
 ) where
 
+import Control.Monad.Catch
 import Control.Monad.Except
 import Control.Monad.Reader
 
@@ -83,7 +84,7 @@ instance CommandArgumentsType [ Text ] where
 
 
 newtype CommandExec a = CommandExec (ReaderT CommandInput IO a)
-    deriving (Functor, Applicative, Monad, MonadIO)
+    deriving (Functor, Applicative, Monad, MonadIO, MonadThrow, MonadCatch, MonadMask)
 
 instance MonadFail CommandExec where
     fail = tfail . T.pack
diff --git a/src/Command/Checkout.hs b/src/Command/Checkout.hs
index 65857b8..397db79 100644
--- a/src/Command/Checkout.hs
+++ b/src/Command/Checkout.hs
@@ -2,6 +2,7 @@ module Command.Checkout (
     CheckoutCommand,
 ) where
 
+import Data.Maybe
 import Data.Text (Text)
 import Data.Text qualified as T
 
@@ -11,7 +12,7 @@ import Command
 import Repo
 
 
-data CheckoutCommand = CheckoutCommand CheckoutOptions (Maybe RepoName) Text
+data CheckoutCommand = CheckoutCommand CheckoutOptions (Maybe RepoName) (Maybe Text)
 
 data CheckoutOptions = CheckoutOptions
     { coPath :: Maybe FilePath
@@ -43,17 +44,19 @@ instance Command CheckoutCommand where
             "repository subtree to checkout"
         ]
 
-    commandInit _ co = uncurry (CheckoutCommand co) . \case
-        (name : revision : _) -> ( Just (RepoName name), revision )
-        [ name ]              -> ( Just (RepoName name), "HEAD" )
-        []                    -> ( Nothing, "HEAD" )
+    commandInit _ co args = CheckoutCommand co
+        (RepoName <$> listToMaybe args)
+        (listToMaybe $ drop 1 args)
     commandExec = cmdCheckout
 
 cmdCheckout :: CheckoutCommand -> CommandExec ()
-cmdCheckout (CheckoutCommand CheckoutOptions {..} name revision) = do
+cmdCheckout (CheckoutCommand CheckoutOptions {..} name mbrev) = do
     repo <- maybe getDefaultRepo getRepo name
-    root <- maybe (fail $ T.unpack $ "revision `" <> revision <> "' not found") getCommitTree =<< readCommit repo revision
+    root <- getCommitTree =<< case mbrev of
+        Just revision -> readCommit repo revision
+        Nothing -> createWipCommit repo
     tree <- case coSubtree of
         Nothing -> return root
-        Just subtree -> maybe (fail $ "subtree `" <> subtree <> "' not found in revision `" <> T.unpack revision <> "'") return =<< getSubtree subtree root
+        Just subtree -> maybe (fail $ "subtree `" <> subtree <> "' not found in " <> maybe "current worktree" (("revision `" <>) . (<> "'") . T.unpack) mbrev) return =<<
+            getSubtree subtree root
     checkoutAt tree $ maybe "." id coPath
diff --git a/src/Command/Run.hs b/src/Command/Run.hs
index 3968196..383276d 100644
--- a/src/Command/Run.hs
+++ b/src/Command/Run.hs
@@ -129,7 +129,7 @@ argumentJobSource names = do
         case find ((name ==) . jobName) (configJobs config) of
             Just job -> return job
             Nothing -> tfail $ "job `" <> textJobName name <> "' not found"
-    Just jobsetCommit <- flip readCommit "HEAD" =<< getDefaultRepo
+    jobsetCommit <- createWipCommit =<< getDefaultRepo
     oneshotJobSource [ JobSet {..} ]
 
 rangeSource :: Repo -> Text -> Text -> IO JobSource
diff --git a/src/Repo.hs b/src/Repo.hs
index ef95b36..0a6c563 100644
--- a/src/Repo.hs
+++ b/src/Repo.hs
@@ -22,6 +22,7 @@ module Repo (
     getSubtree,
 
     checkoutAt,
+    createWipCommit,
     readCommittedFile,
 
     watchBranch,
@@ -30,8 +31,9 @@ module Repo (
 
 import Control.Concurrent
 import Control.Concurrent.STM
-import Control.Exception
+import Control.Exception (IOException)
 import Control.Monad
+import Control.Monad.Catch
 import Control.Monad.IO.Class
 
 import Data.ByteString (ByteString)
@@ -45,6 +47,7 @@ import Data.Text (Text)
 import Data.Text qualified as T
 import Data.Text.Encoding
 
+import System.Environment
 import System.Exit
 import System.FilePath
 import System.INotify
@@ -155,20 +158,20 @@ openRepo path = do
                 Just dir -> return (Just dir)
                 _        -> return Nothing
 
-mkCommit :: Repo -> CommitId -> IO Commit
+mkCommit :: MonadIO m => Repo -> CommitId -> m Commit
 mkCommit commitRepo commitId_ = do
-    commitDetails <- newMVar Nothing
+    commitDetails <- liftIO $ newMVar Nothing
     return $ Commit {..}
 
-readCommit :: MonadIO m => Repo -> Text -> m (Maybe Commit)
+readCommit :: (MonadIO m, MonadFail m) => Repo -> Text -> m Commit
 readCommit repo@GitRepo {..} ref = liftIO $ do
     readProcessWithExitCode "git" [ "--git-dir=" <> gitDir, "rev-parse", "--verify", "--quiet", T.unpack ref <> "^{commit}" ] "" >>= \case
-        ( ExitSuccess, out, _ ) | cid : _ <- lines out -> Just <$> mkCommit repo (CommitId $ BC.pack cid)
-        _                                              -> return Nothing
+        ( ExitSuccess, out, _ ) | cid : _ <- lines out -> mkCommit repo (CommitId $ BC.pack cid)
+        _                                              -> fail $ "revision `" <> T.unpack ref <> "' not found in `" <> gitDir <> "'"
 
 readCommitFromFile :: MonadIO m => Repo -> FilePath -> m (Maybe Commit)
 readCommitFromFile repo@GitRepo {..} path = liftIO $ do
-    try @IOException (BC.readFile $ gitDir </> path) >>= \case
+    try @IO @IOException (BC.readFile $ gitDir </> path) >>= \case
         Right content | (cid : _) <- BC.lines content -> do
             Just <$> mkCommit repo (CommitId cid)
         _ -> do
@@ -263,18 +266,64 @@ getSubtree path tree = liftIO $ do
 checkoutAt :: (MonadIO m, MonadFail m) => Tree -> FilePath -> m ()
 checkoutAt Tree {..} dest = do
     let GitRepo {..} = treeRepo
-    liftIO $ withMVar gitLock $ \_ -> withSystemTempFile "minici-checkout.index" $ \index _ -> do
-        let gitProc args = (proc "git" args)
-                { env = Just
-                    [ ( "GIT_INDEX_FILE", index )
-                    , ( "GIT_DIR", gitDir )
-                    , ( "GIT_WORK_TREE", "." )
-                    ]
-                }
-        "" <- readCreateProcess (gitProc [ "read-tree", showTreeId treeId ]) ""
-        "" <- readCreateProcess (gitProc [ "checkout-index", "--all", "--prefix=" <> addTrailingPathSeparator dest ]) ""
+    liftIO $ withSystemTempFile "minici-checkout.index" $ \index _ -> do
+        curenv <- getEnvironment
+        let readGitProcess args input =
+                withMVar gitLock $ \_ ->
+                    readCreateProcess (proc "git" args)
+                        { env = Just $ concat
+                            [ [ ( "GIT_INDEX_FILE", index ) ]
+                            , [ ( "GIT_DIR", gitDir ) ]
+                            , [ ( "GIT_WORK_TREE", "." ) ]
+                            , curenv
+                            ]
+                        } input
+        "" <- readGitProcess [ "read-tree", showTreeId treeId ] ""
+        "" <- readGitProcess [ "checkout-index", "--all", "--prefix=" <> addTrailingPathSeparator dest ] ""
         return ()
 
+createWipCommit :: (MonadIO m, MonadMask m, MonadFail m) => Repo -> m Commit
+createWipCommit repo@GitRepo {..} = do
+    withSystemTempFile "minici-wip.index" $ \index _ -> do
+        curenv <- liftIO getEnvironment
+        let readGitProcess mbWorkTree args input = liftIO $ do
+                withMVar gitLock $ \_ ->
+                    readCreateProcess (proc "git" args)
+                        { env = Just $ concat
+                            [ [ ( "GIT_INDEX_FILE", index ) ]
+                            , [ ( "GIT_DIR", gitDir ) ]
+                            , map (( "GIT_WORK_TREE", ) . T.unpack) $ maybeToList mbWorkTree
+                            , curenv
+                            ]
+                        } input
+            mkPair = fmap (T.dropWhile (== ' ')) . T.break (== ' ')
+        info <- map mkPair . takeWhile (not . T.null) . T.splitOn "\0". T.pack <$>
+            readGitProcess Nothing [ "worktree", "list", "--porcelain", "-z" ] ""
+        case ( lookup "worktree" info, lookup "HEAD" info ) of
+            ( Just worktree, Just headRev ) -> do
+                let readGitProcessW = readGitProcess (Just worktree)
+
+                headCommit <- mkCommit repo (CommitId $ encodeUtf8 headRev)
+                headTree <- getCommitTree headCommit
+
+                "" <- readGitProcessW [ "read-tree", "--empty" ] ""
+                status <- map mkPair . T.splitOn "\0" . T.pack <$>
+                    readGitProcessW [ "status", "--porcelain=v1", "-z", "--untracked-files=all" ] ""
+                "" <- readGitProcessW [ "update-index", "--add", "-z", "--stdin" ] $ T.unpack $ T.intercalate "\0" $ map snd status
+                [ tid ] <- lines <$> readGitProcessW [ "write-tree" ] ""
+
+                if TreeId (BC.pack tid) == treeId headTree
+                  then return headCommit
+                  else do
+                    headMsg <- getCommitTitle headCommit
+                    let wipMsg = case lookup "branch" info of
+                            Just branch -> "WIP on " <> branch <> ": " <> headMsg
+                            Nothing -> "WIP: " <> headMsg
+                    [ cid ] <- lines <$> readGitProcessW [ "commit-tree", "-m", T.unpack wipMsg, "-p", T.unpack headRev, tid ] ""
+                    mkCommit repo (CommitId $ BC.pack cid)
+
+            _ -> readCommit repo "HEAD"
+
 
 readCommittedFile :: Commit -> FilePath -> IO (Maybe BL.ByteString)
 readCommittedFile Commit {..} path = do