1 files changed, 77 insertions, 39 deletions

diff --git a/Utility/Directory.hs b/Utility/Directory.hs
index e769be7ef..d4e9b358f 100644
--- a/Utility/Directory.hs
+++ b/Utility/Directory.hs
@@ -1,4 +1,4 @@
-{- directory manipulation
+{- directory traversal and manipulation
  -
  - Copyright 2011-2014 Joey Hess <joey@kitenet.net>
  -
@@ -16,7 +16,9 @@ import Control.Monad
 import Control.Monad.IfElse
 import System.FilePath
 import Control.Applicative
+import Control.Concurrent
 import System.IO.Unsafe (unsafeInterleaveIO)
+import Data.Maybe
 
 #ifdef mingw32_HOST_OS
 import qualified System.Win32 as Win32
@@ -31,44 +33,6 @@ import Utility.Exception
 import Utility.Monad
 import Utility.Applicative
 
-{- Unlike getDirectoryContents, this can be used in arbitrarily
- - large directories without using much memory; the list steams lazily.
- -
- - However, any errors that may be encountered while reading the directory
- - contents are *ignored*, rather than throw them in the context of
- - whatever code consumes the lazy list.
- -
- - See https://ghc.haskell.org/trac/ghc/ticket/9266
- -}
-getDirectoryContents' :: FilePath -> IO [FilePath]
-getDirectoryContents' path = loop =<< opendir
-  where
-#ifndef mingw32_HOST_OS
-	opendir = Posix.openDirStream path
-	loop dirp = do
-		v <- tryNonAsync $ Posix.readDirStream dirp
-		case v of
-			(Right ent) | not (null ent) -> do
-				rest <- unsafeInterleaveIO (loop dirp)
-				return (ent:rest)
-			_ -> do
-				void $ tryNonAsync $ Posix.closeDirStream dirp
-				return []
-#else
-	opendir = Win32.findFirstFile (path </> "*")
-	loop (h, fdat) = do
-		-- there is always at least 1 file ("." and "..")
-		ent <- Win32.getFindDataFileName fdat
-		v <- tryNonAsync $ Win32.findNextFile h fdat
-		case v of
-			Right True -> do
-				rest <- unsafeInterleaveIO (loop (h, fdat))
-				return (ent:rest)
-			_ -> do
-				void $ tryNonAsync $ Win32.findClose h
-				return [ent]
-#endif
-
 dirCruft :: FilePath -> Bool
 dirCruft "." = True
 dirCruft ".." = True
@@ -177,3 +141,77 @@ nukeFile file = void $ tryWhenExists go
 #else
 	go = removeFile file
 #endif
+
+#ifndef mingw32_HOST_OS
+data DirectoryHandle = DirectoryHandle IsOpen Posix.DirStream
+#else
+data DirectoryHandle = DirectoryHandle IsOpen Win32.HANDLE Win32.FindData (MVar ())
+#endif
+
+type IsOpen = MVar () -- full when the handle is open
+
+openDirectory :: FilePath -> IO DirectoryHandle
+openDirectory path = do
+#ifndef mingw32_HOST_OS
+	dirp <- Posix.openDirStream path
+	isopen <- newMVar ()
+	return (DirectoryHandle isopen dirp)
+#else
+	(h, fdat) <- Win32.findFirstFile (path </> "*")
+	-- Indicate that the fdat contains a filename that readDirectory
+	-- has not yet returned, by making the MVar be full.
+	-- (There's always at least a "." entry.)
+	alreadyhave <- newMVar ()
+	isopen <- newMVar ()
+	return (DirectoryHandle isopen h fdat alreadyhave)
+#endif
+
+closeDirectory :: DirectoryHandle -> IO ()
+#ifndef mingw32_HOST_OS
+closeDirectory (DirectoryHandle isopen dirp) =
+	whenOpen isopen $
+		Posix.closeDirStream dirp
+#else
+closeDirectory (DirectoryHandle isopen h _ alreadyhave) =
+	whenOpen isopen $ do
+		_ <- tryTakeMVar alreadyhave
+		Win32.findClose h
+#endif
+  where
+	whenOpen :: IsOpen -> IO () -> IO ()
+	whenOpen mv f = do
+		v <- tryTakeMVar mv
+		when (isJust v) f
+
+{- |Reads the next entry from the handle. Once the end of the directory
+is reached, returns Nothing and automatically closes the handle.
+-}
+readDirectory :: DirectoryHandle -> IO (Maybe FilePath)
+#ifndef mingw32_HOST_OS
+readDirectory hdl@(DirectoryHandle _ dirp) = do
+	e <- Posix.readDirStream dirp
+	if null e
+		then do
+			closeDirectory hdl
+			return Nothing
+		else return (Just e)
+#else
+readDirectory hdl@(DirectoryHandle _ h fdat mv) = do
+	-- If the MVar is full, then the filename in fdat has
+	-- not yet been returned. Otherwise, need to find the next
+	-- file.
+	r <- tryTakeMVar mv
+	case r of
+		Just () -> getfn
+		Nothing -> do
+			more <- Win32.findNextFile h fdat
+			if more
+				then getfn
+				else do
+					closeDirectory hdl
+					return Nothing
+  where
+	getfn = do
+		filename <- Win32.getFindDataFileName fdat
+		return (Just filename)
+#endif