gregorycollins · August 29, 2015 14:01 · May 12, 2014 · May 12, 2014 · May 12, 2014
diff --git a/Traversal.hs b/Traversal.hs
@@ -70,6 +70,7 @@ nextDir dt@(DirectoryTraversal ref) = E.mask $ \restore -> go restore
               s <- restore $ D.readDirStream $ levelDirStream l
               if S.null s
                 then do atomicModifyIORef ref $ \ls -> (tl ls, ())
+                        D.closeDirStream $ levelDirStream l
                         go restore
                 else if s == "." || s == ".."
                        then go restore

diff --git a/.gitignore b/.gitignore
@@ -0,0 +1,6 @@
+#*#
+*~
+.cabal-sandbox
+TAGS
+cabal.sandbox.config
+dist/
diff --git a/Setup.hs b/Setup.hs
@@ -0,0 +1,2 @@
+import Distribution.Simple
+main = defaultMain
diff --git a/Traversal.hs b/Traversal.hs
@@ -1,22 +1,96 @@
+{-# LANGUAGE BangPatterns      #-}
 {-# LANGUAGE OverloadedStrings #-}
 
 module Main where
 
+import           Control.Applicative               ((<$>))
 import qualified Control.Exception                 as E
-import           Control.Monad                     ((>=>))
+import           Control.Monad                     (mapM_, when, (>=>))
 import qualified Data.ByteString.Char8             as S
+import           Data.IORef                        (IORef, atomicModifyIORef,
+                                                    newIORef, readIORef)
 import           System.IO.Streams                 (InputStream)
 import qualified System.IO.Streams                 as Streams
 import           System.Posix.ByteString.FilePath  (RawFilePath)
 import qualified System.Posix.Directory.ByteString as D
+import qualified System.Posix.Files.ByteString     as D
 
-traverseDirectory :: RawFilePath -> (InputStream RawFilePath -> IO a) -> IO a
-traverseDirectory fp m = E.bracket (D.openDirStream fp) D.closeDirStream go
+------------------------------------------------------------------------------
+traverseDirectoryRecursive :: RawFilePath
+                           -> (InputStream RawFilePath -> IO a)
+                           -> IO a
+traverseDirectoryRecursive fp m = E.bracket (newDirectoryTraversal fp)
+                                            deleteDirectoryTraversal
+                                            go
   where
-    go d = Streams.makeInputStream (readDirStr d) >>= m
-    readDirStr d = do s <- D.readDirStream d
-                      return $! if S.null s then Nothing else Just s
+    go dt = Streams.makeInputStream (nextDir dt)
+              >>= Streams.lockingInputStream
+              >>= m
 
+
+------------------------------------------------------------------------------
 main :: IO ()
-main = traverseDirectory "." $ Streams.map (`S.append` "\n") >=>
-                               Streams.connectTo Streams.stdout
+main = traverseDirectoryRecursive "." $ Streams.map (`S.append` "\n") >=>
+                                        Streams.connectTo Streams.stdout
+
+
+------------------------------------------------------------------------------
+data Level = Level { levelParent    :: RawFilePath
+                   , levelDirStream :: D.DirStream
+                   }
+
+newtype DirectoryTraversal = DirectoryTraversal (IORef [Level])
+
+newDirectoryTraversal :: RawFilePath -> IO DirectoryTraversal
+newDirectoryTraversal fp = E.mask_ $ do
+    dt <- DirectoryTraversal <$> newIORef []
+    recurseInto fp dt
+    return dt
+
+deleteDirectoryTraversal :: DirectoryTraversal -> IO ()
+deleteDirectoryTraversal (DirectoryTraversal ref) =
+    E.mask_ $ do
+        readIORef ref >>= mapM_ (D.closeDirStream . levelDirStream)
+        writeIORef ref []
+
+recurseInto :: RawFilePath -> DirectoryTraversal -> IO ()
+recurseInto fp (DirectoryTraversal ref) = E.mask_ $ do
+    d <- D.openDirStream fp
+    let lvl = Level fp d
+    atomicModifyIORef ref $ \l -> ((lvl:l), ())
+
+nextDir :: DirectoryTraversal -> IO (Maybe RawFilePath)
+nextDir dt@(DirectoryTraversal ref) = E.mask $ \restore -> go restore
+  where
+    go restore = do
+        lvls <- readIORef ref
+        case lvls of
+          []     -> return Nothing
+          (!l:_) -> do
+              s <- restore $ D.readDirStream $ levelDirStream l
+              if S.null s
+                then do atomicModifyIORef ref $ \ls -> (tl ls, ())
+                        go restore
+                else if s == "." || s == ".."
+                       then go restore
+                       else entry (levelParent l) s
+
+    -- I'm sure there must be a version of this for RawFilePath elsewhere
+    infixr 5 </>
+    a </> b = if a == ""
+                then b
+                else let a'  = fst $! S.spanEnd (== '/') a
+                         a'' = if S.null a' then "/" else a'
+                     in S.concat [a'', "/", b]
+
+    entry parent fp = do
+        let fullPath = parent </> fp
+        dir <- isDir fullPath
+        when dir $ recurseInto fullPath dt
+        return $! Just fullPath
+
+    isDir fp = do s <- D.getFileStatus fp
+                  return $! D.isDirectory s
+
+    tl []     = []
+    tl (_:xs) = xs
diff --git a/directory-traversal.cabal b/directory-traversal.cabal
@@ -0,0 +1,22 @@
+-- Initial directory-traversal.cabal generated by cabal init.  For further
+-- documentation, see http://haskell.org/cabal/users-guide/
+
+name:                directory-traversal
+version:             0.1.0.0
+-- synopsis:
+-- description:
+-- license:
+-- license-file:        LICENSE
+author:              Gregory Collins
+maintainer:          [email protected]
+-- copyright:
+-- category:
+build-type:          Simple
+-- extra-source-files:
+cabal-version:       >=1.10
+
+executable directory-traversal
+  main-is:             Traversal.hs
+  other-extensions:    OverloadedStrings
+  build-depends:       base >=4.5 && <4.6, bytestring >=0.9 && <1.2, io-streams >=1.1 && <1.2, unix >=2.5 && <2.8
+  default-language:    Haskell2010
diff --git a/Traversal.hs b/Traversal.hs
@@ -0,0 +1,22 @@
+{-# LANGUAGE OverloadedStrings #-}
+
+module Main where
+
+import qualified Control.Exception                 as E
+import           Control.Monad                     ((>=>))
+import qualified Data.ByteString.Char8             as S
+import           System.IO.Streams                 (InputStream)
+import qualified System.IO.Streams                 as Streams
+import           System.Posix.ByteString.FilePath  (RawFilePath)
+import qualified System.Posix.Directory.ByteString as D
+
+traverseDirectory :: RawFilePath -> (InputStream RawFilePath -> IO a) -> IO a
+traverseDirectory fp m = E.bracket (D.openDirStream fp) D.closeDirStream go
+  where
+    go d = Streams.makeInputStream (readDirStr d) >>= m
+    readDirStr d = do s <- D.readDirStream d
+                      return $! if S.null s then Nothing else Just s
+
+main :: IO ()
+main = traverseDirectory "." $ Streams.map (`S.append` "\n") >=>
+                               Streams.connectTo Streams.stdout