-- |
-- 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 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 :: OsPath -> IO [(OsPath, FileType)]
getDirectoryContentsRecursive OsPath
root =
  Maybe Int
-> (forall c.
    OsPath
    -> OsPath
    -> Relative OsPath
    -> Basename OsPath
    -> FileType
    -> SymlinkType
    -> ((OsPath, FileType) -> IO c -> IO c)
    -> (IO c -> IO c)
    -> IO c
    -> IO c)
-> (OsPath
    -> OsPath
    -> Relative OsPath
    -> Basename OsPath
    -> FileType
    -> IO (Maybe (OsPath, FileType)))
-> Maybe OsPath
-> IO [(OsPath, FileType)]
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'
    Maybe Int
forall a. Maybe a
Nothing
    (\OsPath
_ OsPath
_ (Relative OsPath
path) Basename OsPath
_ FileType
ft SymlinkType
_ (OsPath, FileType) -> IO c -> IO c
cons IO c -> IO c
prependSubdir IO c
rest -> (OsPath, FileType) -> IO c -> IO c
cons (OsPath
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)
    (\OsPath
_ OsPath
_ (Relative OsPath
path) Basename OsPath
_ FileType
ft -> Maybe (OsPath, FileType) -> IO (Maybe (OsPath, FileType))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((OsPath, FileType) -> Maybe (OsPath, FileType)
forall a. a -> Maybe a
Just (OsPath
path, FileType
ft)))
    (OsPath -> Maybe OsPath
forall a. a -> Maybe a
Just OsPath
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 OsPath -> SymlinkType -> Bool)
-> (Basename OsPath -> Bool) -> OsPath -> IO [(OsPath, FileType)]
getDirectoryContentsWithFilterRecursive Basename OsPath -> SymlinkType -> Bool
visitPred Basename OsPath -> Bool
collectPred OsPath
root =
  Maybe Int
-> (forall c.
    OsPath
    -> OsPath
    -> Relative OsPath
    -> Basename OsPath
    -> FileType
    -> SymlinkType
    -> ((OsPath, FileType) -> IO c -> IO c)
    -> (IO c -> IO c)
    -> IO c
    -> IO c)
-> (OsPath
    -> OsPath
    -> Relative OsPath
    -> Basename OsPath
    -> FileType
    -> IO (Maybe (OsPath, FileType)))
-> Maybe OsPath
-> IO [(OsPath, FileType)]
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'
    Maybe Int
forall a. Maybe a
Nothing
    (\OsPath
_ OsPath
_ (Relative OsPath
path) Basename OsPath
basename FileType
ft SymlinkType
symlink (OsPath, FileType) -> IO c -> IO c
cons IO c -> IO c
prependSubdir IO c
rest ->
       (if Basename OsPath -> Bool
collectPred Basename OsPath
basename then (OsPath, FileType) -> IO c -> IO c
cons (OsPath
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 OsPath -> SymlinkType -> Bool
visitPred Basename OsPath
basename SymlinkType
symlink
         then IO c -> IO c
prependSubdir IO c
rest
         else IO c
rest)
    (\OsPath
_ OsPath
_ (Relative OsPath
path) Basename OsPath
basename FileType
ft ->
      Maybe (OsPath, FileType) -> IO (Maybe (OsPath, FileType))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (OsPath, FileType) -> IO (Maybe (OsPath, FileType)))
-> Maybe (OsPath, FileType) -> IO (Maybe (OsPath, FileType))
forall a b. (a -> b) -> a -> b
$
        if Basename OsPath -> Bool
collectPred Basename OsPath
basename
        then (OsPath, FileType) -> Maybe (OsPath, FileType)
forall a. a -> Maybe a
Just (OsPath
path, FileType
ft)
        else Maybe (OsPath, FileType)
forall a. Maybe a
Nothing)
    (OsPath -> Maybe OsPath
forall a. a -> Maybe a
Just OsPath
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@              - full path to the visited directory relative to root we’re searching in. If the root was absolute then this path would be too, if it was relative then this path would be relative to the same target.
  -- * @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 OsPath) =>
Maybe Int
-> (forall c.
    OsPath
    -> b
    -> Relative OsPath
    -> Basename OsPath
    -> 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 = \Maybe Int
depthLimit forall c.
OsPath
-> b
-> Relative OsPath
-> Basename OsPath
-> SymlinkType
-> (a -> IO c -> IO c)
-> (IO c -> IO c)
-> IO c
-> IO c
foldDir OsPath
-> b
-> Relative OsPath
-> Basename OsPath
-> FileType
-> IO (Maybe a)
filePred f b
input ->
  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]
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' Maybe Int
depthLimit (\OsPath
a b
b Relative OsPath
c Basename OsPath
d FileType
_f SymlinkType
g a -> IO c -> IO c
h IO c -> IO c
i IO c
j -> OsPath
-> b
-> Relative OsPath
-> Basename OsPath
-> SymlinkType
-> (a -> IO c -> IO c)
-> (IO c -> IO c)
-> IO c
-> IO c
forall c.
OsPath
-> b
-> Relative OsPath
-> Basename OsPath
-> SymlinkType
-> (a -> IO c -> IO c)
-> (IO c -> IO c)
-> IO c
-> IO c
foldDir OsPath
a b
b Relative OsPath
c Basename OsPath
d SymlinkType
g a -> IO c -> IO c
h IO c -> IO c
i IO c
j) OsPath
-> b
-> Relative OsPath
-> Basename OsPath
-> 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 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' Maybe Int
depthLimit forall c.
OsPath
-> b
-> Relative OsPath
-> Basename OsPath
-> FileType
-> SymlinkType
-> (a -> IO c -> IO c)
-> (IO c -> IO c)
-> IO c
-> IO c
foldDir OsPath
-> b
-> Relative OsPath
-> Basename OsPath
-> FileType
-> IO (Maybe a)
filePred f b
input =
  (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) ([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
      DirStream
stream <- OsPath -> IO DirStream
Streaming.openDirStream (OsPath -> IO DirStream) -> OsPath -> IO DirStream
forall a b. (a -> b) -> a -> b
$ b -> OsPath
forall a b. Coercible a b => a -> b
coerce b
root
      b -> Int -> IO [a] -> DirStream -> IO [a]
goToplevelDirStream b
root Int
d (DirStream -> IO ()
Streaming.closeDirStream DirStream
stream 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
*> IO [a]
rest) DirStream
stream

    goToplevelDirStream :: b -> Int -> IO [a] -> DirStream -> IO [a]
    goToplevelDirStream :: b -> Int -> IO [a] -> DirStream -> IO [a]
goToplevelDirStream b
_    Int
0     IO [a]
rest DirStream
_      = IO [a]
rest
    goToplevelDirStream 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
          Maybe (OsPath, Basename OsPath, FileType)
x <- DirStream -> IO (Maybe (OsPath, Basename OsPath, FileType))
Streaming.readDirStreamFull DirStream
stream
          case Maybe (OsPath, Basename OsPath, FileType)
x of
            Maybe (OsPath, Basename OsPath, FileType)
Nothing                -> IO [a]
rest
            Just (OsPath
yAbs, Basename OsPath
yBase, FileType
ft) -> do
              let yRel :: Relative OsPath
                  yRel :: Relative OsPath
yRel = Basename OsPath -> Relative OsPath
forall a b. Coercible a b => a -> b
coerce Basename OsPath
yBase
              case FileType
ft of
                Other SymlinkType
_       -> IO (Maybe a) -> IO [a] -> IO [a]
addLazy (OsPath
-> b
-> Relative OsPath
-> Basename OsPath
-> FileType
-> IO (Maybe a)
filePred OsPath
yAbs b
root Relative OsPath
yRel Basename OsPath
yBase FileType
ft) IO [a]
go
                File SymlinkType
_        -> IO (Maybe a) -> IO [a] -> IO [a]
addLazy (OsPath
-> b
-> Relative OsPath
-> Basename OsPath
-> FileType
-> IO (Maybe a)
filePred OsPath
yAbs b
root Relative OsPath
yRel Basename OsPath
yBase FileType
ft) IO [a]
go
                Directory SymlinkType
ft' -> OsPath
-> b
-> Relative OsPath
-> Basename OsPath
-> FileType
-> SymlinkType
-> (a -> IO [a] -> IO [a])
-> (IO [a] -> IO [a])
-> IO [a]
-> IO [a]
forall c.
OsPath
-> b
-> Relative OsPath
-> Basename OsPath
-> FileType
-> SymlinkType
-> (a -> IO c -> IO c)
-> (IO c -> IO c)
-> IO c
-> IO c
foldDir OsPath
yAbs b
root Relative OsPath
yRel Basename OsPath
yBase FileType
ft SymlinkType
ft' a -> IO [a] -> IO [a]
cons (Relative OsPath -> Int -> OsPath -> IO [a] -> IO [a]
goChildDirAcc Relative OsPath
yRel (Int
depth Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) OsPath
yAbs) IO [a]
go

        goChildDirAcc :: Relative OsPath -> Int -> OsPath -> IO [a] -> IO [a]
        goChildDirAcc :: Relative OsPath -> Int -> OsPath -> IO [a] -> IO [a]
goChildDirAcc Relative OsPath
rootAcc !Int
d OsPath
dir IO [a]
rest1 = do
          DirStream
stream1 <- OsPath -> IO DirStream
Streaming.openDirStream OsPath
dir
          Relative OsPath -> Int -> IO [a] -> DirStream -> IO [a]
goChildDirStreamAcc ((OsPath -> OsPath) -> Relative OsPath -> Relative OsPath
forall a b. Coercible a b => a -> b
coerce OsPath -> OsPath
addTrailingPathSeparator Relative OsPath
rootAcc) Int
d (DirStream -> IO ()
Streaming.closeDirStream DirStream
stream1 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
*> IO [a]
rest1) DirStream
stream1

        goChildDirStreamAcc :: Relative OsPath -> Int -> IO [a] -> DirStream -> IO [a]
        goChildDirStreamAcc :: Relative OsPath -> Int -> IO [a] -> DirStream -> IO [a]
goChildDirStreamAcc Relative OsPath
_       Int
0      IO [a]
rest1 DirStream
_       = IO [a]
rest1
        goChildDirStreamAcc Relative OsPath
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
              Maybe (OsPath, Basename OsPath, FileType)
x <- DirStream -> IO (Maybe (OsPath, Basename OsPath, FileType))
Streaming.readDirStreamFull DirStream
stream1
              case Maybe (OsPath, Basename OsPath, FileType)
x of
                Maybe (OsPath, Basename OsPath, FileType)
Nothing                -> IO [a]
rest1
                Just (OsPath
yAbs, Basename OsPath
yBase, FileType
ft) -> do
                  let yRel :: Relative OsPath
                      yRel :: Relative OsPath
yRel = (OsPath -> OsPath -> OsPath)
-> Relative OsPath -> Basename OsPath -> Relative OsPath
forall a b. Coercible a b => a -> b
coerce (OsPath -> OsPath -> OsPath
forall a. Semigroup a => a -> a -> a
(<>) :: OsPath -> OsPath -> OsPath) Relative OsPath
rootAcc Basename OsPath
yBase
                  case FileType
ft of
                    Other SymlinkType
_       -> IO (Maybe a) -> IO [a] -> IO [a]
addLazy (OsPath
-> b
-> Relative OsPath
-> Basename OsPath
-> FileType
-> IO (Maybe a)
filePred OsPath
yAbs b
root Relative OsPath
yRel Basename OsPath
yBase FileType
ft) IO [a]
go1
                    File SymlinkType
_        -> IO (Maybe a) -> IO [a] -> IO [a]
addLazy (OsPath
-> b
-> Relative OsPath
-> Basename OsPath
-> FileType
-> IO (Maybe a)
filePred OsPath
yAbs b
root Relative OsPath
yRel Basename OsPath
yBase FileType
ft) IO [a]
go1
                    Directory SymlinkType
ft' -> OsPath
-> b
-> Relative OsPath
-> Basename OsPath
-> FileType
-> SymlinkType
-> (a -> IO [a] -> IO [a])
-> (IO [a] -> IO [a])
-> IO [a]
-> IO [a]
forall c.
OsPath
-> b
-> Relative OsPath
-> Basename OsPath
-> FileType
-> SymlinkType
-> (a -> IO c -> IO c)
-> (IO c -> IO c)
-> IO c
-> IO c
foldDir OsPath
yAbs b
root Relative OsPath
yRel Basename OsPath
yBase FileType
ft SymlinkType
ft' a -> IO [a] -> IO [a]
cons (Relative OsPath -> Int -> OsPath -> IO [a] -> IO [a]
goChildDirAcc Relative OsPath
yRel (Int
depth1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) OsPath
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
      Maybe a
x' <- IO (Maybe a)
x
      case Maybe a
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