src/advent21/advent21parallel.hs

   1 {-# LANGUAGE NegativeLiterals #-}
   2 {-# LANGUAGE FlexibleContexts #-}
   3 {-# LANGUAGE OverloadedStrings #-}
   4 {-# LANGUAGE TypeFamilies #-}
   5 {-# LANGUAGE BangPatterns #-}
   6
   7 import Data.Text (Text)
   8 import qualified Data.Text as T
   9 import qualified Data.Text.IO as TIO
  10
  11 import Text.Megaparsec hiding (State)
  12 import qualified Text.Megaparsec.Lexer as L
  13 import Text.Megaparsec.Text (Parser)
  14 import qualified Control.Applicative as CA
  15 -- import qualified Data.Functor as F
  16
  17 import qualified Data.Map.Strict as M
  18 import Data.Map.Strict ((!))
  19
  20 import Data.List
  21
  22 import Control.Parallel.Strategies (parMap, rpar)
  23
  24
  25 type Grid = M.Map (Int, Int) Bool
  26 type ExplodedGrid = M.Map (Int, Int) Grid
  27
  28 data Rule = Rule Grid Grid deriving (Eq, Show)
  29
  30 rulePre (Rule g _) = g
  31 rulePost (Rule _ g) = g
  32
  33
  34 initialGrid = case parse gridP "" ".#./..#/###" of
  35                 Left _ -> M.empty
  36                 Right g -> g
  37
  38
  39 main :: IO ()
  40 main = do
  41         text <- TIO.readFile "data/advent21.txt"
  42         let rules = readRules text
  43         print $ countLit $ nthApplication rules 5
  44         print $ countLit $ nthApplication rules 18
  45
  46
  47 -- Read the rules, and expand them to all equivalent left hand sides
  48 readRules :: Text -> [Rule]
  49 readRules = expandRules . successfulParse
  50
  51
  52 expandRules :: [Rule] -> [Rule]
  53 expandRules = concatMap expandRule
  54
  55 expandRule :: Rule -> [Rule]
  56 expandRule rule = [Rule l (rulePost rule) | l <- allArrangements (rulePre rule)]
  57
  58
  59 reflectH :: Grid -> Grid
  60 reflectH g = M.fromList [((r, c) , M.findWithDefault False (rm - r, c) g) | r <- [0..rm], c <- [0..cm] ]
  61     where (rm, cm) = bounds g
  62
  63 reflectV :: Grid -> Grid
  64 reflectV g = M.fromList [((r, c) , M.findWithDefault False (r, cm - c) g) | r <- [0..rm], c <- [0..cm] ]
  65     where (rm, cm) = bounds g
  66
  67 -- awkward naming to avoid clashing with Prelude
  68 transposeG :: Grid -> Grid
  69 transposeG g = M.fromList [((c, r) , M.findWithDefault False (r, c) g) | r <- [0..rm], c <- [0..cm] ]
  70     where (rm, cm) = bounds g
  71
  72
  73 -- Find all the arrangments of a grid, including reflection and rotation.
  74 allArrangements :: Grid -> [Grid]
  75 allArrangements grid = map (\f -> f grid) [ id
  76                                           , reflectH
  77                                           , reflectV
  78                                           , transposeG
  79                                           , reflectH . transposeG
  80                                           , reflectV . transposeG
  81                                           , reflectH . reflectV . transposeG
  82                                           , reflectV . reflectH
  83                                           ]
  84
  85
  86
  87 -- Count number of lit pixels
  88 countLit :: Grid -> Int
  89 countLit = M.size . M.filter id
  90
  91 -- apply the rules _n_ times
  92 nthApplication :: [Rule] -> Int -> Grid
  93 nthApplication rules n = (!! n) $ iterate (applyOnce rules) initialGrid
  94
  95 -- Apply one step of the expansion
  96 applyOnce :: [Rule] -> Grid -> Grid
  97 -- applyOnce rules g = contractExploded $ M.map (apply rules) $ explodeGrid g
  98 applyOnce rules g = contractExploded $ M.unions $ parMap rpar (M.map (apply rules)) $ M.splitRoot $ explodeGrid g
  99
 100
 101 -- find the appropriate rule and apply it to a grid
 102 apply :: [Rule] -> Grid -> Grid
 103 apply rules grid = rulePost thisRule
 104     where ri = head $ findIndices (\r -> rulePre r == grid) rules
 105           thisRule = rules!!ri
 106
 107
 108 -- create the appropriate subgrids of a grid
 109 explodeGrid :: Grid -> ExplodedGrid
 110 explodeGrid g = if (rm + 1) `rem` 2 == 0
 111                 then explodeGrid' 2 g
 112                 else explodeGrid' 3 g
 113     where (rm, cm) = bounds g
 114
 115 explodeGrid' :: Int -> Grid -> ExplodedGrid
 116 explodeGrid' n g = M.fromList [((bigR, bigC), subGrid n g bigR bigC) | bigR <- [0..bigRm], bigC <- [0..bigCm]]
 117     where (rm, cm) = bounds g
 118           bigRm = (rm + 1) `div` n - 1
 119           bigCm = (cm + 1) `div` n - 1
 120
 121
 122 subGrid :: Int -> Grid -> Int -> Int -> Grid
 123 subGrid n g bigR bigC = M.fromList [ ((r, c),
 124                                       M.findWithDefault False (r + rStep, c + cStep) g)
 125                                    | r <- [0..(n - 1)], c <- [0..(n - 1)]
 126                                    ]
 127     where rStep = bigR * n
 128           cStep = bigC * n
 129
 130 -- merge a set of subgrids into one
 131 contractExploded :: ExplodedGrid -> Grid
 132 -- contractExploded gs = foldl1 (>|<) $ map (foldl1 (>-<)) rows
 133 contractExploded gs = foldl1 (>|<) $ parMap rpar (foldl1 (>-<)) rows
 134     where rows = explodedRows gs
 135
 136 -- find the rows of an exploded grid
 137 explodedRows :: ExplodedGrid -> [ExplodedGrid]
 138 explodedRows eg = [M.filterWithKey (\(r, _) _ -> r == row) eg | row <- [0..rowMax] ]
 139     where (rowMax, _) = bounds eg
 140
 141 -- merge two grids horizontally
 142 (>-<) :: Grid -> Grid -> Grid
 143 (>-<) g1 g2 = M.union g1 g2'
 144     where (_, cm) = bounds g1
 145           g2' = M.mapKeys (\(r, c) -> (r, c + cm + 1)) g2
 146
 147 -- merge two grids vertically
 148 (>|<) :: Grid -> Grid -> Grid
 149 (>|<) g1 g2 = M.union g1 g2'
 150     where (rm, _) = bounds g1
 151           g2' = M.mapKeys (\(r, c) -> (r + rm + 1, c)) g2
 152
 153
 154
 155
 156 bounds :: M.Map (Int, Int) a -> (Int, Int)
 157 bounds grid = (maximum $ map fst $ M.keys grid, maximum $ map snd $ M.keys grid)
 158
 159
 160 showGrid :: Grid -> String
 161 showGrid g = unlines [[showGChar $ M.findWithDefault False (r, c) g |
 162                 c <- [0..cm] ] | r <- [0..rm] ]
 163     where (rm, cm) = bounds g
 164           showGChar True = '#'
 165           showGChar False = '.'
 166
 167
 168
 169 -- really persuade Megaparsec not to include newlines in how it consume spaces.
 170 onlySpace = (char ' ') <|> (char '\t')
 171
 172 sc :: Parser ()
 173 sc = L.space (skipSome onlySpace) CA.empty CA.empty
 174
 175 lexeme  = L.lexeme sc
 176
 177 symbol = L.symbol sc
 178 rowSep = symbol "/"
 179 ruleJoin = symbol "=>"
 180
 181 present = id True <$ symbol "#"
 182 absent = id False <$ symbol "."
 183
 184 rulesP = ruleP `sepBy` space
 185 ruleP = Rule <$> gridP <* ruleJoin <*> gridP
 186
 187 gridP = gridify <$> rowP `sepBy` rowSep
 188     where gridify g = M.fromList $ concat
 189                                     [map (\(c, v) -> ((r, c), v)) nr |
 190                                              (r, nr) <- zip [0..]
 191                                                             [zip [0..] r | r <- g]]
 192
 193
 194 rowP = some (present <|> absent)
 195
 196 successfulParse :: Text -> [Rule]
 197 successfulParse input =
 198         case parse rulesP "input" input of
 199                 Left  _error -> []
 200                 Right instructions  -> instructions