Day 7 - Foldable instance for summing up directories

src/Day7.hs

@@ -13,15 +13,12 @@ day7 = do
    let parsed = forceRight $ parseTerminalLines input
    let tree = buildTree parsed
    let small = filterDirectories (<= 100000) tree
-   let sumOfSmall = sumUp small
+   let sumOfSmall = sum small
    putStrLn ("Sum of those directories is " ++ show sumOfSmall)
-   let rootSize = snd . head $ sizeOfDirectories tree
-   putStrLn ("Root size is " ++ show rootSize)
+   let rootSize = head $ sizeOfDirectories tree
    let freeSpace = 70000000 - rootSize
    let requiredSpace = 30000000 - freeSpace
-   putStrLn ("Required space is " ++ show requiredSpace)
    let bigEnough = filterDirectories (>= requiredSpace) tree
-   let onlySizes = map snd bigEnough
-   let smallest = minimum onlySizes
+   let smallest = minimum bigEnough
    putStrLn ("Smallest directory that is still large enough " ++ show smallest)
 

src/Day7/Interpreter.hs

@@ -1,29 +1,40 @@
 module Day7.Interpreter (
   buildTree,
   mkdir,
-  Directory (..),
+  FileTree (..),
+  Directory,
   calculateSize,
   filterDirectories,
   sizeOfDirectories,
-  sumUp,
 ) where
 
 import Day7.Parser
 import Data.Map (Map)
 import qualified Data.Map as Map
+import qualified Data.Foldable as F
 
-data Directory = Directory
-  { sub :: Map String Directory
-  , files :: Map String Int
+
+data FileTree a = FileTree
+  { sub :: Map String (FileTree a)
+  , files :: Map String a
   , isRoot :: Bool
 } deriving (Eq)
 
+-- Directory is a FileTree with file sizes of Int
+type Directory = FileTree Int
 
-instance Show Directory where
+-- pretty much only useful for summing up file sizes 
+instance F.Foldable FileTree where
+  foldMap f dir = foldMap f (files dir) <> mconcat mappedSubfolders
+    where
+      subFolders = Map.elems $ sub dir
+      mappedSubfolders = map (foldMap f) subFolders
+
+instance Show a => Show (FileTree a) where
   show d = "DIR " ++ show (Map.toList $ files d) ++ " - " ++ show (Map.toList $ sub d) ++ "\n"
 
 mkdir :: Directory
-mkdir = Directory { sub = Map.empty, files = Map.empty, isRoot = False }
+mkdir = FileTree { sub = Map.empty, files = Map.empty, isRoot = False }
 
 rootDirectory :: Directory
 rootDirectory = mkdir { isRoot = True  }
@@ -57,24 +68,14 @@ toFiles (_:rest) = toFiles rest
 toFiles [] = []
 
 calculateSize :: Directory -> Int
-calculateSize dir = sum (Map.elems (files dir)) + sum sizes
-  where
-    subFolders = Map.elems (sub dir)
-    sizes = map calculateSize subFolders
+calculateSize = F.foldl' (+) 0
 
-flatten :: Directory -> [(String, Directory)]
-flatten dir = flatten' ("/", dir)
-  where
-  flatten' (name, d) = (name, d) : concatMap flatten' (Map.toList (sub d))
+flatten :: Directory -> [Directory]
+flatten d = d : concatMap flatten (Map.elems (sub d))
 
-sizeOfDirectories :: Directory -> [(String, Int)]
-sizeOfDirectories dir = map withSize allDirectories
-  where
-    allDirectories = flatten dir
-    withSize (name, directory) = (name, calculateSize directory)
+sizeOfDirectories :: Directory -> [Int]
+sizeOfDirectories = map calculateSize . flatten
 
-filterDirectories :: (Int -> Bool) -> Directory -> [(String, Int)]
-filterDirectories predicate dir = filter (\(_, size) -> predicate size) $ sizeOfDirectories dir
+filterDirectories :: (Int -> Bool) -> Directory -> [Int]
+filterDirectories predicate dir = filter predicate $ sizeOfDirectories dir
 
-sumUp :: [(String, Int)] -> Int
-sumUp = sum . map snd

test/Day7Spec.hs

@@ -126,5 +126,5 @@ spec =
            let parsed = forceRight $ parseTerminalLines inputPart1
            let tree = buildTree parsed
            let filtered = filterDirectories (<= 100000) tree
-           let summed = sumUp filtered
+           let summed = sum filtered
            summed `shouldBe` 95437