using Comonads to solve Advent of Code 2021 day 1

Day1.hs

-- in response to https://twitter.com/BartoszMilewski/status/1466106578524786690
-- 
-- Can we solve https://adventofcode.com/2021/day/1 using the NonEmpty Comonad?
{-# LANGUAGE LambdaCase, ScopedTypeVariables #-}
import Prelude hiding (sum)
import Control.Category ((>>>))
import Control.Comonad (Comonad(extract, extend))
import Control.Monad (guard)
import Data.Either (rights)
import Data.List.NonEmpty (NonEmpty((:|)))
import Data.Maybe (mapMaybe)
import qualified Data.List.NonEmpty as NonEmpty


-- There are a bunch of generic functions like 'mapM' and 'whileM' which work
-- with any Monad. Correspondingly, there should be generic functions which
-- work with any Comonad. I couldn't find them on hackage though, so let's
-- define our own.

-- I often want to propagate data from one part of the data structure to the
-- other; in this case from the tail of the list to the head. We can do that
-- using a function which gathers the data from its immediate neighbours, by
-- repeatedly calling 'extend' until the data has jumped from neighbour to
-- neighbour to our desired destination. In this implementation, the desired
-- destination is the focus.
propagate
  :: forall w a r. Comonad w
  => (w (Either a r) -> Maybe r)
  -> w a -> r
propagate weither2maybe
   -- w a
    = fmap Left
   -- w (Either a r)
  >>> iterate (extend weither2either)
   -- [w (Either a r)]
  >>> fmap extract
   -- [Either a r]
  >>> rights
   -- [r]
  >>> head  -- the list cannot be empty, but it can be _|_ if the data never
            -- reaches the focus
  where
    weither2either
      :: w (Either a r) -> Either a r
    weither2either weither
      = case weither2maybe weither of
          Just r
            -> Right r
          Nothing
            -> extract weither

-- For example, we can sum all the values in the NonEmpty by propagating the
-- sum so far from the tail to the head. Note that pattern-matching failures
-- cause the Maybe computation to fail with 'Nothing', thus indicating that the
-- data has not yet reached the current cell.
sum
  :: NonEmpty Int -> Int
sum = propagate $ \case
  current :| [] -> do
    Left x <- pure current
    pure x
  current :| neighbour : _ -> do
    Left x <- pure current
    Right sumSoFar <- pure neighbour
    pure (sumSoFar + x)

-- The two exercise in Day 1 use 'sum' in a very specific way: to count the
-- number of places in which a particular pattern occurs. So let's create
-- another generic function for that use case. 'sum' only works with NonEmpty,
-- but we can make 'count' work with any Comonad for which we know how to sum
-- its contents.
count
  :: Comonad w
  => (w Int -> Int)
  -> (w a -> Maybe ())
  -> w a -> Int
count wsum detect
  = wsum . fmap maybe2int . extend detect
  where
    maybe2int
      :: Maybe () -> Int
    maybe2int (Just ())
      = 1
    maybe2int Nothing
      = 0

-- We can now solve the first part of the Day 1 exercise, counting the number
-- of times the depth increases, by having each cell check whether they are a
-- location where the depth increases (one time) or not (zero times), and then
-- taking the sum.
day1a
  :: NonEmpty Int -> Int
day1a = count sum detect
  where
    detect
      :: NonEmpty Int -> Maybe ()
    detect nonEmpty = do
      x :| y : _ <- pure nonEmpty
      guard (x < y)

-- The second part is very similar, except that we're using a sliding window.
-- There are two cells at the end which don't have enough data to begin a
-- sliding window of length 3, so we need to use a 'Maybe'.
day1b :: NonEmpty Int -> Int
day1b
   -- NonEmpty Int
    = extend slidingWindow
   -- NonEmpty (Maybe Int)
  >>> count sum detect
   -- Int
  where
    slidingWindow
      :: NonEmpty Int -> Maybe Int
    slidingWindow nonEmpty = do
      x1 :| x2 : x3 : _ <- pure nonEmpty
      pure (x1 + x2 + x3)

    detect
      :: NonEmpty (Maybe Int) -> Maybe ()
    detect nonEmpty = do
      Just x :| Just y : _ <- pure nonEmpty
      guard (x < y)

-- Compare the above Comonad-based solution to the one I originally wrote for
-- Advent of Code:
-- 
-- > day1a :: [Int] -> Int
-- > day1a xs
-- >   = length
-- >   $ filter id
-- >   $ zipWith (>) (drop 1 xs) xs
-- > 
-- > day1b :: [Int] -> Int
-- >   = day1a
-- >   . fmap sum
-- >   . mapMaybe take3
-- >   . tails
-- >   where
-- >     take3 :: [a] -> Maybe [a]
-- >     take3 xs = do
-- >       x1 : x2 : x3 : _ <- pure xs
-- >       pure [x1, x2, x3]
-- 
-- One thing I notice is that my original solution reuses 'day1a' in 'day1b',
-- while my Comonad-based solution doesn't. That's because Comonad operations
-- cannot change the overall shape of a structure, they can only change the
-- value of each cell. Thus, when the sliding window forces us to ignore the
-- last two cells, this has a cascading impact on the rest of the computation,
-- which now needs to be adapted to handle 'Maybe's.
-- 
-- Of course, that's only a problem if we restrict ourselves to Comonad
-- operations. I chose to do that as part of Bartosz's challenge, but under
-- normal circumstances, it makes more sense to mix and match Comonad and
-- non-Comonad operations as needed, like this:

day1b' :: NonEmpty Int -> Int
day1b'
   -- NonEmpty Int
    = extend slidingWindow
   -- NonEmpty (Maybe Int)
  >>> NonEmpty.toList
   -- [Maybe Int]
  >>> mapMaybe id
   -- [Int]
  >>> NonEmpty.fromList  -- _|_ if the original list has less than 3 entries
   -- NonEmpty Int
  >>> day1a
   -- Int
  where
    slidingWindow
      :: NonEmpty Int -> Maybe Int
    slidingWindow nonEmpty = do
      x1 :| x2 : x3 : _ <- pure nonEmpty
      pure (x1 + x2 + x3)

main :: IO ()
main = do
  input <- readFile "example"
  -- input <- readFile "input"
  print $ day1a $ NonEmpty.fromList $ fmap read $ lines $ input
  print $ day1b $ NonEmpty.fromList $ fmap read $ lines $ input
  print $ day1b' $ NonEmpty.fromList $ fmap read $ lines $ input