advent23/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2023/01/02/optimising-haskell-example-2/
   2
   3 -- import Debug.Trace
   4
   5 import AoC
   6 import Linear
   7 import Data.Ix
   8 import Data.Monoid
   9 import Control.Monad.State.Strict
  10 import Control.Monad.ST
  11 import qualified Data.Array.IArray as A
  12 import qualified Data.Array.MArray as M
  13 import Data.Array.ST
  14 import Data.Maybe
  15
  16
  17 type Position = V2 Int -- x, y
  18
  19 data Direction = North | South | West | East
  20   deriving (Show, Eq, Ord, Enum, Bounded)
  21
  22 newtype Elf = Elf Position
  23   deriving (Eq, Ord, Show)
  24
  25 type Population = A.Array Position (Maybe Elf)
  26
  27 type MPopulation s = STArray s Position (Maybe Elf)
  28 type MClashCounts s = STArray s Position Int
  29
  30 data Grove = Grove { currentGrove :: Population, proposalDirections :: [Direction], elapsedRounds :: Int}
  31   deriving (Eq)
  32
  33 instance Show Grove where
  34   show grove = (showElves $ currentGrove grove) ++ ", " ++ (show $ take 4 $ proposalDirections grove) ++ ", e = " ++ (show $ elapsedRounds grove)
  35     where showElves g = "Grove " ++ (show $ A.bounds g) ++ " " ++ (show $ filter (isJust . snd) $ A.assocs g)
  36
  37 type GroveState = State Grove
  38
  39 main :: IO ()
  40 main =
  41   do  dataFileName <- getDataFileName
  42       text <- readFile dataFileName
  43       let grove = Grove (mkGrove text) (cycle [North .. East]) 0
  44       -- print grove
  45       -- print $ runState simulateOnce grove
  46       -- print $ execState (simulateN 4) grove
  47       print $ part1 grove
  48       print $ part2 grove
  49
  50 part1, part2 :: Grove -> Int
  51 part1 grove = countEmpty grove' bounds
  52   where grove' = currentGrove $ execState (simulateN 10) grove
  53         bounds = findBounds grove'
  54
  55 part2 grove = elapsedRounds grove'
  56   where grove' = execState simulateToCompletion grove
  57
  58
  59 simulateToCompletion, simulateOnce, growGrove, updateDirections, updateCount :: GroveState ()
  60
  61 simulateToCompletion =
  62   do oldGrove <- gets currentGrove
  63      simulateOnce
  64      newGrove <- gets currentGrove
  65      if oldGrove == newGrove
  66      then return ()
  67      else simulateToCompletion
  68
  69 simulateN :: Int -> GroveState ()
  70 simulateN 0 = return ()
  71 simulateN n =
  72   do simulateOnce
  73      simulateN (n - 1)
  74
  75 simulateOnce =
  76   do  grove <- gets currentGrove
  77       proposalsInf <- gets proposalDirections
  78       let proposals = take 4 proposalsInf
  79       let newGrove =
  80             runSTArray $
  81               do mPopulation <- M.thaw grove
  82                  mCounts <- M.mapArray (const 0) mPopulation
  83                  proposeMoves mPopulation mCounts proposals
  84                  removeClashes mPopulation mCounts
  85                  moveElves mPopulation
  86                  return mPopulation
  87       modify' (\g -> g { currentGrove = newGrove})
  88       growGrove
  89       updateDirections
  90       updateCount
  91
  92 growGrove =
  93   do grove <- gets currentGrove
  94      let (b0, b1) = findBounds grove
  95      let bounds' = (b0 ^+^ (V2 -1 -1), b1 ^+^ (V2 1 1))
  96      let grove' = A.accumArray (flip const) Nothing bounds' $ filter ((inRange bounds') . fst ) $ A.assocs grove
  97      modify' (\g -> g { currentGrove = grove'})
  98
  99 updateDirections = modify' (\g -> g { proposalDirections = tail (proposalDirections g)})
 100 updateCount = modify' (\g -> g { elapsedRounds = (elapsedRounds g) + 1})
 101
 102 -- position changing utilities
 103
 104 anyNeighbour :: [Position]
 105 anyNeighbour =  [ V2 dx dy
 106                 | dx <- [-1, 0, 1]
 107                 , dy <- [-1, 0, 1]
 108                 , not ((dx == 0) && (dy == 0))
 109                 ]
 110
 111 directionNeighbour :: Direction -> [Position]
 112 directionNeighbour North = filter (\(V2 _x y) -> y ==  1) anyNeighbour
 113 directionNeighbour South = filter (\(V2 _x y) -> y == -1) anyNeighbour
 114 directionNeighbour West  = filter (\(V2 x _y) -> x == -1) anyNeighbour
 115 directionNeighbour East  = filter (\(V2 x _y) -> x ==  1) anyNeighbour
 116
 117 stepDelta :: Direction -> Position
 118 stepDelta North = V2  0  1
 119 stepDelta South = V2  0 -1
 120 stepDelta West  = V2 -1  0
 121 stepDelta East  = V2  1  0
 122
 123 noElves :: MPopulation s -> [Position] -> ST s Bool
 124 noElves elves tests =
 125   do others <- mapM (M.readArray elves) tests
 126      return $ all isNothing others
 127
 128 isolated :: MPopulation s -> Position -> ST s Bool
 129 isolated elves here = noElves elves $ fmap (here ^+^) anyNeighbour
 130
 131 -- get elves to make proposals
 132
 133 proposeMoves :: MPopulation s -> MClashCounts s -> [Direction] -> ST s ()
 134 proposeMoves mPopulation mCounts proposals =
 135   do assocs <- M.getAssocs mPopulation
 136      mapM_ (makeProposal mPopulation mCounts proposals) assocs
 137
 138 makeProposal :: MPopulation s -> MClashCounts s -> [Direction] -> (Position, Maybe Elf) -> ST s ()
 139 makeProposal elves clashes directions (here, elf)
 140   | isNothing elf = return ()
 141   | otherwise = do isIsolated <- isolated elves here
 142                    unless isIsolated
 143                       do  proposals <- mapM (proposedStep elves here) directions
 144                           let step = fromMaybe (V2 0 0) $ getFirst $ mconcat $ fmap First proposals
 145                           let there = here ^+^ step
 146                           thereCount <- M.readArray clashes there
 147                           M.writeArray clashes there (thereCount + 1)
 148                           M.writeArray elves here (Just (Elf there))
 149
 150 proposedStep :: MPopulation s -> Position -> Direction -> ST s (Maybe Position)
 151 proposedStep elves here direction =
 152   do isFree <- noElves elves interfering
 153      if isFree
 154      then return $ Just $ stepDelta direction
 155      else return Nothing
 156   where interfering = fmap (here ^+^) $ directionNeighbour direction
 157
 158 -- find clashing elves and prevent them moving
 159
 160 removeClashes :: MPopulation s -> MClashCounts s -> ST s ()
 161 removeClashes elves counts =
 162   do cts <- M.getAssocs counts
 163      let clashes = fmap fst $ filter ((> 1) . snd) cts
 164      stopClashingElves clashes elves
 165
 166 stopClashingElves :: [Position] -> MPopulation s -> ST s ()
 167 stopClashingElves clashes elves = mapM_ stopClash targets
 168   where targets = concatMap findNbrs clashes
 169         findNbrs c = fmap (^+^ c) $ fmap stepDelta [North .. East]
 170         stopClash here =
 171           do target <- M.readArray elves here
 172              when (isJust target) $ M.writeArray elves here (Just (Elf here))
 173
 174 -- the elves move
 175
 176 moveElves :: MPopulation s -> ST s ()
 177 moveElves elves =
 178   do assocs <- M.getAssocs elves
 179      mapM_ (moveElf elves) assocs
 180
 181 moveElf :: MPopulation s -> (Position, Maybe Elf) -> ST s ()
 182 moveElf _elves (_here, Nothing) = return ()
 183 moveElf elves (here, Just (Elf there)) =
 184   do M.writeArray elves here Nothing
 185      M.writeArray elves there (Just (Elf there))
 186
 187 -- reset the array bounds
 188
 189 findBounds :: Population -> (Position, Position)
 190 findBounds grove = boundsR
 191   where bounds0 = A.bounds grove
 192         boundsT = shrink grove topStrip    topShrink    bounds0
 193         boundsB = shrink grove bottomStrip bottomShrink boundsT
 194         boundsL = shrink grove leftStrip   leftShrink   boundsB
 195         boundsR = shrink grove rightStrip  rightShrink  boundsL
 196
 197 shrink :: Population
 198           -> ((Position, Position) -> (Position, Position))
 199           -> (Position, Position)
 200           -> (Position, Position)
 201           -> (Position, Position)
 202 shrink grove findStrip stripDirection currentBounds
 203   | emptyStrip grove (findStrip currentBounds) =
 204         shrink grove findStrip stripDirection (shiftBounds currentBounds stripDirection)
 205   | otherwise =   currentBounds
 206   where shiftBounds (b0, b1) (d0, d1) = (b0 ^+^ d0, b1 ^+^ d1)
 207
 208 emptyStrip :: Population -> (Position, Position) -> Bool
 209 emptyStrip grove strip = all isNothing $ fmap (grove A.!) $ range strip
 210
 211 topStrip, bottomStrip, leftStrip, rightStrip :: (Position, Position) -> (Position, Position)
 212 topStrip    (V2 minX _minY, V2 maxX maxY) = (V2 minX maxY, V2 maxX maxY)
 213 bottomStrip (V2 minX minY, V2 maxX _maxY) = (V2 minX minY, V2 maxX minY)
 214 leftStrip   (V2 minX minY, V2 _maxX maxY) = (V2 minX minY, V2 minX maxY)
 215 rightStrip  (V2 _minX minY, V2 maxX maxY) = (V2 maxX minY, V2 maxX maxY)
 216
 217 topShrink, bottomShrink, leftShrink, rightShrink :: (Position, Position)
 218 topShrink = (V2 0 0, V2 0 -1)
 219 bottomShrink = (V2 0 1, V2 0 0)
 220 leftShrink = (V2 1 0, V2 0 0)
 221 rightShrink = (V2 0 0, V2 -1 0)
 222
 223 countEmpty :: Population -> (Position, Position) -> Int
 224 countEmpty grove bounds = length $ filter isNothing $ fmap (grove A.!) cells
 225   where cells = range bounds
 226
 227 -- Parse the input file
 228
 229 mkGrove :: String -> Population
 230 mkGrove text = A.accumArray
 231       (\_ e -> e)
 232       Nothing
 233       (V2 -1 -1, V2 maxX maxY)
 234       [ mkElf x y -- Elf (V2 x y) (V2 x y)
 235       | x <- [0..(maxX - 1)], y <- [0..(maxY - 1)]
 236       -- , isElf x y
 237       ]
 238   where rows = reverse $ lines text
 239         maxY = length rows
 240         maxX = (length $ head rows)
 241         mkElf x y
 242           | ((rows !! y) !! x) == '#' = ((V2 x y), Just ( Elf (V2 x y) ))
 243           | otherwise = ((V2 x y), Nothing)