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