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