Day 12

2023-09-24 20:20:47 +02:00
parent 0402af987a
commit f9e3cf2f75
6 changed files with 192 additions and 3 deletions
--- a/src/Days/D12.hs
+++ b/src/Days/D12.hs
@@ -0,0 +1,143 @@
+{-# LANGUAGE ImportQualifiedPost #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+
+module Days.D12 where
+
+import Common
+import Control.Monad (filterM, foldM)
+import Control.Monad.Extra (concatMapM)
+import Control.Monad.ST (ST, runST)
+import Data.ByteString (ByteString)
+import Data.ByteString.Char8 qualified as BS
+import Data.ByteString.Unsafe qualified as BS
+import Data.Char (ord)
+import Data.Functor ((<&>))
+import Data.List (findIndex)
+import Data.Maybe (fromMaybe)
+import Data.Vector.Storable (Vector)
+import Data.Vector.Storable qualified as Vec
+import Data.Vector.Storable.Mutable qualified as MutVec
+import Foreign.C (castCCharToChar)
+import Foreign.ForeignPtr (newForeignPtr_)
+import Foreign.Storable (Storable)
+import System.IO.Unsafe (unsafePerformIO)
+
+data Point = P !Int !Int
+  deriving (Show, Eq, Ord)
+
+data Grid t = Grid Int Int (Vector t)
+
+data MGridSet s = MGridSet Int Int (MutVec.MVector s Bool)
+
+neigh4 :: Point -> [Point]
+neigh4 (P x y) =
+  [ P (x - 1) y,
+    P (x + 1) y,
+    P x (y - 1),
+    P x (y + 1)
+  ]
+
+(@) :: (Storable t) => Grid t -> Point -> t
+(@) (Grid _ w m) (P x y) = m Vec.! (w * x + y)
+
+get :: MGridSet s -> Point -> ST s Bool
+get (MGridSet _ w m) (P x y) = MutVec.unsafeRead m (w * x + y)
+
+set :: MGridSet s -> Point -> ST s Bool
+set (MGridSet _ w m) (P x y) = MutVec.unsafeExchange m (w * x + y) True
+
+dims :: Grid t -> (Int, Int)
+dims (Grid h w _) = (h, w)
+
+mkGridSet :: (Int, Int) -> ST s (MGridSet s)
+mkGridSet (h, w) = MGridSet h w <$> MutVec.replicate (h * w) False
+
+contains :: (Storable t) => Grid t -> Point -> Bool
+contains (Grid h w _) (P x y) =
+  0 <= x
+    && 0 <= y
+    && x < h
+    && y < w
+
+bfs :: (Storable t) => Grid t -> (Point -> Point -> Bool) -> Point -> [[Point]]
+bfs g f start =
+  runST $
+    mkGridSet (dims g)
+      >>= (\mgs -> set mgs start >> return mgs)
+      >>= flip go [start]
+  where
+    go :: MGridSet s -> [Point] -> ST s [[Point]]
+    go seen l
+      | null l = pure []
+      | otherwise = (l :) <$> (advance seen l >>= go seen)
+
+    -- advance :: MGridSet s -> [Point] -> ST s [Point]
+    -- advance seen points =
+    --   points
+    --     & concatMap (\point -> [p | p <- neigh4 point, g `contains` p, f point p])
+    --     & foldM
+    --       (\l p -> set seen p <&> (\isNew -> if isNew then l else p : l))
+    --       []
+
+    -- Uglier, but unfortunately significantly faster
+    advance :: MGridSet s -> [Point] -> ST s [Point]
+    advance seen =
+      concatMapM
+        ( \point ->
+            filterM (fmap not . get seen) [p | p <- neigh4 point, g `contains` p]
+              >>= foldM
+                (\l p -> set seen p <&> \wasSeen -> if wasSeen then l else p : l)
+                []
+                . filter (f point)
+        )
+
+find :: (Storable t, Eq t) => Grid t -> t -> Maybe Point
+find (Grid _ w m) v = (\pos -> P (pos `div` w) (pos `mod` w)) <$> Vec.elemIndex v m
+
+elevation :: Char -> Int
+elevation 'S' = ord 'a'
+elevation 'E' = ord 'z'
+elevation c = ord c
+
+needGear :: Grid Char -> Point -> Point -> Bool
+needGear grid from to = elevation (grid @ to) <= elevation (grid @ from) + 1
+
+part1 :: Grid Char -> Maybe Int
+part1 grid = do
+  start <- find grid 'S'
+  end <- find grid 'E'
+  let bfsResult = bfs grid (needGear grid) start
+  findIndex (elem end) bfsResult
+
+part2 :: Grid Char -> Maybe Int
+part2 grid =
+  find grid 'E'
+    >>= findIndex (any (\p -> elevation (grid @ p) == elevation 'a'))
+      . bfs grid (flip $ needGear grid)
+
+gridify :: ByteString -> Grid Char
+gridify bytes =
+  let numBytes = BS.length bytes
+      w = fromMaybe numBytes (BS.elemIndex '\n' bytes)
+      h = (numBytes + 1) `div` (w + 1)
+      joinedBytes = BS.filter (/= '\n') bytes
+   in Grid
+        h
+        w
+        ( unsafePerformIO $
+            BS.unsafeUseAsCStringLen
+              joinedBytes
+              ( \(mem, len) ->
+                  newForeignPtr_ mem
+                    >>= Vec.unsafeFreeze . flip MutVec.unsafeFromForeignPtr0 len
+                    <&> Vec.map castCCharToChar
+              )
+        )
+
+day :: Day
+day =
+  ByteStringDay
+    ( \text ->
+        let grid = gridify text
+         in Right (Answer <$> part1 grid, Answer <$> part2 grid)
+    )