2015 Day 12 in Haskell

2015/day12/day12.hs

@@ -0,0 +1,67 @@
+import Aoc
+import Data.Char ( isDigit )
+import Control.Arrow ( first, second )
+
+data JSON a = Int a | Str String | Array [JSON a] | Object [(String, JSON a)]
+    deriving ( Show, Eq )
+
+parseElements :: (String -> (p, String)) -> Char -> String -> ([p], String)
+parseElements parseElement closing cs
+    | head cs == closing = ([], tail cs)
+    | head cs == ',' = parseElements parseElement closing (tail cs)
+    | otherwise =
+        let (element, rest0) = parseElement cs in
+        let (elements, rest1) = parseElements parseElement closing rest0 in
+        (element:elements, rest1)
+
+parseString :: String -> (String, String)
+parseString rest0 =
+    let (string, rest1) = span (/= '"') rest0 in -- no escaping
+    let ('"':rest2) = rest1 in
+    (string, rest2)
+
+parseKeyValue :: Read a => [Char] -> ((String, JSON a), String)
+parseKeyValue ('"':rest0) =
+    let (string, rest1) = parseString rest0 in
+    let (':':rest2) = rest1 in
+    let (object, rest3) = parseJSON rest2 in
+    ((string, object), rest3)
+
+parseJSON :: Read a => String -> (JSON a, String)
+parseJSON ('[':rest0) =
+    let (elements, rest1) = parseElements parseJSON ']' rest0 in
+    (Array elements, rest1)
+parseJSON ('{':rest0) =
+    let (elements, rest1) = parseElements parseKeyValue '}' rest0 in
+    (Object elements, rest1)
+parseJSON ('"':rest0) = first Str $ parseString rest0
+parseJSON rest = if isNum (head rest) then first (Int . read) $ span isNum rest else error rest where isNum c = isDigit c || c == '-'
+
+parseDone :: (p, String) -> p
+parseDone (p, "") = p
+
+instance Foldable JSON where
+  foldMap :: Monoid m => (a -> m) -> JSON a -> m
+  foldMap f (Int a) = f a
+  foldMap f (Str _) = mempty
+  foldMap f (Array xs) = foldMap (foldMap f) xs
+  foldMap f (Object kvs) = foldMap (foldMap f . snd) kvs
+
+hasRed :: Eq a => JSON a -> Bool
+hasRed (Object kvs) = any ((==Str "red") . snd) kvs
+hasRed _ = False
+
+removeRed :: Eq a => JSON a -> JSON a
+removeRed json = case json of
+    Object kvs -> if hasRed json then Object [] else Object $ map (second removeRed) kvs
+    Array xs -> Array $ map removeRed xs
+    _ -> json
+
+part1 :: JSON Int -> Int
+part1 json = sum $ foldMap (:[]) json
+
+part2 :: JSON Int -> Int
+part2 = part1 . removeRed
+
+main :: IO ()
+main = aocMain (parseDone . parseJSON . single) part1 part2