advent07/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2022/12/07/advent-of-code-2022-day-7/
   2
   3 import AoC
   4 import Data.Text (Text)
   5 import qualified Data.Text.IO as TIO
   6 import Data.Attoparsec.Text hiding (take)
   7 import Control.Applicative
   8 import Data.Char
   9 import Data.Maybe
  10 import Data.Tree
  11 import Data.Tree.Zipper hiding (tree)
  12 import qualified Data.Map.Strict as M
  13 import Data.List (foldl', sort)
  14
  15 data ParsedObject = CD String
  16                   | LS
  17                   | PDirectory String
  18                   | PFile Int String
  19                   deriving Show
  20
  21 data Directory = Dir String (M.Map String Int)
  22   deriving (Show, Eq)
  23
  24 data ContainedSize = CSize String Integer
  25   deriving (Show, Eq)
  26
  27 type DTree = Tree Directory
  28 type ZDTree = TreePos Full Directory
  29 type STree = Tree ContainedSize
  30
  31 reportingThreshold, spaceAvailable, spaceRequired :: Integer
  32 reportingThreshold =   100000
  33 spaceAvailable     = 70000000
  34 spaceRequired      = 30000000
  35
  36 main :: IO ()
  37 main =
  38   do  dataFileName <- getDataFileName
  39       text <- TIO.readFile dataFileName
  40       let trace = successfulParse text
  41       let tree = mkTree trace emptyTree
  42       let sizedTree = transitiveSizes $ containingSizes tree
  43       print $ part1 sizedTree
  44       print $ part2 sizedTree
  45
  46 part1, part2 :: STree -> Integer
  47 part1 = foldTree (\x xs -> sum (x:xs)) . fmap cancelLarge
  48
  49 part2 tree = spaceFreed
  50   where nodes = fmap extractCSize $ flatten tree
  51         spaceUsed = extractCSize $ rootLabel tree
  52         spaceUnused = spaceAvailable - spaceUsed
  53         spaceToFree = spaceRequired - spaceUnused
  54         viableNodes = filter (>= spaceToFree) nodes
  55         spaceFreed = head $ sort viableNodes
  56
  57 emptyTree :: DTree
  58 emptyTree = Node {rootLabel = (Dir "/" M.empty), subForest = []}
  59
  60 mkTree :: [ParsedObject] -> DTree -> DTree
  61 mkTree trace tree = toTree $ root $ makeTree trace $ fromTree tree
  62
  63 makeTree :: [ParsedObject] -> ZDTree -> ZDTree
  64 makeTree trace tree = foldl' processCommand tree trace
  65
  66 processCommand :: ZDTree -> ParsedObject -> ZDTree
  67 processCommand tree (CD name)
  68   | name == "/" = root tree
  69   | name == ".." = fromJust $ parent tree
  70   | otherwise = fromJust $ childWithName name tree
  71 processCommand tree LS = tree
  72 processCommand tree (PFile size name) =
  73   modifyLabel (\ (Dir n fs) -> Dir n (M.insert name size fs)) tree
  74 processCommand tree (PDirectory name) =
  75   if (isJust $ childWithName name tree)
  76   then tree
  77   else fromJust $ parent $ insert (Node { rootLabel = (Dir name M.empty)
  78                                         , subForest = []
  79                                         }) $ children tree
  80
  81
  82 childWithName :: String -> ZDTree -> Maybe ZDTree
  83 childWithName name tree = searchForChild name (firstChild tree)
  84
  85 searchForChild :: String -> Maybe ZDTree -> Maybe ZDTree
  86 searchForChild _name Nothing = Nothing
  87 searchForChild name (Just tree)
  88   | name == labelName = Just tree
  89   | otherwise = searchForChild name (next tree)
  90   where (Dir labelName _) = label tree
  91
  92 containingSizes :: DTree -> STree
  93 containingSizes (Node {rootLabel = (Dir name files), subForest = sf}) =
  94   (Node {rootLabel = (CSize name sizeHere), subForest = sizedTrees})
  95   where sizeHere = M.foldl (+) 0 $ M.map fromIntegral files
  96         sizedTrees = fmap containingSizes sf
  97
  98 transitiveSizes :: STree -> STree
  99 transitiveSizes (Node {rootLabel = (CSize name sizeHere), subForest = sf}) =
 100   (Node {rootLabel = (CSize name (sizeHere + subSizes)), subForest = sizedTrees })
 101   where sizedTrees = fmap transitiveSizes sf
 102         subSizes = sum $ fmap (extractCSize . rootLabel) sizedTrees
 103
 104 extractCSize, cancelLarge :: ContainedSize -> Integer
 105 extractCSize (CSize _ s) = s
 106
 107 cancelLarge (CSize _ s)
 108   | s <= reportingThreshold = s
 109   | otherwise = 0
 110
 111 -- Parse the input file
 112
 113 traceP :: Parser [ParsedObject]
 114 lineP :: Parser ParsedObject
 115 cdP :: Parser ParsedObject
 116 lsP :: Parser ParsedObject
 117 fileP :: Parser ParsedObject
 118 directoryP :: Parser ParsedObject
 119 letterP :: Parser Char
 120 nameP :: Parser String
 121
 122 traceP = lineP `sepBy` endOfLine
 123
 124 lineP = cdP <|> lsP <|> directoryP <|> fileP
 125 cdP = CD <$> ("$ cd " *> nameP)
 126 lsP = LS <$ "$ ls"
 127 fileP = PFile <$> (decimal <* " ") <*> nameP
 128 directoryP = PDirectory <$> ("dir " *> nameP)
 129
 130 nameP = many1 letterP
 131 letterP = satisfy (not . isSpace)
 132
 133 successfulParse :: Text -> [ParsedObject]
 134 successfulParse input =
 135   case parseOnly traceP input of
 136     Left  _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
 137     Right commands -> commands