Part 1 down

day7/main.hs

@@ -1,10 +1,22 @@
 #! /usr/bin/env -S"ANSWER=42" nix-shell
 #! nix-shell -p ghcid
-#! nix-shell -p "haskellPackages.ghcWithPackages (p: with p; [shower])"
+#! nix-shell -p "haskellPackages.ghcWithPackages (p: with p; [attoparsec multiset containers pretty-simple])"
 #! nix-shell -i "ghcid -c 'ghci -Wall' -T main"
 
 {-# LANGUAGE OverloadedStrings #-}
 
+import Control.Applicative
+import Data.Attoparsec.Text (Parser, parseOnly)
+import Data.Map (Map)
+import Data.MultiSet (MultiSet)
+import Data.Tree (Tree)
+import Text.Pretty.Simple
+import qualified Data.Attoparsec.Text as A
+import qualified Data.Map.Strict as M
+import qualified Data.MultiSet as B
+import qualified Data.Text as Text
+import qualified Data.Tree as T
+
 exampleData :: [String]
 exampleData =
   [ "light red bags contain 1 bright white bag, 2 muted yellow bags."
@@ -18,8 +30,94 @@ exampleData =
   , "dotted black bags contain no other bags."
   ]
 
+data Color = Color String String
+  deriving (Show, Ord, Eq)
+
+-- An efficient implementation of multisets, also sometimes called bags.
+type Bag = MultiSet Color
+
+parseFullColor :: Parser Color
+parseFullColor = do
+  adjective <- A.many' A.letter
+  A.skipSpace
+  colorName <- A.many' A.letter
+  A.skipSpace
+  _ <- ("bags" <|> "bag")
+  pure $ Color adjective colorName
+
+parseBag' :: Parser Bag
+parseBag' = do
+  i <- A.decimal
+  A.skipSpace
+  c <- parseFullColor
+  pure $ B.insertMany c i B.empty
+
+parseNoBag :: Parser Bag
+parseNoBag = do
+  _ <- "no other bags"
+  pure $ B.empty
+
+parseBag :: Parser Bag
+parseBag = parseNoBag <|> parseBag'
+
+type Rules = Map Color Bag
+
+parseRule :: Parser Rules
+parseRule = do
+  k <- parseFullColor
+  A.skipSpace
+  _ <- "contain"
+  A.skipSpace
+  vx <- parseBag `A.sepBy` (", ")
+  _ <- "."
+  let v = B.unions vx
+  pure $ M.insert k v M.empty
+
+
+eitherToMaybe :: Either b a -> Maybe a
+eitherToMaybe (Right v) = Just v
+eitherToMaybe (Left _) = Nothing
+
+parseStrs :: [String] -> Maybe Rules
+parseStrs s = M.unions <$> parse s
+ where
+  parse = eitherToMaybe
+        . sequence
+        . map (parseOnly (parseRule <* A.endOfInput) . Text.pack)
+
+mkTree :: (Color,Int) -> Rules -> Tree (Color, Int)
+mkTree color rules = T.unfoldTree mkNode color
+  where
+    mkNode :: (Color, Int) -> ((Color, Int), [(Color, Int)])
+    mkNode n@(c,_) = (n, concat $ B.toOccurList <$> M.lookup c rules)
+
+mkSimpleTree :: Color -> Rules -> Tree Color
+mkSimpleTree color rules = T.unfoldTree mkNode color
+  where
+    mkNode :: Color -> (Color, [Color])
+    mkNode c = (c, fmap fst <$> concat $ B.toOccurList <$> M.lookup c rules)
+
+solvePart1 :: Color -> [String] -> Maybe Int
+solvePart1 color strs = do
+  rules <- parseStrs strs
+  let trees = map (\c -> mkSimpleTree c rules) (filter (/= color) $ M.keys rules)
+  let containColor = map (elem color) trees
+  pure $ length $ filter (== True) containColor
+
 main :: IO ()
 main = do
-  _input <- lines <$> readFile "day7/input"
-  putStrLn "Day 7 - Part 1"
+  input <- lines <$> readFile "day7/input"
+  putStrLn ":: Day 7 - Part 1"
   print $ exampleData
+  putStrLn ":: Tests"
+  print $ parseOnly (parseFullColor <* A.endOfInput) "light red bag"
+  print $ parseOnly (parseBag       <* A.endOfInput) "2 muted blue bags"
+  print $ parseOnly (parseBag       <* A.endOfInput) "3 muted red bages"
+  print $ parseOnly (parseRule      <* A.endOfInput) "dark orange bags contain 3 bright white bags, 4 muted yellow bags."
+  putStrLn ":: Parsing test data"
+  print $ parseStrs exampleData
+  putStrLn ":: Building tree"
+  pPrint $ mkTree (Color "shiny" "gold", 1) <$> (parseStrs exampleData)
+  putStrLn ":: Solving Part 1"
+  print $ solvePart1 (Color "shiny" "gold") exampleData
+  print $ solvePart1 (Color "shiny" "gold") input