    {- streaming directory traversal
     -
     - Copyright 2011-2018 Joey Hess <id@joeyh.name>
     -
     - License: BSD-2-clause
     -}

    {-# LANGUAGE CPP #-}
    {-# LANGUAGE LambdaCase #-}
    {-# OPTIONS_GHC -fno-warn-tabs #-}

    module Utility.Directory.Stream where

    import Control.Monad
    import System.FilePath
    import System.IO.Unsafe (unsafeInterleaveIO)
    import Control.Concurrent
    import Data.Maybe
    import Prelude

    #ifdef mingw32_HOST_OS
    import qualified System.Win32 as Win32
    #else
    import qualified System.Posix as Posix
    #endif

    import Utility.Directory
    import Utility.Exception

    #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 :: FilePath -> IO DirectoryHandle
    openDirectory FilePath
    path = do
    #ifndef mingw32_HOST_OS
    	DirStream
    dirp <- FilePath -> IO DirStream
    Posix.openDirStream FilePath
    path
    	MVar ()
    isopen <- () -> IO (MVar ())
    forall a. a -> IO (MVar a)
    newMVar ()
    	DirectoryHandle -> IO DirectoryHandle
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (MVar () -> DirStream -> DirectoryHandle
    DirectoryHandle MVar ()
    isopen DirStream
    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 -> IO ()
    closeDirectory (DirectoryHandle MVar ()
    isopen DirStream
    dirp) =
    	MVar () -> IO () -> IO ()
    whenOpen MVar ()
    isopen (IO () -> IO ()) -> IO () -> IO ()
    forall a b. (a -> b) -> a -> b
    $
    		DirStream -> IO ()
    Posix.closeDirStream DirStream
    dirp
    #else
    closeDirectory (DirectoryHandle isopen h _ alreadyhave) =
    	whenOpen isopen $ do
    		_ <- tryTakeMVar alreadyhave
    		Win32.findClose h
    #endif
      where
    	whenOpen :: IsOpen -> IO () -> IO ()
    	whenOpen :: MVar () -> IO () -> IO ()
    whenOpen MVar ()
    mv IO ()
    f = do
    		Maybe ()
    v <- MVar () -> IO (Maybe ())
    forall a. MVar a -> IO (Maybe a)
    tryTakeMVar MVar ()
    mv
    		Bool -> IO () -> IO ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    when (Maybe () -> Bool
    forall a. Maybe a -> Bool
    isJust Maybe ()
    v) IO ()
    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 :: DirectoryHandle -> IO (Maybe FilePath)
    readDirectory hdl :: DirectoryHandle
    hdl@(DirectoryHandle MVar ()
    _ DirStream
    dirp) = do
    	FilePath
    e <- DirStream -> IO FilePath
    Posix.readDirStream DirStream
    dirp
    	if FilePath -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null FilePath
    e
    		then do
    			DirectoryHandle -> IO ()
    closeDirectory DirectoryHandle
    hdl
    			Maybe FilePath -> IO (Maybe FilePath)
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return Maybe FilePath
    forall a. Maybe a
    Nothing
    		else Maybe FilePath -> IO (Maybe FilePath)
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (FilePath -> Maybe FilePath
    forall a. a -> Maybe a
    Just FilePath
    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

    -- | Like getDirectoryContents, but rather than buffering the whole
    -- directory content in memory, lazily streams.
    --
    -- This is like lazy readFile in that the handle to the directory remains
    -- open until the whole list is consumed, or until the list is garbage
    -- collected. So use with caution particularly when traversing directory
    -- trees.
    streamDirectoryContents :: FilePath -> IO [FilePath]
    streamDirectoryContents :: FilePath -> IO [FilePath]
    streamDirectoryContents FilePath
    d = FilePath -> IO DirectoryHandle
    openDirectory FilePath
    d IO DirectoryHandle
    -> (DirectoryHandle -> IO [FilePath]) -> IO [FilePath]
    forall a b. IO a -> (a -> IO b) -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= DirectoryHandle -> IO [FilePath]
    collect
      where
    	collect :: DirectoryHandle -> IO [FilePath]
    collect DirectoryHandle
    hdl = DirectoryHandle -> IO (Maybe FilePath)
    readDirectory DirectoryHandle
    hdl IO (Maybe FilePath)
    -> (Maybe FilePath -> IO [FilePath]) -> IO [FilePath]
    forall a b. IO a -> (a -> IO b) -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= \case
    		Maybe FilePath
    Nothing -> [FilePath] -> IO [FilePath]
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return []
    		Just FilePath
    f -> do
    			[FilePath]
    rest <- IO [FilePath] -> IO [FilePath]
    forall a. IO a -> IO a
    unsafeInterleaveIO (DirectoryHandle -> IO [FilePath]
    collect DirectoryHandle
    hdl)
    			[FilePath] -> IO [FilePath]
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (FilePath
    fFilePath -> [FilePath] -> [FilePath]
    forall a. a -> [a] -> [a]
    :[FilePath]
    rest)

    -- | True only when directory exists and contains nothing.
    -- Throws exception if directory does not exist.
    isDirectoryEmpty :: FilePath -> IO Bool
    isDirectoryEmpty :: FilePath -> IO Bool
    isDirectoryEmpty FilePath
    d = IO DirectoryHandle
    -> (DirectoryHandle -> IO ())
    -> (DirectoryHandle -> IO Bool)
    -> IO Bool
    forall (m :: * -> *) a c b.
    (HasCallStack, MonadMask m) =>
    m a -> (a -> m c) -> (a -> m b) -> m b
    bracket (FilePath -> IO DirectoryHandle
    openDirectory FilePath
    d) DirectoryHandle -> IO ()
    closeDirectory DirectoryHandle -> IO Bool
    check
      where
    	check :: DirectoryHandle -> IO Bool
    check DirectoryHandle
    h = do
    		Maybe FilePath
    v <- DirectoryHandle -> IO (Maybe FilePath)
    readDirectory DirectoryHandle
    h
    		case Maybe FilePath
    v of
    			Maybe FilePath
    Nothing -> Bool -> IO Bool
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return Bool
    True
    			Just FilePath
    f
    				| Bool -> Bool
    not (FilePath -> Bool
    dirCruft FilePath
    f) -> Bool -> IO Bool
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return Bool
    False
    				| Bool
    otherwise -> DirectoryHandle -> IO Bool
    check DirectoryHandle
    h