refactor

2 files changed, 37 insertions, 3 deletions

diff --git a/Annex/Init.hs b/Annex/Init.hs
index 9e0361daf..5aff4cf39 100644
--- a/Annex/Init.hs
+++ b/Annex/Init.hs
@@ -31,6 +31,7 @@ import Annex.Version
 import Annex.Difference
 import Annex.UUID
 import Annex.Link
+import Annex.WorkTree
 import Config
 import Annex.Direct
 import Annex.AdjustedBranch
@@ -39,7 +40,6 @@ import Annex.Hook
 import Annex.InodeSentinal
 import Upgrade
 import Annex.Perms
-import qualified Database.Keys
 import Utility.UserInfo
 #ifndef mingw32_HOST_OS
 import Utility.FileMode
@@ -90,7 +90,7 @@ initialize' mversion = do
 		setVersion (fromMaybe defaultVersion mversion)
 	whenM versionSupportsUnlockedPointers $ do
 		configureSmudgeFilter
-		Database.Keys.scanAssociatedFiles
+		scanUnlockedFiles
 	v <- checkAdjustedClone
 	case v of
 		NeedUpgradeForAdjustedClone -> 
diff --git a/Annex/WorkTree.hs b/Annex/WorkTree.hs
index b2c8cb7f3..fe42ab726 100644
--- a/Annex/WorkTree.hs
+++ b/Annex/WorkTree.hs
@@ -1,6 +1,6 @@
 {- git-annex worktree files
  -
- - Copyright 2013-2015 Joey Hess <id@joeyh.name>
+ - Copyright 2013-2016 Joey Hess <id@joeyh.name>
  -
  - Licensed under the GNU GPL version 3 or higher.
  -}
@@ -12,6 +12,13 @@ import Annex.Link
 import Annex.CatFile
 import Annex.Version
 import Config
+import qualified Git.Ref
+import qualified Git.Branch
+import qualified Git.LsTree
+import qualified Git.Types
+import Database.Types
+import qualified Database.Keys
+import qualified Database.Keys.SQL
 
 {- Looks up the key corresponding to an annexed file in the work tree,
  - by examining what the file links to.
@@ -41,3 +48,30 @@ whenAnnexed a file = ifAnnexed file (a file) (return Nothing)
 
 ifAnnexed :: FilePath -> (Key -> Annex a) -> Annex a -> Annex a
 ifAnnexed file yes no = maybe no yes =<< lookupFile file
+
+{- Find all unlocked files and update the keys database for them. 
+ - 
+ - This is expensive, and so normally the associated files are updated
+ - incrementally when changes are noticed. So, this only needs to be done
+ - when initializing/upgrading a v6 mode repository.
+ -}
+scanUnlockedFiles :: Annex ()
+scanUnlockedFiles = whenM (isJust <$> inRepo Git.Branch.current) $ 
+	Database.Keys.runWriter $ \h -> do
+		showSideAction "scanning for unlocked files"
+		liftIO $ Database.Keys.SQL.dropAllAssociatedFiles h
+		(l, cleanup) <- inRepo $ Git.LsTree.lsTree Git.Ref.headRef
+		forM_ l $ \i -> 
+			when (isregfile i) $
+				maybe noop (add h i)
+					=<< catKey (Git.LsTree.sha i)
+		liftIO $ void cleanup
+  where
+	isregfile i = case Git.Types.toBlobType (Git.LsTree.mode i) of
+		Just Git.Types.FileBlob -> True
+		Just Git.Types.ExecutableBlob -> True
+		_ -> False
+	add h i k = liftIO $ Database.Keys.SQL.addAssociatedFileFast 
+		(toIKey k)
+		(Git.LsTree.file i)
+		h