Mastering JSON with the plate and cosmos lenses

Day 12 of Advent of Code.

I have been attacking the Advent of Code puzzles in Haskell. Sometimes you can arrive at a solution incredibly quickly by using a modern, pure, functional programming language. Day 12 is a great example of this.

The problem in Day 12 is this: you are given a JSON blob as a string and asked to sum up all the numbers inside it. You can do this in one line with shell and regexps, but the second part of the problem (filter out all objects with the string "red" as a value) forces you to parse the string.

So it’s midnight on Friday. You’re sleepy and wildly inebriated. What’s a Haskeller to do? There exists what I call the “brute force” solution, even though it is still fairly elegant thanks to the niceties of the Haskell ecosystem. In this solution we parse the string into an Aeson (as per our treasure map) Value and then fold over the data structure.¹

data Value = Object (HashMap Text Value)
           | Array (Vector Value)
           | String Text
           | Number Scientific
           | Bool Bool
           | Null
             deriving (Eq, Read, Show, Typeable, Data)

This is a recursive datatype, so writing a fold comes pretty natural:

import Data.Foldable (toList)

count (Object hashmap) = sum (map count (toList hashmap))
count (Array vector)   = sum (map count (toList vector))
count (Number sci)     = sci
count _                = 0

You gotta love the Foldable typeclass. We were able to destruct a hashmap and a vector without ever importing their modules or reading their Hackage pages.

This covers part one. Part two, to continue the brute-force approach, would involve putting a guard on the count (Object hashmap) | noReds hashmap declaration, and the guard would filter over the hashmap looking for any values of "red". We would arrive at a working answer but perhaps not an interesting one.

Rose-colored lens

But: Faithful readers of Spoke Proof will remember our post “Lens and JSON are best friends in Haskell.” What does a lensy solution look like? Let’s try sumOf. It takes two arguments: a fold (so it knows what to sum over) and the object to traverse.

tally = sumOf someFold

What should our fold be? It does not immediately seem obvious that such a lens exists. We would need a lens that targets every number in the JSON object. Fortunately, there is something very close: Control.Lens.Plated.

plate

The idea behind the plate :: (Plated a) => Traversal' a a traversal is simple: given a data structure of type a, the plate lens traverses all the data inside that is also of type a. Plated is the Swiss-army knife of recursive data structures. It lets you treat any recursive data type a as a container of a’s.

A quick example:

λ> x
Array (fromList
  [ Number 1.0
  , Object (fromList [("b", Number 2.0), ("c", String "red")])
  , Object (fromList [("b", Number 2.0), ("c", String "blue")])
  , Number 3.0])

λ> x ^.. plate
[ Number 1.0
, Object (fromList [("b", Number 2.0), ("c", String "red")])
, Object (fromList [("b", Number 2.0), ("c", String "blue")])
, Number 3.0]

λ> x ^.. (plate . plate)
[Number 2.0, String "red" , Number 2.0, String "blue"]

λ> x ^.. (plate . plate . plate)
[]

Hopefully that helps in developing an intuition about plate.

cosmos

There is another library function here, called cosmos:

λ> x ^.. cosmos
[Array (fromList
  [ Number 1.0
  , Object (fromList [("b", Number 2.0), ("c", String "red")])
  , Object (fromList [("b", Number 2.0), ("c", String "blue")])
  , Number 3.0])
, Number 1.0
, Object (fromList [("b", Number 2.0), ("c", String "red")])
, Number 2.0
, String "red"
, Object (fromList [("b", Number 2.0), ("c", String "blue")])
, Number 2.0
, String "blue"
, Number 3.0]

Can you guess what cosmos does, just based off of this input? The documentation says this:

-- Fold over all transitive descendants of a Plated container,
-- including itself.
cosmos :: Plated a => Fold a a

“Transitive descendants” means all the children of the container, then all the children’s children, then all the children’s children’s children, ad infinitum until there are no more children left to be folded. With cosmos, we now have enough to solve part one of the problem! The _Number prism will compose with cosmos to produce

λ> :t cosmos . _Number
   (Applicative f, Plated a, Contravariant f, AsNumber a) =>
   (Scientific -> f Scientific) -> a -> f a

Which is a fold!²

λ> x ^.. (cosmos . _Number)
[1.0, 2.0, 2.0, 3.0]

λ> sumOf (cosmos . _Number) x
8.0

λ> productOf (cosmos . _Number) x
12.0

Pretty neat, right?

cosmosOf

For part two, we will need the power user’s edition of cosmos, which is cosmosOf someFold. Whereas cosmos assumes you want all the children of a container, cosmosOf someFold asks uses someFold to decide which next children to target. In fact, cosmos is defined as cosmos = cosmosOf plate. It’s all coming together, babies!

nonred (Object o) | "red" `elem` o = False
nonred _ = True

This means we can pass in a custom fold that filters out all the riff-raff objects with red:

λ> x ^.. cosmosOf (plate . filtered nonred)
[Array (fromList
  [ Number 1.0
  , Object (fromList [("b", Number 2.0), ("c", String "blue")])
  , Number 3.0])
, Number 1.0
, Object (fromList [("b", Number 2.0), ("c", String "blue")])
, Number 2.0
, String "blue"
, Number 3.0]

λ> x ^.. cosmosOf (plate . filtered nonred) . _Number
[1.0, 2.0, 3.0]

λ> sumOf (cosmosOf (plate . filtered nonred) . _Number) x
6.0

Notice that we no longer fold over the object with “red”, which means _Number is no longer able to target numbers inside objects with red values, which means foldMapOf no longer sums it up!

Altogether now

module Main where

import BasePrelude
import Control.Lens
import Data.Aeson
import Data.Aeson.Lens

nonred (Object o) | "red" `elem` o = False
nonred _ = True

input      = (^?! _Value) <$> readFile "<snip>"
tally fold = sumOf (fold . _Number) <$> input
main       = (,) <$> tally plate <*> tally (plate . filtered nonred)

I think lens is great because it lets you write a program that encodes what you mean and little more. Today we wrote some code to define what it means to have a red-poisoned object, and we wrote some code to fold our values into a sum, and we wrote some code to read and parse a file. But that’s about it! The handy-dandy, all-in-one, lemon-fresh lens package took care of the rest.