1 files changed, 71 insertions, 0 deletions

diff --git a/src/Data/Torrent/Tree.hs b/src/Data/Torrent/Tree.hs
new file mode 100644
index 00000000..e9a337a1
--- /dev/null
+++ b/src/Data/Torrent/Tree.hs
@@ -0,0 +1,71 @@
+-- |
+--   Copyright   :  (c) Sam Truzjan 2013
+--   License     :  BSD3
+--   Maintainer  :  pxqr.sta@gmail.com
+--   Stability   :  experimental
+--   Portability :  portable
+--
+{-# LANGUAGE FlexibleInstances  #-}
+{-# LANGUAGE TemplateHaskell    #-}
+{-# LANGUAGE DeriveDataTypeable #-}
+module Data.Torrent.Tree
+       ( DirTree (..)
+       , build
+
+       , Data.Torrent.Tree.lookup
+       , lookupDir
+
+       , fileCount
+       , dirCount
+       ) where
+
+import Control.Arrow
+import Data.ByteString as BS
+import Data.ByteString.Char8 as BC
+import Data.Foldable
+import Data.List as L
+import Data.Map  as M
+import Data.Monoid
+
+import Data.Torrent.Layout
+
+
+data DirTree a = Dir  { children :: Map ByteString (DirTree a) }
+               | File { node     :: FileInfo a                 }
+                 deriving Show
+
+build :: LayoutInfo -> DirTree ()
+build SingleFile {liFile = FileInfo {..}} = Dir
+    { children = M.singleton fiName (File fi) }
+  where
+    fi = FileInfo fiLength fiMD5Sum ()
+build MultiFile {..} = Dir $ M.singleton liDirName files
+  where
+    files = Dir $ M.fromList $ L.map mkFileEntry liFiles
+    mkFileEntry FileInfo {..} = (L.head fiName, ent) -- TODO FIXME
+      where
+        ent = File $ FileInfo fiLength fiMD5Sum ()
+
+decompress :: DirTree () -> [FileInfo ()]
+decompress = undefined
+
+lookup :: [FilePath] -> DirTree a -> Maybe (DirTree a)
+lookup []        t      = Just t
+lookup (p : ps) (Dir m) | Just subTree <- M.lookup (BC.pack p) m
+                        = Data.Torrent.Tree.lookup ps subTree
+lookup _         _      = Nothing
+
+lookupDir :: [FilePath] -> DirTree a -> Maybe [(ByteString, DirTree a)]
+lookupDir ps d
+  | Just subTree <- Data.Torrent.Tree.lookup ps d =
+    case subTree of
+      File _  -> Nothing
+      Dir  es -> Just $ M.toList es
+
+fileCount :: DirTree a -> Sum Int
+fileCount File {..} = Sum 1
+fileCount Dir  {..} = foldMap fileCount children
+
+dirCount :: DirTree a -> Sum Int
+dirCount File {..} = Sum 0
+dirCount Dir  {..} = Sum 1 <> foldMap dirCount children