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 = solveA bags
      ansB = solveB bags
   in do
     putStrLn $ "day7a: " ++ (show ansA)
     putStrLn $ "day7b: " ++ (show ansB)

solveA :: [Bag] -> Int
solveA bs = pred
       $ length
       $ filter (== True)
       $ map snd
       $ loeb
       $ map loebifyA bs

solveB :: [Bag] -> Int
solveB bs = pred
       $ unsafeLookup "shiny gold"
       $ loeb
       $ map loebifyB bs

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 [f ls | f <- innerbags]))
  where
    lookupify = unsafeLookup . snd
    innerbags = map lookupify bs

loebifyB :: Bag -> ([(Label, Int)] -> (Label, Int))
loebifyB (Bag l []) = const (l, 1)
loebifyB (Bag l bs) =
  (\ls -> (l, sum $ 1 : [q*(f ls) | (f,q) <- innerbags]))
  where
    lookupify = unsafeLookup . snd
    innerbags = zip (map lookupify bs)
                    (map fst bs)

unsafeLookup :: Eq a => a -> [(a,b)] -> b
unsafeLookup k = unMaybe . (lookup k)

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

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'))