Support for git worktree creation from a thunk using local git repo

src/Nix/Thunk.hs

@@ -13,6 +13,8 @@ module Nix.Thunk
   , packThunk
   , createThunk
   , createThunk'
+  , createWorktree
+  , CreateWorktreeConfig (..)
   , ThunkPackConfig (..)
   , ThunkConfig (..)
   , updateThunkToLatest

src/Nix/Thunk/Command.hs

@@ -45,24 +45,37 @@ thunkCreateConfig = ThunkCreateConfig
     source = (ThunkCreateSource_Absolute <$> maybeReader (parseGitUri . T.pack))
          <|> (ThunkCreateSource_Relative <$> str)
 
+createWorktreeConfig :: Parser CreateWorktreeConfig
+createWorktreeConfig = CreateWorktreeConfig
+  <$> optional (strOption (short 'b' <> long "branch" <> metavar "BRANCH" <> help "create a new branch"))
+  <*> switch (long "detach" <> short 'd' <> help "detach HEAD at the commit specified in thunk")
+
 data ThunkCommand
   = ThunkCommand_Update ThunkUpdateConfig (NonEmpty FilePath)
   | ThunkCommand_Unpack (NonEmpty FilePath)
+  | ThunkCommand_Worktree CreateWorktreeConfig (FilePath, FilePath)
   | ThunkCommand_Pack ThunkPackConfig (NonEmpty FilePath)
   | ThunkCommand_Create ThunkCreateConfig
   deriving Show
 
 thunkDirList :: Parser (NonEmpty FilePath)
 thunkDirList = (:|)
-  <$> thunkDirArg (metavar "THUNKDIRS..." <> help "Paths to directories containing thunk data")
-  <*> many (thunkDirArg mempty)
-  where
-    thunkDirArg opts = fmap (dropTrailingPathSeparator . normalise) $ strArgument $ action "directory" <> opts
+  <$> dirArg (metavar "THUNKDIRS..." <> help "Paths to directories containing thunk data")
+  <*> many (dirArg mempty)
+
+createWorktreeArgs :: Parser (FilePath, FilePath)
+createWorktreeArgs = (,)
+  <$> dirArg (metavar "THUNKDIR" <> help "Path to directory containing thunk data")
+  <*> dirArg (metavar "GITDIR" <> help "Path to local git repo")
+
+dirArg :: Mod ArgumentFields FilePath -> Parser FilePath
+dirArg opts = fmap (dropTrailingPathSeparator . normalise) $ strArgument $ action "directory" <> opts
 
 thunkCommand :: Parser ThunkCommand
 thunkCommand = hsubparser $ mconcat
   [ command "update" $ info (ThunkCommand_Update <$> thunkUpdateConfig <*> thunkDirList) $ progDesc "Update packed thunk to latest revision available on the tracked branch"
   , command "unpack" $ info (ThunkCommand_Unpack <$> thunkDirList) $ progDesc "Unpack thunk into git checkout of revision it points to"
+  , command "worktree" $ info (ThunkCommand_Worktree <$> createWorktreeConfig <*> createWorktreeArgs) $ progDesc "Create a git worktree of the thunk using the specified local git repo"
   , command "pack" $ info (ThunkCommand_Pack <$> thunkPackConfig <*> thunkDirList) $ progDesc "Pack git checkout or unpacked thunk into thunk that points at the current branch's upstream"
   , command "create" $ info (ThunkCommand_Create <$> thunkCreateConfig) $ progDesc "Create a packed thunk without cloning the repository first"
   ]
@@ -79,5 +92,6 @@ runThunkCommand
 runThunkCommand = \case
   ThunkCommand_Update config dirs -> mapM_ (updateThunkToLatest config) dirs
   ThunkCommand_Unpack dirs -> mapM_ unpackThunk dirs
+  ThunkCommand_Worktree config (thunkDir, gitDir) -> createWorktree thunkDir gitDir config
   ThunkCommand_Pack config dirs -> mapM_ (packThunk config) dirs
   ThunkCommand_Create config -> createThunk' config

src/Nix/Thunk/Internal.hs

@@ -11,6 +11,7 @@
 {-# LANGUAGE QuasiQuotes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TemplateHaskell #-}
+{-# LANGUAGE TupleSections #-}
 {-# LANGUAGE TypeApplications #-}
 module Nix.Thunk.Internal where
 
@@ -220,6 +221,11 @@ data ThunkCreateConfig = ThunkCreateConfig
   , _thunkCreateConfig_destination :: Maybe FilePath
   } deriving Show
 
+data CreateWorktreeConfig = CreateWorktreeConfig
+  { _createWorktreeConfig_branch :: Maybe String
+  , _createWorktreeConfig_detach :: Bool
+  } deriving Show
+
 -- | Convert a GitHub source to a regular Git source. Assumes no submodules.
 forgetGithub :: Bool -> GitHubSource -> GitSource
 forgetGithub useSsh s = GitSource
@@ -1137,6 +1143,69 @@ gitCloneForThunkUnpack gitSrc commit dir = do
   when (_gitSource_fetchSubmodules gitSrc) $
     void $ readGitProcess dir ["submodule", "update", "--recursive", "--init"]
 
+createWorktree :: MonadNixThunk m => FilePath -> FilePath -> CreateWorktreeConfig -> m ()
+createWorktree thunkDir gitDir config = checkThunkDirectory thunkDir *> readThunk thunkDir >>= \case
+  Left err -> failReadThunkErrorWhile "while creating worktree" err
+  Right ThunkData_Checkout -> failWith [i|Thunk at ${thunkDir} is already unpacked|]
+  Right (ThunkData_Packed _ tptr) -> do
+
+    ensureGitRevExist gitDir tptr
+
+    let (thunkParent, thunkName) = splitFileName thunkDir
+    withTempDirectory thunkParent thunkName $ \tmpThunk -> do
+      withSpinner' ("Creating worktree for " <> T.pack thunkName)
+                   (Just (const $ "Created worktree for " <> T.pack thunkName)) $ do
+        currentDir <- liftIO getCurrentDirectory
+        let worktreePath = currentDir </> tmpThunk </> unpackedDirName
+            thunkFullPath = currentDir </> thunkDir </> unpackedDirName
+
+            -- Create a new branch with the user specified name if provided
+            -- else fallback to the branch specified in thunk
+            -- If a local branch already exists in gitDir, the worktree creation will fail
+            -- In which case the user should specify an alternate branch or use "-d"
+            mBranchName = case _createWorktreeConfig_branch config of
+              Just b -> Just b
+              _ -> T.unpack . untagName <$> (_gitSource_branch $ thunkSourceToGitSource $ _thunkPtr_source tptr)
+
+        _ <- readGitProcess gitDir $
+          [ "worktree", "add"
+          , worktreePath
+          , refToHexString (_thunkRev_commit $ _thunkPtr_rev tptr)
+          ] ++ (if _createWorktreeConfig_detach config
+                then ["-d"]
+                else maybe [] (\b -> ["-b",  b]) mBranchName)
+
+        liftIO $ removePathForcibly thunkDir
+
+        _ <- readGitProcess gitDir $
+          [ "worktree", "move"
+          , normalise worktreePath
+          , normalise thunkFullPath]
+        pure ()
+
+-- | Ensures that the git repo contains the revision specified in the ThunkPtr
+-- by doing fetch from remote if necessary.
+ensureGitRevExist :: MonadNixThunk m => FilePath -> ThunkPtr -> m ()
+ensureGitRevExist gitDir tptr = do
+  isdir <- liftIO $ doesDirectoryExist gitDir
+  -- check .git
+  unless isdir $ failWith $ "Git directory does not exist: " <> T.pack gitDir
+
+  (exitCode, _, _) <- readCreateProcessWithExitCode $
+    gitProc gitDir $
+      [ "reflog"
+      , "exists"
+      , refToHexString (_thunkRev_commit $ _thunkPtr_rev tptr)
+      ]
+
+  when (exitCode /= ExitSuccess) $ do
+    void $ readGitProcess gitDir $
+      [ "fetch"
+      , T.unpack $ gitUriToText (_gitSource_url $ thunkSourceToGitSource $ _thunkPtr_source tptr)
+      , refToHexString (_thunkRev_commit $ _thunkPtr_rev tptr)
+      ]
+
+
 -- | Read a git process ignoring the global configuration (according to 'ignoreGitConfig').
 readGitProcess :: MonadNixThunk m => FilePath -> [String] -> m Text
 readGitProcess dir = readProcessAndLogOutput (Notice, Notice) . ignoreGitConfig . gitProc dir
@@ -1172,8 +1241,18 @@ packThunk' noTrail (ThunkPackConfig force thunkConfig) thunkDir = checkThunkDire
     (finalMsg noTrail $ const $ "Packed thunk " <> T.pack thunkDir) $
     do
       let checkClean = if force then CheckClean_NoCheck else CheckClean_FullCheck
-      thunkPtr <- modifyThunkPtrByConfig thunkConfig <$> getThunkPtr checkClean thunkDir (_thunkConfig_private thunkConfig)
-      liftIO $ removePathForcibly thunkDir
+      (thunkPtr, isWorktree) <- first (modifyThunkPtrByConfig thunkConfig)
+        <$> getThunkPtr checkClean thunkDir (_thunkConfig_private thunkConfig)
+      if isWorktree
+        then void $ do
+          -- Remove the branch locally, and then remove the worktree
+          case _gitSource_branch $ thunkSourceToGitSource $ _thunkPtr_source thunkPtr of
+            Just branch -> do
+              void $ readGitProcess thunkDir ["switch", "--detach"]
+              void $ readGitProcess thunkDir ["branch", "-d", T.unpack $ untagName branch]
+            Nothing -> pure () -- Should never happen
+          readGitProcess thunkDir ["worktree", "remove", "."]
+        else liftIO $ removePathForcibly thunkDir
       createThunk thunkDir $ Right thunkPtr
       pure thunkPtr
 
@@ -1193,14 +1272,22 @@ data CheckClean
   | CheckClean_NoCheck
   -- ^ Don't check that the repo is clean
 
-getThunkPtr :: forall m. MonadNixThunk m => CheckClean -> FilePath -> Maybe Bool -> m ThunkPtr
+getThunkPtr :: forall m. MonadNixThunk m => CheckClean -> FilePath -> Maybe Bool -> m (ThunkPtr, Bool)
 getThunkPtr gitCheckClean dir mPrivate = do
   let repoLocations = nubOrd $ map (first normalise)
         [(".git", "."), (unpackedDirName </> ".git", unpackedDirName)]
   repoLocation' <- liftIO $ flip findM repoLocations $ doesDirectoryExist . (dir </>) . fst
-  thunkDir <- case repoLocation' of
-    Nothing -> failWith [i|Can't find an unpacked thunk in ${dir}|]
-    Just (_, path) -> pure $ normalise $ dir </> path
+  (thunkDir, isWorktree) <- case repoLocation' of
+    Nothing -> do
+      ff <- liftIO $ flip findM repoLocations $ doesFileExist . (dir </>) . fst
+      case ff of
+        Nothing -> failWith [i|Can't find an unpacked thunk in ${dir}|]
+        Just (gitPath, path) -> do
+          putLog Informational "Couldn't find .git dir, looking for a worktree instead"
+          fileContents <- liftIO $ T.readFile (dir </> gitPath)
+          unless (T.isPrefixOf "gitdir: " fileContents) $ failWith [i|Can't find an unpacked thunk or worktree in ${dir}|]
+          pure $ (normalise $ dir </> path, True)
+    Just (_, path) -> pure $ (normalise $ dir </> path, False)
 
   let (checkClean, checkIgnored) = case gitCheckClean of
         CheckClean_FullCheck -> (True, True)
@@ -1210,7 +1297,7 @@ getThunkPtr gitCheckClean dir mPrivate = do
     "thunk pack: thunk checkout contains unsaved modifications"
 
   -- Check whether there are any stashes
-  when checkClean $ do
+  when (checkClean && not isWorktree) $ do
     stashOutput <- readGitProcess thunkDir ["stash", "list"]
     unless (T.null stashOutput) $
       failWith $ T.unlines $
@@ -1230,12 +1317,16 @@ getThunkPtr gitCheckClean dir mPrivate = do
         ]
       _ -> return (b, c)
 
-  -- Get information on all branches and their (optional) designated upstream
-  -- correspondents
+  let refs = if isWorktree
+        -- Get information on current branch only
+        then "refs/heads/" <> maybe "" T.unpack mCurrentBranch
+        -- Get information on all branches and their (optional) designated
+        -- upstream correspondents
+        else "refs/heads/"
   headDump :: [Text] <- T.lines <$> readGitProcess thunkDir
     [ "for-each-ref"
     , "--format=%(refname:short) %(upstream:short) %(upstream:remotename)"
-    , "refs/heads/"
+    , refs
     ]
 
   (headInfo :: Map Text (Maybe (Text, Text)))
@@ -1319,7 +1410,7 @@ getThunkPtr gitCheckClean dir mPrivate = do
   remoteUri <- case parseGitUri remoteUri' of
     Nothing -> failWith $ "Could not identify git remote: " <> remoteUri'
     Just uri -> pure uri
-  uriThunkPtr remoteUri mPrivate mCurrentBranch mCurrentCommit
+  (, isWorktree) <$> uriThunkPtr remoteUri mPrivate mCurrentBranch mCurrentCommit
 
 -- | Get the latest revision available from the given source
 getLatestRev :: MonadNixThunk m => ThunkSource -> m ThunkRev

tests.nix

@@ -243,5 +243,83 @@ in
           cmp ~/code/myapp/git.json ~/code/myapp-remote-merge-master/git.json;
           nix-thunk unpack ~/code/myapp
         """)
+
+      with subtest("can create worktree using existing repo, doing detached HEAD when no branch is specified in thunk"):
+        client.succeed("""
+          nix-thunk create root@githost:/root/myorg/myapp.git ~/code/myapp-2;
+          git clone root@githost:/root/myorg/myapp.git ~/code/myapp-mainrepo;
+          nix-thunk worktree ~/code/myapp-2 ~/code/myapp-mainrepo;
+          branch=$(git -C ~/code/myapp-2 branch --show-current);
+          if [ ! -z $branch ]; then
+             exit 1
+          fi
+        """);
+
+      with subtest("gives error when packing worktree on detached HEAD"):
+        client.fail("""
+          nix-thunk pack ~/code/myapp-2;
+        """)
+
+      with subtest("can pack worktree with branch specified, and removes the local branch after packing"):
+        client.succeed("""
+          git -C ~/code/myapp-mainrepo checkout -b temp-branch;
+          git -C ~/code/myapp-2 checkout master;
+          nix-thunk pack ~/code/myapp-2;
+        """);
+        client.fail("""
+          git -C ~/code/myapp-mainrepo rev-parse --verify master;
+        """)
+
+      with subtest("can create worktree, and checkout the default branch"):
+        client.succeed("""
+          nix-thunk worktree ~/code/myapp-2 ~/code/myapp-mainrepo;
+          git -C ~/code/myapp-mainrepo rev-parse --verify master;
+        """);
+
+      with subtest("fails if the branch is already checked out"):
+        client.succeed("""
+          git -C ~/code/myapp-2 branch --set-upstream-to origin/master;
+          nix-thunk pack ~/code/myapp-2;
+          git -C ~/code/myapp-mainrepo checkout -b master;
+        """);
+        client.fail("""
+          nix-thunk worktree ~/code/myapp-2 ~/code/myapp-mainrepo;
+        """);
+
+      with subtest("can create worktree, when a new branch is specified"):
+        client.succeed("""
+          nix-thunk worktree ~/code/myapp-2 ~/code/myapp-mainrepo -b somebranch-2;
+          git -C ~/code/myapp-mainrepo rev-parse --verify somebranch-2;
+        """);
+
+      with subtest("fails when packing worktree with unpushed branch"):
+        client.fail("""
+          nix-thunk pack ~/code/myapp-2; # has somebranch-2 checked out
+        """)
+
+      with subtest("can pack worktree having unpushed branches"):
+        client.succeed("""
+          git -C ~/code/myapp-mainrepo checkout temp-branch;
+          git -C ~/code/myapp-2 checkout master; # repo still contains somebranch-2, having no remote
+          git -C ~/code/myapp-2 branch --set-upstream-to origin/master;
+          nix-thunk pack ~/code/myapp-2;
+        """)
+
+      with subtest("fails to pack worktree containing modifications"):
+        client.succeed("""
+          nix-thunk worktree ~/code/myapp-2 ~/code/myapp-mainrepo;
+          touch ~/code/myapp-2/extra-file;
+        """)
+        client.fail("""
+          nix-thunk pack ~/code/myapp-2;
+        """)
+
+      with subtest("can pack worktree with stashed changes"):
+        client.succeed("""
+          git -C ~/code/myapp-2 add extra-file;
+          git -C ~/code/myapp-2 stash;
+          git -C ~/code/myapp-2 branch --set-upstream-to origin/master;
+          nix-thunk pack ~/code/myapp-2;
+        """)
       '';
   }) {}