Advent of Code Day 3 (Part 1 & 2)

{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TupleSections #-}

module AdventOfCode.Day3 where

import Data.List (find)
import Data.Maybe (isJust, fromJust)
import Control.Monad (forM_, forM)
import Control.Monad.IO.Class (MonadIO(..))
import Data.Array.ST (Ix(range), readArray, writeArray, MArray(newArray))
import Data.Array.IO.Internals (IOArray)
import Data.Char (isDigit)
import Data.Hashable (Hashable)
import qualified Data.HashTable.IO as HT

-- part 1

data Matrix arr e = Matrix
  { unItems :: arr Int e
  , unDim :: (Int, Int)
  }

createMatrix :: (MArray arr e m) => Int -> Int -> e -> m (Matrix arr e)
createMatrix w h defaultValue = do
  items <- newArray (0, w * h - 1) defaultValue
  return $ Matrix items (w, h)

getIndex :: Matrix arr e -> Int -> Int -> Maybe Int
getIndex m row col =
  let (w, h) = unDim m
      itemIdx = row * w + col
  in if (row < 0 || row >= w) || (col < 0 || col >= h)
    then Nothing
    else Just itemIdx

writeValue :: (MArray arr e m) => Matrix arr e -> Int -> Int -> e -> m Bool
writeValue m row col val = case getIndex m row col of
  Nothing -> return False
  Just idx -> writeArray (unItems m) idx val >> return True

enumerate :: (Enum b) => b -> [a] -> [(b, a)]
enumerate = zip . enumFrom

adjacentIdxs :: Matrix arr e -> Int -> Int -> [Int]
adjacentIdxs m row col = [fromJust idx | r <- [row-1 .. row+1]
                                       , c <- [col-1 .. col+1]
                                       , r /= row || c /= col
                                       , let idx = getIndex m r c
                                       , isJust idx]

adjacentItems :: (MArray arr e m) => Matrix arr e -> Int -> Int -> m [e]
adjacentItems m row col = mapM (readArray $ unItems m) $ adjacentIdxs m row col

findSpans :: (a -> Bool) -> [a] -> [((Int, Int), [a])]
findSpans p = go . enumerate 0
  where go [] = []
        go xs = let (found, rest) = span (p . snd) xs
                    (idxs, item) = foldr (\(i, x) (is, xs') -> (i : is, x : xs')) ([], []) found
                    spanTuple = (head idxs, last idxs)
                in if null found
                  then go $ tail rest
                  else (spanTuple, item) : go rest

arraySpan :: (MArray arr e m, Ix i) => arr i e -> (i, i) -> m [e]
arraySpan arr = mapM (readArray arr) . range

findNumberSpans :: (MArray arr Char m) => Matrix arr Char -> m [(Int, ((Int, Int), String))]
findNumberSpans (Matrix items (w, h)) = do
  enumeratedRows <- enumerate 0 <$> sequence [arraySpan items (i, i+w-1) | i <- [0, w .. w*h-1]]
  fmap concat $ forM enumeratedRows $ \(rowIdx, rowVals) -> do
    let spans = findSpans isDigit rowVals
    return $ map (rowIdx,) spans

sumPartNumbers :: (MArray arr Char m) => Matrix arr Char -> [(Int, ((Int, Int), String))] -> m Int
sumPartNumbers matrix numberSpans = do
  fmap sum $ forM numberSpans $ \(rowIdx, (colStartAndEnd, partNo)) -> do
    isValidPartNo <- or <$> forM (range colStartAndEnd) (fmap (any (\ c -> c /= '.' && not (isDigit c))) . adjacentItems matrix rowIdx)
    return $ if isValidPartNo then read partNo else 0

writeRows :: (MArray arr e m) => Matrix arr e -> [(Int, [e])] -> m ()
writeRows m rows = do
  forM_ rows $ \(rowIndex, rowVals) -> do
    forM_ (enumerate 0 rowVals) $ \(colIndex, val) -> do
      writeValue m rowIndex colIndex val

readMatrix :: IO (Matrix IOArray Char)
readMatrix = do
  enumeratedLines <- enumerate 0 . lines <$> getContents :: IO [(Int, String)]
  m <- createMatrix (length enumeratedLines) (1 + fst (last enumeratedLines)) '\0' :: IO (Matrix IOArray Char)
  writeRows m enumeratedLines >> return m

partOne :: IO Int
partOne = readMatrix >>= \m -> findNumberSpans m >>= sumPartNumbers m

-- part 2

type HashTable k v = HT.BasicHashTable k v

insertWith :: forall k v m. (Hashable k, MonadIO m) => HashTable k v -> (v -> v) -> k -> v -> m ()
insertWith ht t k v = do
  existingValMb <- liftIO $ HT.lookup ht k :: m (Maybe v)
  liftIO $ case existingValMb of
    Nothing -> HT.insert ht k v
    Just v' -> HT.insert ht k (t v')

sumGears :: forall arr m. (MArray arr Char m, MonadIO m) => Matrix arr Char -> [(Int, ((Int, Int), String))] -> m Int
sumGears matrix numberSpans = do
    ht <- liftIO HT.new :: m (HashTable Int [Int])
    processSpans numberSpans ht
    sum . map product . filter ((== 2) . length) . map snd <$> liftIO (HT.toList ht)
  where processSpans [] _ = return ()
        processSpans ((rowIdx, (colStartAndEnd, partNo)) : rest) ht = do
          let colBounds = range colStartAndEnd
          adjItems <- concat <$> mapM (adjacentItems matrix rowIdx) colBounds
          let adjIdxs = concatMap (adjacentIdxs matrix rowIdx) colBounds
              starMb = find ((== '*') . snd) $ zip adjIdxs adjItems
          case starMb of
            Nothing -> processSpans rest ht
            Just (starIdx, _) -> insertWith ht (read partNo :) starIdx [read partNo] >> processSpans rest ht

partTwo :: IO Int
partTwo = readMatrix >>= \m -> findNumberSpans m >>= sumGears m

main :: IO ()
main = partTwo >>= print