Optimize file traversal

haskell · Apr 29, 2024 · 2984109 · 2984109
1 parent de8c9f0
commit 2984109
Show file tree

Hide file tree

Showing 3 changed files with 82 additions and 1 deletion.
diff --git a/lib/GHCup/Prelude/File.hs b/lib/GHCup/Prelude/File.hs
@@ -188,7 +188,7 @@ getDirectoryContentsRecursive (fromGHCupPath -> fp) = getDirectoryContentsRecurs
 getDirectoryContentsRecursiveUnsafe :: (MonadResource m)
                                     => FilePath
                                     -> ConduitT i FilePath m ()
-getDirectoryContentsRecursiveUnsafe fp = sourceDirectoryDeep False fp .| C.map rmPrefix
+getDirectoryContentsRecursiveUnsafe fp = sourceDirectoryDeep' fp .| C.map rmPrefix
  where
   rmPrefix = joinPath . fromJust . L.stripPrefix (splitDirectories fp) . splitDirectories
 

diff --git a/lib/GHCup/Prelude/File/Posix.hs b/lib/GHCup/Prelude/File/Posix.hs
@@ -35,8 +35,11 @@ import qualified System.Posix.Files            as PF
 import qualified System.Posix.IO               as SPI
 import qualified System.Posix as Posix
 import qualified GHCup.Prelude.File.Posix.Foreign as FD
+import GHCup.Prelude.File.Posix.Traversals
 import GHC.IO.Exception (IOException(ioe_type), IOErrorType (..))
 
+import qualified Data.Conduit.Combinators as C
+
 
 -- | On unix, we can use symlinks, so we just get the
 -- symbolic link target.
@@ -261,3 +264,33 @@ catchErrno en a1 a2 =
 removeEmptyDirectory :: FilePath -> IO ()
 removeEmptyDirectory = PD.removeDirectory
 
+sourceDirectory' :: MonadResource m => FilePath -> ConduitT i (FD.DirType, FilePath) m ()
+sourceDirectory' dir =
+    bracketP (openDirStreamPortable dir) closeDirStreamPortable go
+  where
+    go ds =
+        loop
+      where
+        loop = do
+            (typ, e) <- liftIO $ readDirEntPortable ds
+            if
+              | null e    -> return ()
+              | "." == e  -> loop
+              | ".." == e -> loop
+              | otherwise -> do
+                  yield $ (typ, dir </> e)
+                  loop
+
+sourceDirectoryDeep' :: MonadResource m
+                     => FilePath -- ^ Root directory
+                     -> ConduitT i FilePath m ()
+sourceDirectoryDeep' fp' = start fp' .| C.map snd
+  where
+    start :: MonadResource m => FilePath -> ConduitT i (FD.DirType, FilePath) m ()
+    start dir = sourceDirectory' dir .| awaitForever go
+
+    go :: MonadResource m => (FD.DirType, FilePath) -> ConduitT (FD.DirType, FilePath) (FD.DirType, FilePath) m ()
+    go (typ, fp)
+      | FD.dtDir == typ = start fp
+      | otherwise       = yield (typ, fp)
+
diff --git a/lib/GHCup/Prelude/File/Windows.hs b/lib/GHCup/Prelude/File/Windows.hs
@@ -29,6 +29,8 @@ import qualified System.Win32.File             as WS
 
 import           Data.Bits ((.&.))
 
+import           Conduit
+import qualified Data.Conduit.Combinators as C
 
 
 -- | On unix, we can use symlinks, so we just get the
@@ -223,3 +225,49 @@ fromExtendedLengthPath ePath =
       not ('/' `elem` path ||
            "." `elem` splitDirectories path ||
            ".." `elem` splitDirectories path)
+
+
+sourceDirectory' :: MonadResource m => FilePath -> ConduitT i (WS.FileAttributeOrFlag, FilePath) m ()
+sourceDirectory' dir =
+    bracketP (alloc dir) dealloc go
+  where
+    go (topdir, b, h, fd) =
+        loop b
+      where
+        loop False = return ()
+        loop True = do
+          f <- liftIO $ WS.getFindDataFileName fd
+           more <- liftIO $ WS.findNextFile h fd
+
+           -- can't get file attribute from FindData yet (needs Win32 PR)
+           fattr <- liftIO $ WS.getFileAttributes (topdir </> f)
+
+           if | f == "." || f == ".." -> loop more
+              | otherwise             -> do
+                                           yield (fattr, f)
+                                           loop more
+
+    alloc topdir = do
+      query <- liftIO $ furnishPath (topdir </> "*")
+      (h, fd) <- liftIO $ WS.findFirstFile query
+      pure (topdir, True, h, fd)
+
+    dealloc (_, _, fd, _) = liftIO $ WS.findClose fd
+
+sourceDirectoryDeep' :: MonadResource m
+                     => FilePath -- ^ Root directory
+                     -> ConduitT i FilePath m ()
+sourceDirectoryDeep' fp' = start fp' .| C.map snd
+  where
+    isDir attrs = attrs .&. WS.fILE_ATTRIBUTE_DIRECTORY /= 0
+
+    start :: MonadResource m => FilePath -> ConduitT i (WS.FileAttributeOrFlag, FilePath) m ()
+    start dir = sourceDirectory' dir .| awaitForever go
+
+    go :: MonadResource m
+       => (WS.FileAttributeOrFlag, FilePath)
+       -> ConduitT (WS.FileAttributeOrFlag, FilePath) (WS.FileAttributeOrFlag, FilePath) m ()
+    go (typ, fp)
+      | isDir typ = start fp
+      | otherwise = yield (typ, fp)
+