{-# LANGUAGE FlexibleContexts, NondecreasingIndentation, TupleSections, TypeApplications, ViewPatterns #-}
module Day17 (day17a, day17b) where
import Control.Arrow ((&&&), (***))
import Control.Monad (filterM)
import Control.Monad.Loops (andM, dropWhileM, takeWhileM)
import Data.Array.IArray (Array, IArray, Ix, bounds, accumArray, elems, rangeSize)
import Data.Array.ST (MArray, getBounds, thaw, readArray, runSTArray, writeArray)
import Data.Bool (bool)
import Data.List.NonEmpty (nonEmpty)
import Data.List.Split (chunksOf)
import Data.Maybe (fromMaybe, listToMaybe, maybeToList)
import Data.Semigroup (Max(Max), Min(Min), sconcat)
import Data.Tuple (swap)
import Text.Megaparsec (MonadParsec, (<|>), parseMaybe, sepEndBy)
import Text.Megaparsec.Char (newline, string)
import Text.Megaparsec.Char.Lexer (decimal)
data Element = Empty | Wall | Stagnant | Flowing
deriving (Eq)
debug :: (IArray a Element, Integral i, Ix i) => a (i, i) Element -> String
debug scene@(bounds -> ((_, minX), (_, maxX))) =
unlines $ map (bool '.' '+' . (== 500)) [minX..maxX] :
chunksOf (rangeSize (minX, maxX)) (elemChar <$> elems scene)
where elemChar Empty = '.'
elemChar Wall = '#'
elemChar Stagnant = '~'
elemChar Flowing = '|'
parser :: (MonadParsec e String m, IArray a Element, Integral i, Ix i) => m (a (i, i) Element)
parser = sepEndBy lineParser newline >>= makeScene . concat
where lineParser = do
(secondAxis, order) <- ("x=", id) <$ string "y=" <|> ("y=", swap) <$ string "x="
y <- decimal <* string ", "
x0 <- string secondAxis *> decimal
x1 <- string ".." *> decimal
return [order (y, x) | x <- [x0..x1]]
makeScene points@(nonEmpty -> Just points') = return $
accumArray (flip const) Empty ((minY, minX), (maxY, maxX)) $ (, Wall) <$> points
where ((Min minY, Min minX), (Max maxY, Max maxX)) =
sconcat $ ((Min *** Min . pred) &&& (Max *** Max . succ)) <$> points'
makeScene _ = fail "expecting points"
ranges :: (Num a, Ord a) => [a] -> [(a, a)]
ranges = ranges' Nothing
where ranges' k [] = maybeToList k
ranges' Nothing (x:xs) = ranges' (Just (x, x)) xs
ranges' (Just (l, r)) (x:xs)
| r + 1 < x = (l, r) : ranges' (Just (x, x)) xs
| otherwise = ranges' (Just (l, max r x)) xs
flood :: (MArray a Element m, Integral i, Ix i, Show i) => i -> a (i, i) Element -> m Bool
flood startX scene = do
((minY, minX), (maxY, maxX)) <- getBounds scene
let pour y (x0, x1)
| y > maxY = return False
| otherwise = do
empties <- filterM (\x -> (== Empty) <$> readArray scene (y, x)) [x0..x1]
blocked0 <- andM [(/= Flowing) <$> readArray scene (y, x) | x <- [x0..x1]]
blocked1 <- and <$> mapM (fill y) (ranges empties)
return $ blocked0 && blocked1
fill y (x0, x1) = do
blocked <- pour (y + 1) (x0, x1)
if not blocked
then False <$ sequence_ [writeArray scene (y, x) Flowing | x <- [x0..x1]]
else do
let isEmpty x = (== Empty) <$> readArray scene (y, x)
isSupported x =
if y < maxY
then (`notElem` [Empty, Flowing]) <$> readArray scene (y + 1, x)
else return False
lefts <- takeWhileM isEmpty [x0 - 1, x0 - 2..minX]
rights <- takeWhileM isEmpty [x1 + 1, x1 + 2..maxX]
l <- fromMaybe (last $ x0 : lefts) . listToMaybe <$> dropWhileM isSupported lefts
r <- fromMaybe (last $ x1 : rights) . listToMaybe <$> dropWhileM isSupported rights
if l >= x0 && r <= x1
then True <$ sequence_ [writeArray scene (y, x) Stagnant | x <- [x0..x1]]
else do
sequence_ [writeArray scene (y, x) Wall | x <- [x0..x1]]
blockedL <- fill y (l, x0 - 1)
blockedR <- fill y (x1 + 1, r)
let blocked'@(bool Flowing Stagnant -> elt) = blockedL && blockedR
blocked' <$ sequence_ [writeArray scene (y, x) elt | x <- [l..r]]
pour minY (startX, startX)
day17a :: String -> Maybe Int
day17a input = do
scene <- parseMaybe @() parser input :: Maybe (Array (Int, Int) Element)
let scene' = runSTArray $ thaw scene >>= \a -> a <$ flood 500 a
return $ length $ filter (`elem` [Stagnant, Flowing]) $ elems scene'
day17b :: String -> Maybe Int
day17b input = do
scene <- parseMaybe @() parser input :: Maybe (Array (Int, Int) Element)
let scene' = runSTArray $ thaw scene >>= \a -> a <$ flood 500 a
return $ length $ filter (== Stagnant) $ elems scene'