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