1{-# LANGUAGE StandaloneDeriving #-} 2 3module TreeDB( 4 DirList, 5 dlEmpty, dlByExt, dlByExts, dlAdd, dlAddByExt, 6 7 TreeDB, 8 tdbEmpty, tdbByDir, tdbAdd, tdbAddDir, 9 tdbBuild, tdbMerge, 10 11 tdbByDirExt, tdbByDirExts 12 ) 13where 14 15import qualified Data.ByteString.Char8 as C 16import Data.List 17import Data.Trie(Trie) 18import qualified Data.Trie as T 19import Data.Typeable 20 21import System.FilePath 22 23-- 24-- The files in a directory, partitioned by extension. 25-- 26type DirList = [(String, [String])] 27 28dlEmpty :: DirList 29dlEmpty = [] 30 31-- Linear search for files by extension, in a single directory. 32dlByExt :: String -> DirList -> [String] 33dlByExt _ [] = [] 34dlByExt ext ((ext', names) : dirlist) 35 | ext' == ext = [n <.> ext' | n <- names] 36 | otherwise = dlByExt ext dirlist 37 38-- Search for multiple extensions at once. 'exts' must be sorted, with no 39-- duplicates. 40dlByExts :: [String] -> DirList -> [String] 41dlByExts _ [] = [] 42dlByExts [] _ = [] 43dlByExts (ext:exts) ((ext', names):dirlist) = 44 case compare ext ext' of 45 -- 'ext' isn't in the list. 46 LT -> dlByExts exts ((ext', names):dirlist) 47 -- 'ext' is right here. 48 EQ -> [n <.> ext' | n <- names] ++ dlByExts exts dirlist 49 -- 'ext' may be in the remainder. Nothing else can match here. 50 GT -> dlByExts (ext:exts) dirlist 51 52-- Insert a file, given its extension. Again linear. 53dlAdd :: FilePath -> DirList -> DirList 54dlAdd file dirList = 55 dlAddByExt (takeExtension file) (dropExtension file) dirList 56 57-- Keeps the list sorted by extension 58dlAddByExt :: String -> String -> DirList -> DirList 59dlAddByExt ext name [] = [(ext, [name])] 60dlAddByExt ext name ((ext', names):dirlist) = 61 case compare ext ext' of 62 LT -> (ext, [name]):(ext', names):dirlist 63 EQ -> (ext', name:names):dirlist 64 GT -> (ext', names):(dlAddByExt ext name dirlist) 65 66-- 67-- A map from directory to contents, excluding subdirectories. 68-- 69type TreeDB = Trie DirList 70 71deriving instance Typeable Trie 72 73tdbEmpty :: TreeDB 74tdbEmpty = T.empty 75 76-- Get directory contents by directory path 77tdbByDir :: FilePath -> TreeDB -> Maybe DirList 78tdbByDir path treeDB = T.lookup (C.pack path) treeDB 79 80-- Add a file 81tdbAdd :: FilePath -> TreeDB -> TreeDB 82tdbAdd path treeDB 83 | T.member dirS treeDB = 84 T.adjust (\dirList -> dlAdd file dirList) dirS treeDB 85 | otherwise = 86 T.insert dirS (dlAdd file dlEmpty) treeDB 87 where 88 dir = takeDirectory path 89 file = takeFileName path 90 dirS = C.pack dir 91 92-- Add a directory, complete with (relative) contents 93tdbAddDir :: FilePath -> [FilePath] -> TreeDB -> TreeDB 94tdbAddDir dir files treeDB 95 | T.member dirS treeDB = 96 T.adjust (\dirList -> foldr dlAdd dirList files) dirS treeDB 97 | otherwise = 98 T.insert dirS (foldr dlAdd dlEmpty files) treeDB 99 where 100 dirS = C.pack dir 101 102tdbBuild :: [FilePath] -> TreeDB 103tdbBuild files = foldr tdbAdd tdbEmpty files 104 105tdbMerge :: TreeDB -> TreeDB -> TreeDB 106tdbMerge = T.unionL 107 108-- 109-- Combined queries 110-- 111 112-- Find files by directory and extension 113tdbByDirExt :: FilePath -> String -> TreeDB -> Maybe [FilePath] 114tdbByDirExt path ext treeDB = do 115 dirList <- tdbByDir path treeDB 116 let filenames = dlByExt ext dirList 117 return [ path </> file | file <- filenames ] 118 119-- Look for multiple extensions. 'exts' need not be sorted. 120tdbByDirExts :: FilePath -> [String] -> TreeDB -> Maybe [FilePath] 121tdbByDirExts path exts treeDB = do 122 dirList <- tdbByDir path treeDB 123 let filenames = dlByExts (sort exts) dirList 124 return [ path </> file | file <- filenames ] 125