1 files changed, 56 insertions, 18 deletions

diff --git a/Annex/Direct.hs b/Annex/Direct.hs
index 7b0dbc1e0..d4b73860e 100644
--- a/Annex/Direct.hs
+++ b/Annex/Direct.hs
@@ -15,6 +15,7 @@ import qualified Git.Merge
 import qualified Git.DiffTree as DiffTree
 import qualified Git.Config
 import qualified Git.Ref
+import qualified Git.Branch
 import Git.Sha
 import Git.FilePath
 import Git.Types
@@ -236,28 +237,65 @@ changedDirect oldk f = do
 	whenM (pure (null locs) <&&> not <$> inAnnex oldk) $
 		logStatus oldk InfoMissing
 
-{- Since direct mode repositories use core.bare=true, pushes are allowed
- - that overwrite the master branch (or whatever branch is currently
- - checked out) at any time. But committing when a change has been pushed
- - to the current branch and not merged into the work tree will have the
- - effect of reverting the pushed changes.
- -
- - To avoid this problem, when git annex commits in a direct mode
- - repository, it does not commit to HEAD, but instead to annexhead.
- - This ref always contains the last local commit.
- -}
-annexheadRef :: Ref
-annexheadRef = Ref $ "refs" </> "annexhead"
-
 {- Enable/disable direct mode. -}
 setDirect :: Bool -> Annex ()
 setDirect wantdirect = do
-	when wantdirect $ do
-		f <- fromRepo $ Git.Ref.file annexheadRef
-		v <- inRepo $ Git.Ref.sha Git.Ref.headRef
-		liftIO $ maybe (nukeFile f) (writeFile f . show) v
+	if wantdirect
+		then do
+			switchHEAD
+			setbare
+		else do
+			setbare
+			switchHEADBack
 	setConfig (annexConfig "direct") val
-	setConfig (ConfigKey Git.Config.coreBare) val
 	Annex.changeGitConfig $ \c -> c { annexDirect = wantdirect }
   where
 	val = Git.Config.boolConfig wantdirect
+	setbare = setConfig (ConfigKey Git.Config.coreBare) val
+
+{- Since direct mode sets core.bare=true, incoming pushes could change
+ - the currently checked out branch. To avoid this problem, HEAD
+ - is changed to a internal ref that nothing is going to push to.
+ -
+ - For refs/heads/master, use refs/heads/annex/direct/master;
+ - this way things that show HEAD (eg shell prompts) will
+ - hopefully show just "master". -}
+directBranch :: Ref -> Ref
+directBranch orighead = case split "/" $ show orighead of
+	("refs":"heads":"annex":"direct":_) -> orighead
+	("refs":"heads":rest) ->
+		Ref $ "refs/heads/annex/direct/" ++ intercalate "/" rest
+	_ -> Ref $ "refs/heads/" ++ show (Git.Ref.base orighead)
+
+{- Converts a directBranch back to the original branch.
+ -
+ - Any other ref is left unchanged.
+ -}
+fromDirectBranch :: Ref -> Ref
+fromDirectBranch directhead = case split "/" $ show directhead of
+	("refs":"heads":"annex":"direct":rest) -> 
+		Ref $ "refs/heads/" ++ intercalate "/" rest
+	_ -> directhead
+
+switchHEAD :: Annex ()
+switchHEAD = maybe noop switch =<< inRepo Git.Branch.currentUnsafe
+  where
+	switch orighead = do
+		let newhead = directBranch orighead
+		maybe noop (inRepo . Git.Branch.update newhead)
+			=<< inRepo (Git.Ref.sha orighead)
+		inRepo $ Git.Branch.checkout newhead
+
+switchHEADBack :: Annex ()
+switchHEADBack = maybe noop switch =<< inRepo Git.Branch.currentUnsafe
+  where
+	switch currhead = do
+		let orighead = fromDirectBranch currhead
+		v <- inRepo $ Git.Ref.sha currhead
+		case v of
+			Just headsha
+				| orighead /= currhead -> do
+					inRepo $ Git.Branch.update orighead headsha
+					inRepo $ Git.Branch.checkout orighead
+					inRepo $ Git.Branch.delete currhead
+			_ -> inRepo $ Git.Branch.checkout orighead