-- |
-- Module:     System.Directory.OsPath.Contents
-- Copyright:  (c) Sergey Vinokurov 2024
-- License:    Apache-2.0 (see LICENSE)
-- Maintainer: serg.foo@gmail.com

{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}

module System.Directory.OsPath.Contents
  ( getDirectoryContentsRecursive
  , getDirectoryContentsWithFilterRecursive

  , listContentsRecFold
  ) where

import Control.Exception (onException)
import Data.Coerce (coerce, Coercible)
import System.IO.Unsafe (unsafeInterleaveIO)
import System.OsPath

import System.Directory.OsPath.Streaming.Internal (DirStream)
import qualified System.Directory.OsPath.Streaming.Internal as Streaming
import qualified System.Directory.OsPath.Streaming.Internal.Raw as Raw
import System.Directory.OsPath.Types

-- | Recursively list all the files and directories in a directory and all subdirectories.
--
-- The directory structure is traversed depth-first.
--
-- The result is generated lazily so is not well defined if the source
-- directory structure changes before the list is fully consumed.
--
-- Symlinks within directory structure may cause result to be infinitely long.
getDirectoryContentsRecursive
  :: OsPath
  -> IO [(OsPath, FileType)]
getDirectoryContentsRecursive :: OsString -> IO [(OsString, FileType)]
getDirectoryContentsRecursive OsString
root =
  Maybe Int
-> (forall c.
    OsString
    -> OsString
    -> Relative OsString
    -> Basename OsString
    -> FileType
    -> SymlinkType
    -> ((OsString, FileType) -> IO c -> IO c)
    -> (IO c -> IO c)
    -> IO c
    -> IO c)
-> (OsString
    -> OsString
    -> Relative OsString
    -> Basename OsString
    -> FileType
    -> IO (Maybe (OsString, FileType)))
-> Maybe OsString
-> IO [(OsString, FileType)]
forall (f :: * -> *) a b.
(Foldable f, Coercible b OsString) =>
Maybe Int
-> (forall c.
    OsString
    -> b
    -> Relative OsString
    -> Basename OsString
    -> FileType
    -> SymlinkType
    -> (a -> IO c -> IO c)
    -> (IO c -> IO c)
    -> IO c
    -> IO c)
-> (OsString
    -> b
    -> Relative OsString
    -> Basename OsString
    -> FileType
    -> IO (Maybe a))
-> f b
-> IO [a]
listContentsRecFold'
    Maybe Int
forall a. Maybe a
Nothing
    (\OsString
_ OsString
_ (Relative OsString
path) Basename OsString
_ FileType
ft SymlinkType
_ (OsString, FileType) -> IO c -> IO c
cons IO c -> IO c
prependSubdir IO c
rest -> (OsString, FileType) -> IO c -> IO c
cons (OsString
path, FileType
ft) (IO c -> IO c) -> IO c -> IO c
forall a b. (a -> b) -> a -> b
$ IO c -> IO c
prependSubdir IO c
rest)
    (\OsString
_ OsString
_ (Relative OsString
path) Basename OsString
_ FileType
ft -> Maybe (OsString, FileType) -> IO (Maybe (OsString, FileType))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((OsString, FileType) -> Maybe (OsString, FileType)
forall a. a -> Maybe a
Just (OsString
path, FileType
ft)))
    (OsString -> Maybe OsString
forall a. a -> Maybe a
Just OsString
root)

-- | Recursively list all the files and directories that satisfy given
-- predicate in a directory and all subdirectories. Descending into
-- some subdirectories may be avoided by filtering them out with a
-- visiting predicate.
--
-- Not visited directory entry may still be reported depending on the
-- collection predicate.
--
-- The directory structure is traversed depth-first.
--
-- The result is generated lazily so is not well defined if the source
-- directory structure changes before the list is fully consumed.
--
-- Symlinks within directory structure may cause result to be infinitely long, but
-- they can be filtered out with a suitable directory visiting predicate.
getDirectoryContentsWithFilterRecursive
  :: (Basename OsPath -> SymlinkType -> Bool) -- ^ Whether to visit a directory
  -> (Basename OsPath ->                Bool) -- ^ Whether to collect given directory element, either file or directory.
  -> OsPath
  -> IO [(OsPath, FileType)]
getDirectoryContentsWithFilterRecursive :: (Basename OsString -> SymlinkType -> Bool)
-> (Basename OsString -> Bool)
-> OsString
-> IO [(OsString, FileType)]
getDirectoryContentsWithFilterRecursive Basename OsString -> SymlinkType -> Bool
visitPred Basename OsString -> Bool
collectPred OsString
root =
  Maybe Int
-> (forall c.
    OsString
    -> OsString
    -> Relative OsString
    -> Basename OsString
    -> FileType
    -> SymlinkType
    -> ((OsString, FileType) -> IO c -> IO c)
    -> (IO c -> IO c)
    -> IO c
    -> IO c)
-> (OsString
    -> OsString
    -> Relative OsString
    -> Basename OsString
    -> FileType
    -> IO (Maybe (OsString, FileType)))
-> Maybe OsString
-> IO [(OsString, FileType)]
forall (f :: * -> *) a b.
(Foldable f, Coercible b OsString) =>
Maybe Int
-> (forall c.
    OsString
    -> b
    -> Relative OsString
    -> Basename OsString
    -> FileType
    -> SymlinkType
    -> (a -> IO c -> IO c)
    -> (IO c -> IO c)
    -> IO c
    -> IO c)
-> (OsString
    -> b
    -> Relative OsString
    -> Basename OsString
    -> FileType
    -> IO (Maybe a))
-> f b
-> IO [a]
listContentsRecFold'
    Maybe Int
forall a. Maybe a
Nothing
    (\OsString
_ OsString
_ (Relative OsString
path) Basename OsString
basename FileType
ft SymlinkType
symlink (OsString, FileType) -> IO c -> IO c
cons IO c -> IO c
prependSubdir IO c
rest ->
       (if Basename OsString -> Bool
collectPred Basename OsString
basename then (OsString, FileType) -> IO c -> IO c
cons (OsString
path, FileType
ft) else IO c -> IO c
forall a. a -> a
id) (IO c -> IO c) -> IO c -> IO c
forall a b. (a -> b) -> a -> b
$
         if Basename OsString -> SymlinkType -> Bool
visitPred Basename OsString
basename SymlinkType
symlink
         then IO c -> IO c
prependSubdir IO c
rest
         else IO c
rest)
    (\OsString
_ OsString
_ (Relative OsString
path) Basename OsString
basename FileType
ft ->
      Maybe (OsString, FileType) -> IO (Maybe (OsString, FileType))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (OsString, FileType) -> IO (Maybe (OsString, FileType)))
-> Maybe (OsString, FileType) -> IO (Maybe (OsString, FileType))
forall a b. (a -> b) -> a -> b
$
        if Basename OsString -> Bool
collectPred Basename OsString
basename
        then (OsString, FileType) -> Maybe (OsString, FileType)
forall a. a -> Maybe a
Just (OsString
path, FileType
ft)
        else Maybe (OsString, FileType)
forall a. Maybe a
Nothing)
    (OsString -> Maybe OsString
forall a. a -> Maybe a
Just OsString
root)

{-# INLINE listContentsRecFold #-}
-- | The most general form of gathering directory contents.
--
-- Treats symlinks the same as regular files and directories. Folding functions can
-- decide how to handle symlinks.
--
-- Both directory and file actions can throw exceptions and this function
-- will try to close finished directory streams promptly (they’ll be closed
-- by GC in the worst case).
listContentsRecFold
  :: forall f a b. (Foldable f, Coercible b OsPath)
  => Maybe Int
  -- ^ Depth limit if specified, negative values treated the same as positive ones.
  -> (forall c. OsPath -> b -> Relative OsPath -> Basename OsPath -> SymlinkType -> (a -> IO c -> IO c) -> (IO c -> IO c) -> IO c -> IO c)
  -- ^ Decide how to fold directory and its children given its path.
  --
  -- Can do IO actions to plan what to do and typically should derive its
  -- result from last @IO c@ argument.
  --
  -- Returns @IO c@ where @c@ is hidden from the user so the only way
  -- to make it is to construct from the passed @IO c@ action.
  --
  -- Arguments:
  --
  -- * @OsPath@              - absolute path to the visited directory
  -- * @b@                   - root of the visited directory as passed originally in @f b@ to the bigger fold function
  -- * @Relative OsPath@     - path to the visited directory relative to the previous @b@ argument
  -- * @Basename OsPath@     - name of the visited directory without slashes
  -- * @SymlinkType@         - symlink status of the visited directory
  -- * @(a -> IO c -> IO c)@ - can be used to record some output (@a@) about the directory itself
  -- * @(IO c -> IO c)@      - traverse inside this directory, can be ignored to skip its children
  -- * @IO c@                - continue scanning not yet visited parts, must be used to construct return value (otherwise it won’t typecheck!)
  --
  -- The passed @(IO c -> IO c)@ argument function should (but is not required to)
  -- be applied in the returned function and it will prepend results for subdirectories
  -- of the directory being analyzed. If not applied these subdirectories will be skipped,
  -- this way ignoring particular directory and all its children can be achieved.
  -> (OsPath -> b -> Relative OsPath -> Basename OsPath -> FileType -> IO (Maybe a))
  -- ^ What to do with file
  -> f b
  -- ^ Roots to search in, either absolute or relative
  -> IO [a]
listContentsRecFold :: forall (f :: * -> *) a b.
(Foldable f, Coercible b OsString) =>
Maybe Int
-> (forall c.
    OsString
    -> b
    -> Relative OsString
    -> Basename OsString
    -> SymlinkType
    -> (a -> IO c -> IO c)
    -> (IO c -> IO c)
    -> IO c
    -> IO c)
-> (OsString
    -> b
    -> Relative OsString
    -> Basename OsString
    -> FileType
    -> IO (Maybe a))
-> f b
-> IO [a]
listContentsRecFold = \Maybe Int
depthLimit forall c.
OsString
-> b
-> Relative OsString
-> Basename OsString
-> SymlinkType
-> (a -> IO c -> IO c)
-> (IO c -> IO c)
-> IO c
-> IO c
foldDir OsString
-> b
-> Relative OsString
-> Basename OsString
-> FileType
-> IO (Maybe a)
filePred f b
input ->
  Maybe Int
-> (forall c.
    OsString
    -> b
    -> Relative OsString
    -> Basename OsString
    -> FileType
    -> SymlinkType
    -> (a -> IO c -> IO c)
    -> (IO c -> IO c)
    -> IO c
    -> IO c)
-> (OsString
    -> b
    -> Relative OsString
    -> Basename OsString
    -> FileType
    -> IO (Maybe a))
-> f b
-> IO [a]
forall (f :: * -> *) a b.
(Foldable f, Coercible b OsString) =>
Maybe Int
-> (forall c.
    OsString
    -> b
    -> Relative OsString
    -> Basename OsString
    -> FileType
    -> SymlinkType
    -> (a -> IO c -> IO c)
    -> (IO c -> IO c)
    -> IO c
    -> IO c)
-> (OsString
    -> b
    -> Relative OsString
    -> Basename OsString
    -> FileType
    -> IO (Maybe a))
-> f b
-> IO [a]
listContentsRecFold' Maybe Int
depthLimit (\OsString
a b
b Relative OsString
c Basename OsString
d FileType
_f SymlinkType
g a -> IO c -> IO c
h IO c -> IO c
i IO c
j -> OsString
-> b
-> Relative OsString
-> Basename OsString
-> SymlinkType
-> (a -> IO c -> IO c)
-> (IO c -> IO c)
-> IO c
-> IO c
forall c.
OsString
-> b
-> Relative OsString
-> Basename OsString
-> SymlinkType
-> (a -> IO c -> IO c)
-> (IO c -> IO c)
-> IO c
-> IO c
foldDir OsString
a b
b Relative OsString
c Basename OsString
d SymlinkType
g a -> IO c -> IO c
h IO c -> IO c
i IO c
j) OsString
-> b
-> Relative OsString
-> Basename OsString
-> FileType
-> IO (Maybe a)
filePred f b
input

{-# INLINE listContentsRecFold' #-}
-- Actual worker with slightly worse type signature that we don’t want to expose to the users.
-- But it’s better candidate for implementing getDirectoryContentsRecursive here than
-- listContentsRecFold.
listContentsRecFold'
  :: forall f a b. (Foldable f, Coercible b OsPath)
  => Maybe Int
  -> (forall c. OsPath -> b -> Relative OsPath -> Basename OsPath -> FileType -> SymlinkType -> (a -> IO c -> IO c) -> (IO c -> IO c) -> IO c -> IO c)
  -> (OsPath -> b -> Relative OsPath -> Basename OsPath -> FileType -> IO (Maybe a))
  -> f b
  -> IO [a]
listContentsRecFold' :: forall (f :: * -> *) a b.
(Foldable f, Coercible b OsString) =>
Maybe Int
-> (forall c.
    OsString
    -> b
    -> Relative OsString
    -> Basename OsString
    -> FileType
    -> SymlinkType
    -> (a -> IO c -> IO c)
    -> (IO c -> IO c)
    -> IO c
    -> IO c)
-> (OsString
    -> b
    -> Relative OsString
    -> Basename OsString
    -> FileType
    -> IO (Maybe a))
-> f b
-> IO [a]
listContentsRecFold' Maybe Int
depthLimit forall c.
OsString
-> b
-> Relative OsString
-> Basename OsString
-> FileType
-> SymlinkType
-> (a -> IO c -> IO c)
-> (IO c -> IO c)
-> IO c
-> IO c
foldDir OsString
-> b
-> Relative OsString
-> Basename OsString
-> FileType
-> IO (Maybe a)
filePred f b
input =
  DirReadCache -> IO [a]
goCache (DirReadCache -> IO [a]) -> IO DirReadCache -> IO [a]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IO DirReadCache
Raw.allocateDirReadCache
  where
    goCache :: DirReadCache -> IO [a]
goCache DirReadCache
cache =
      (b -> IO [a] -> IO [a]) -> IO [a] -> f b -> IO [a]
forall a b. (a -> b -> b) -> b -> f a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Int -> b -> IO [a] -> IO [a]
goNewDir Int
initLimit) (DirReadCache -> IO ()
Raw.releaseDirReadCache DirReadCache
cache IO () -> IO [a] -> IO [a]
forall a b. IO a -> IO b -> IO b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> [a] -> IO [a]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []) f b
input
      where
        !initLimit :: Int
initLimit = case Maybe Int
depthLimit of
          Maybe Int
Nothing -> -Int
1 -- Loop until overflow, basically infinitely
          Just Int
x  -> Int -> Int
forall a. Num a => a -> a
abs Int
x

        goNewDir :: Int -> b -> IO [a] -> IO [a]
        goNewDir :: Int -> b -> IO [a] -> IO [a]
goNewDir !Int
d b
root IO [a]
rest = do
          stream <- OsString -> IO DirStream
Streaming.openDirStream (OsString -> IO DirStream) -> OsString -> IO DirStream
forall a b. (a -> b) -> a -> b
$ b -> OsString
forall a b. Coercible a b => a -> b
coerce b
root
          goDirStream root d (Streaming.closeDirStream stream *> rest) stream

        goDirStream :: b -> Int -> IO [a] -> DirStream -> IO [a]
        goDirStream :: b -> Int -> IO [a] -> DirStream -> IO [a]
goDirStream b
_    Int
0     IO [a]
rest DirStream
_      = IO [a]
rest
        goDirStream b
root Int
depth IO [a]
rest DirStream
stream = IO [a]
go
          where
            go :: IO [a]
            go :: IO [a]
go = (IO [a] -> IO () -> IO [a]
forall a b. IO a -> IO b -> IO a
`onException` DirStream -> IO ()
Streaming.closeDirStream DirStream
stream) (IO [a] -> IO [a]) -> IO [a] -> IO [a]
forall a b. (a -> b) -> a -> b
$ do
              x <- DirReadCache
-> DirStream -> IO (Maybe (OsString, Basename OsString, FileType))
Streaming.readDirStreamWithCache DirReadCache
cache DirStream
stream
              case x of
                Maybe (OsString, Basename OsString, FileType)
Nothing                -> IO [a]
rest
                Just (OsString
yAbs, Basename OsString
yBase, FileType
ft) -> do
                  let yRel :: Relative OsPath
                      yRel :: Relative OsString
yRel = Basename OsString -> Relative OsString
forall a b. Coercible a b => a -> b
coerce Basename OsString
yBase
                  case FileType
ft of
                    Other SymlinkType
_       -> IO (Maybe a) -> IO [a] -> IO [a]
addLazy (OsString
-> b
-> Relative OsString
-> Basename OsString
-> FileType
-> IO (Maybe a)
filePred OsString
yAbs b
root Relative OsString
yRel Basename OsString
yBase FileType
ft) IO [a]
go
                    File SymlinkType
_        -> IO (Maybe a) -> IO [a] -> IO [a]
addLazy (OsString
-> b
-> Relative OsString
-> Basename OsString
-> FileType
-> IO (Maybe a)
filePred OsString
yAbs b
root Relative OsString
yRel Basename OsString
yBase FileType
ft) IO [a]
go
                    Directory SymlinkType
ft' -> OsString
-> b
-> Relative OsString
-> Basename OsString
-> FileType
-> SymlinkType
-> (a -> IO [a] -> IO [a])
-> (IO [a] -> IO [a])
-> IO [a]
-> IO [a]
forall c.
OsString
-> b
-> Relative OsString
-> Basename OsString
-> FileType
-> SymlinkType
-> (a -> IO c -> IO c)
-> (IO c -> IO c)
-> IO c
-> IO c
foldDir OsString
yAbs b
root Relative OsString
yRel Basename OsString
yBase FileType
ft SymlinkType
ft' a -> IO [a] -> IO [a]
cons (Relative OsString -> Int -> OsString -> IO [a] -> IO [a]
goNewDirAcc Relative OsString
yRel (Int
depth Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) OsString
yAbs) IO [a]
go

            goNewDirAcc :: Relative OsPath -> Int -> OsPath -> IO [a] -> IO [a]
            goNewDirAcc :: Relative OsString -> Int -> OsString -> IO [a] -> IO [a]
goNewDirAcc Relative OsString
rootAcc !Int
d OsString
dir IO [a]
rest1 = do
              stream1 <- OsString -> IO DirStream
Streaming.openDirStream OsString
dir
              goDirStreamAcc rootAcc d (Streaming.closeDirStream stream1 *> rest1) stream1

            goDirStreamAcc :: Relative OsPath -> Int -> IO [a] -> DirStream -> IO [a]
            goDirStreamAcc :: Relative OsString -> Int -> IO [a] -> DirStream -> IO [a]
goDirStreamAcc Relative OsString
_       Int
0      IO [a]
rest1 DirStream
_       = IO [a]
rest1
            goDirStreamAcc Relative OsString
rootAcc Int
depth1 IO [a]
rest1 DirStream
stream1 = IO [a]
go1
              where
                go1 :: IO [a]
                go1 :: IO [a]
go1 = (IO [a] -> IO () -> IO [a]
forall a b. IO a -> IO b -> IO a
`onException` DirStream -> IO ()
Streaming.closeDirStream DirStream
stream1) (IO [a] -> IO [a]) -> IO [a] -> IO [a]
forall a b. (a -> b) -> a -> b
$ do
                  x <- DirReadCache
-> DirStream -> IO (Maybe (OsString, Basename OsString, FileType))
Streaming.readDirStreamWithCache DirReadCache
cache DirStream
stream1
                  case x of
                    Maybe (OsString, Basename OsString, FileType)
Nothing                -> IO [a]
rest1
                    Just (OsString
yAbs, Basename OsString
yBase, FileType
ft) -> do
                      let yRel :: Relative OsPath
                          yRel :: Relative OsString
yRel = (OsString -> OsString -> OsString)
-> Relative OsString -> Basename OsString -> Relative OsString
forall a b. Coercible a b => a -> b
coerce OsString -> OsString -> OsString
(</>) Relative OsString
rootAcc Basename OsString
yBase
                      case FileType
ft of
                        Other SymlinkType
_       -> IO (Maybe a) -> IO [a] -> IO [a]
addLazy (OsString
-> b
-> Relative OsString
-> Basename OsString
-> FileType
-> IO (Maybe a)
filePred OsString
yAbs b
root Relative OsString
yRel Basename OsString
yBase FileType
ft) IO [a]
go1
                        File SymlinkType
_        -> IO (Maybe a) -> IO [a] -> IO [a]
addLazy (OsString
-> b
-> Relative OsString
-> Basename OsString
-> FileType
-> IO (Maybe a)
filePred OsString
yAbs b
root Relative OsString
yRel Basename OsString
yBase FileType
ft) IO [a]
go1
                        Directory SymlinkType
ft' -> OsString
-> b
-> Relative OsString
-> Basename OsString
-> FileType
-> SymlinkType
-> (a -> IO [a] -> IO [a])
-> (IO [a] -> IO [a])
-> IO [a]
-> IO [a]
forall c.
OsString
-> b
-> Relative OsString
-> Basename OsString
-> FileType
-> SymlinkType
-> (a -> IO c -> IO c)
-> (IO c -> IO c)
-> IO c
-> IO c
foldDir OsString
yAbs b
root Relative OsString
yRel Basename OsString
yBase FileType
ft SymlinkType
ft' a -> IO [a] -> IO [a]
cons (Relative OsString -> Int -> OsString -> IO [a] -> IO [a]
goNewDirAcc Relative OsString
yRel (Int
depth1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) OsString
yAbs) IO [a]
go1

        addLazy :: IO (Maybe a) -> IO [a] -> IO [a]
        addLazy :: IO (Maybe a) -> IO [a] -> IO [a]
addLazy IO (Maybe a)
x IO [a]
y = do
          x' <- IO (Maybe a)
x
          case x' of
            Maybe a
Nothing  -> IO [a]
y
            Just a
x'' -> a -> IO [a] -> IO [a]
cons a
x'' IO [a]
y

        cons :: a -> IO [a] -> IO [a]
        cons :: a -> IO [a] -> IO [a]
cons a
x IO [a]
y =
          (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
:) ([a] -> [a]) -> IO [a] -> IO [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO [a] -> IO [a]
forall a. IO a -> IO a
unsafeInterleaveIO IO [a]
y