index file recovery

3 files changed, 42 insertions, 0 deletions

diff --git a/Git/RecoverRepository.hs b/Git/RecoverRepository.hs
index 8dce04f1e..c2cad53f2 100644
--- a/Git/RecoverRepository.hs
+++ b/Git/RecoverRepository.hs
@@ -10,6 +10,7 @@ module Git.RecoverRepository (
 	retrieveMissingObjects,
 	resetLocalBranches,
 	removeTrackingBranches,
+	rewriteIndex,
 	emptyGoodCommits,
 ) where
 
@@ -19,17 +20,21 @@ import Git.Command
 import Git.Fsck
 import Git.Objects
 import Git.Sha
+import Git.Types
 import qualified Git.Config
 import qualified Git.Construct
 import qualified Git.LsTree as LsTree
+import qualified Git.LsFiles as LsFiles
 import qualified Git.Ref as Ref
 import qualified Git.RefLog as RefLog
+import qualified Git.UpdateIndex as UpdateIndex
 import Utility.Tmp
 import Utility.Rsync
 
 import qualified Data.Set as S
 import qualified Data.ByteString.Lazy as L
 import System.Log.Logger
+import Data.Tuple.Utils
 
 {- Finds and removes corrupt objects from the repository, returning a list
  - of all such objects, which need to be found elsewhere to finish
@@ -349,6 +354,30 @@ verifyTree missing treesha r
 			-- as long as ls-tree succeeded, we're good
 			else cleanup
 
+{- Rewrites the index file, removing from it any files whose blobs are
+ - missing. Returns the list of affected files. -}
+rewriteIndex :: MissingObjects -> Repo -> IO [FilePath]
+rewriteIndex missing r
+	| repoIsLocalBare r = return []
+	| otherwise = do
+		(indexcontents, cleanup) <- LsFiles.stagedDetails [Git.repoPath r] r
+		let (missing, present) = partition ismissing indexcontents
+		unless (null missing) $ do
+			nukeFile (localGitDir r </> "index")
+			UpdateIndex.streamUpdateIndex r
+				=<< (catMaybes <$> mapM reinject present)
+		void cleanup
+		return $ map fst3 missing
+  where
+	getblob (_file, Just sha, Just _mode) = Just sha
+	getblob _ = Nothing
+	ismissing = maybe False (`S.member` missing) . getblob
+	reinject (file, Just sha, Just mode) = case toBlobType mode of
+		Nothing -> return Nothing
+		Just blobtype -> Just <$>
+			UpdateIndex.stageFile sha blobtype file r
+	reinject _ = return Nothing
+
 newtype GoodCommits = GoodCommits (S.Set Sha)
 
 emptyGoodCommits :: GoodCommits
diff --git a/Git/Types.hs b/Git/Types.hs
index 4765aad6c..abfb99f9f 100644
--- a/Git/Types.hs
+++ b/Git/Types.hs
@@ -9,6 +9,7 @@ module Git.Types where
 
 import Network.URI
 import qualified Data.Map as M
+import System.Posix.Types
 
 {- Support repositories on local disk, and repositories accessed via an URL.
  -
@@ -81,3 +82,9 @@ readBlobType "100644" = Just FileBlob
 readBlobType "100755" = Just ExecutableBlob
 readBlobType "120000" = Just SymlinkBlob
 readBlobType _ = Nothing
+
+toBlobType :: FileMode -> Maybe BlobType
+toBlobType 0o100644 = Just FileBlob
+toBlobType 0o100755 = Just ExecutableBlob
+toBlobType 0o120000 = Just SymlinkBlob
+toBlobType _ = Nothing
diff --git a/Git/UpdateIndex.hs b/Git/UpdateIndex.hs
index 5d07e2011..3b33ac846 100644
--- a/Git/UpdateIndex.hs
+++ b/Git/UpdateIndex.hs
@@ -13,6 +13,7 @@ module Git.UpdateIndex (
 	streamUpdateIndex,
 	lsTree,
 	updateIndexLine,
+	stageFile,
 	unstageFile,
 	stageSymlink
 ) where
@@ -61,6 +62,11 @@ updateIndexLine :: Sha -> BlobType -> TopFilePath -> String
 updateIndexLine sha filetype file =
 	show filetype ++ " blob " ++ show sha ++ "\t" ++ indexPath file
 
+stageFile :: Sha -> BlobType -> FilePath -> Repo -> IO Streamer
+stageFile sha filetype file repo = do
+	p <- toTopFilePath file repo
+	return $ pureStreamer $ updateIndexLine sha filetype p
+
 {- A streamer that removes a file from the index. -}
 unstageFile :: FilePath -> Repo -> IO Streamer
 unstageFile file repo = do