Today’s problem was all about reading numbers from the input file, so I first added the words function from Haskell to my utility library:

def String.words (s : String) : List String :=
  s.splitOn " " |> List.filter (!·.isEmpty)

In the problem, we also read in the operations to be performed on the numbers. I initially parsed '+' and '*' into a pair of the function and the neutral element for the operation, but then found out that List.sum will use the Add and Zero instance for the type instance and do the right thing. However, there was no List.prod in Lean’s standard library¹, so I defined it myself:

def List.prod [Mul α] [One α] (xs : List α) : α :=
  xs.foldl (· * ·) 1

I followed the implementation of List.sum just using One and Mul instead of Zero and Add. It’s still not completely clear to me why the literal * and 1 work here (like the + and 0 in the definition of List.sum), but I guess they resolve to the one and mul from the respective type class instances.

I also expected the operation and the neutral element to be “packaged together” in something like a Monoid type class, but I guess (and I really mean guess here) that’s not used here as the extra structure is not needed for e.g. List.sum, i.e. we could “sum” a list of values of a type that don’t fulfil the monoid laws.

Things I (re-)learned today

• The “batteries” library had my beloved (from earlier AoC seasons) transpose function for List (List α).
• That Lean’s type class hierarchy is more fine-grained than Haskell’s with separate type classes for e.g. Add, Mul, Zero, One. Definitely something I want to learn more about!

Solution

import Batteries.Data.String
import Aoc
open Aoc

def main : IO Unit := do
  let lines : List String ← (·.trim.splitOn "\n") <$> readInput "Day06/input.txt"
  let (nums, ops) := match lines.splitAt (lines.length - 1) with
    | (ns, [os]) =>
      let ns := ns.map (λ (l : String) => l.words.map String.toInt!)
      let ops := os.words.map (λ o =>
        match o with
        | "+" => List.sum
        | "*" => List.prod
        | _     => panic! "Unexpected operation: {o}"
      )
      (ns, ops)
    | _ => panic! "Unexpected input format"

  -- part 1
  let results := nums.transpose.zipWith (λ ns op => op ns) ops
  IO.println s!"Part 1: {results.sum}"

  -- part 2
  let cephaloNums := lines.map (String.toList)
    -- work column by column
    |>.transpose
    -- split on the empty columns
    |> List.splitOnP (·.all (· == ' '))
    -- for each number in the group, collect the digits
    |> List.map (List.map (List.filter (fun c => c.isDigit)))
    -- … and finally assemble the digits into an integer
    |> List.map (List.map (String.mk · |> String.toInt!))

  let results := cephaloNums.zipWith (λ ns op => op ns) ops
  IO.println s!"Part 2: {results.sum}"