aoc2020/day7.hs

{-# LANGUAGE FlexibleContexts #-}

import Text.Parsec

data Bag = Bag Label [(Quantity, Label)]
  deriving Show
type Label = String
type Quantity = Int

main :: IO ()
main = do
  raw <- readFile "day7.txt"
  let bags = validate $ map (parse lineP []) $ lines raw
      ansA = pred $ length $ filter (\x -> (snd x) == True)
                $ loeb $ map loebifyA bags
      ansB = pred $ unMaybe $ lookup "shiny gold"
                $ loeb $ map loebifyB bags
   in do
     putStrLn $ "day7a: " ++ (show ansA)
     putStrLn $ "day7b: " ++ (show ansB)

loeb :: Functor f => f (f a -> a) -> f a
loeb x = go where go = fmap ($ go) x

loebifyA :: Bag -> ([(Label, Bool)] -> (Label, Bool))
loebifyA (Bag "shiny gold" _) = const ("shiny gold", True)
loebifyA (Bag l bs) =
  (\ls -> (l, or (unMaybes (map (\f -> f ls) (map lookupify bs)))))
  where
    lookupify (_,s) = lookup s

loebifyB :: Bag -> ([(Label, Int)] -> (Label, Int))
loebifyB (Bag l []) = const (l, 1)
loebifyB (Bag l bs) =
  (\ls -> (l, sum (1:(map (\f -> f ls) (map lookupify bs)))))
  where
    lookupify (q,s) = (\ms -> q * (unMaybe $ lookup s ms))

unMaybe :: Maybe a -> a
unMaybe Nothing = error "invalid bag label"
unMaybe (Just v) = v

unMaybes :: [Maybe a] -> [a]
unMaybes = map unMaybe

validate :: [Either ParseError Bag] -> [Bag]
validate [] = []
validate ((Left _):_) = error "invalid input"
validate ((Right b):bs) = b:(validate bs)

lineP :: Parsec String () Bag
lineP = do
  l <- initlabelP
  ls <- baglistP
  eof
  return $ Bag l ls

baglistP :: Parsec String () [(Quantity, Label)]
baglistP = do
  ((string "no other bags.") >> return []) <|>
    (endBy qlabelP $ oneOf [',', '.'] >> whitespaces)

initlabelP :: Parsec String () String
initlabelP = manyTill anyChar
  (try $ whitespaces >> (string "bags contain") >> whitespaces)

whitespaces :: Parsec String () String
whitespaces = many $ char ' '

qlabelP :: Parsec String () (Quantity, Label)
qlabelP = do
  q <- many1 digit
  whitespaces
  s <- labelP
  return (read q,s)

labelP :: Parsec String () String
labelP = manyTill anyChar
  (try $ whitespaces >> (string "bag") >> optional (char 's'))