-- |
-- 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

  , listContentsRecFold
  ) where

import Control.Exception (mask, 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 :: 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)

{-# 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 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 that
-- 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 =
  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 =
          ((forall a. IO a -> IO a) -> IO [a]) -> IO [a]
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask (((forall a. IO a -> IO a) -> IO [a]) -> IO [a])
-> ((forall a. IO a -> IO a) -> IO [a]) -> IO [a]
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
restore -> 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
            (IO [a] -> IO [a]
forall a. IO a -> IO a
restore
              (b -> Int -> IO [a] -> DirStream -> IO [a]
goDirStream 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))

        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
              Maybe (OsPath, Basename OsPath, FileType)
x <- DirReadCache
-> DirStream -> IO (Maybe (OsPath, Basename OsPath, FileType))
Streaming.readDirStreamWithCache DirReadCache
cache 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]
goNewDirAcc Relative OsPath
yRel (Int
depth Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) OsPath
yAbs) IO [a]
go

            goNewDirAcc :: Relative OsPath -> Int -> OsPath -> IO [a] -> IO [a]
            goNewDirAcc :: Relative OsPath -> Int -> OsPath -> IO [a] -> IO [a]
goNewDirAcc Relative OsPath
rootAcc !Int
d OsPath
dir IO [a]
rest1 =
              ((forall a. IO a -> IO a) -> IO [a]) -> IO [a]
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask (((forall a. IO a -> IO a) -> IO [a]) -> IO [a])
-> ((forall a. IO a -> IO a) -> IO [a]) -> IO [a]
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
restore -> do
                DirStream
stream1 <- OsPath -> IO DirStream
Streaming.openDirStream OsPath
dir
                (IO [a] -> IO [a]
forall a. IO a -> IO a
restore
                  (Relative OsPath -> Int -> IO [a] -> DirStream -> IO [a]
goDirStreamAcc 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))

            goDirStreamAcc :: Relative OsPath -> Int -> IO [a] -> DirStream -> IO [a]
            goDirStreamAcc :: Relative OsPath -> Int -> IO [a] -> DirStream -> IO [a]
goDirStreamAcc Relative OsPath
_       Int
0      IO [a]
rest1 DirStream
_       = IO [a]
rest1
            goDirStreamAcc 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 <- DirReadCache
-> DirStream -> IO (Maybe (OsPath, Basename OsPath, FileType))
Streaming.readDirStreamWithCache DirReadCache
cache 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
(</>) 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]
goNewDirAcc 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