advent19/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2023/12/22/advent-of-code-2023-day-18/
   2
   3 import AoC
   4
   5 import Data.Text (Text, unpack)
   6 import qualified Data.Text.IO as TIO
   7 import Data.Attoparsec.Text
   8 import qualified Data.Attoparsec.Text as AT
   9 import Control.Applicative
  10 import qualified Data.Map.Strict as M
  11 import Data.Map.Strict ((!))
  12 import Control.Lens
  13 import Data.Semigroup
  14 import Data.Monoid
  15 import Data.Semigroup (Semigroup)
  16
  17 data Interval = Interval { _low :: Int, _high :: Int }
  18   deriving (Eq, Ord, Show)
  19 makeLenses ''Interval
  20
  21 data Part a = Part { _x :: a, _m :: a, _a :: a, _s :: a }
  22   deriving (Eq, Ord, Show)
  23 makeLenses ''Part
  24
  25 data Register = X | M | A | S
  26   deriving (Eq, Ord, Show)
  27
  28 data Comparator = Lt | Gt
  29   deriving (Eq, Ord, Show)
  30
  31 type RuleBase = M.Map String [RuleElement]
  32
  33 data Destination = Accept | Reject | Rule String | Continue
  34   deriving (Eq, Ord, Show)
  35
  36 data RuleElement = WithTest Test Destination
  37                  | WithoutTest Destination
  38   deriving (Eq, Ord, Show)
  39
  40 data Test = Test { _register :: Register
  41                  , _comparator :: Comparator
  42                  , _threshold :: Int
  43                  }
  44   deriving (Eq, Ord, Show)
  45 makeLenses ''Test
  46
  47 data WaitingPart = WaitingPart String (Part Interval)
  48   deriving (Eq, Ord, Show)
  49
  50 data Evaluation = Evaluation
  51   { _accepted :: [Part Interval]
  52   , _waiting :: [WaitingPart]
  53   } deriving (Eq, Ord, Show)
  54 makeLenses ''Evaluation
  55
  56 instance Semigroup Evaluation where
  57   (Evaluation a1 w1) <> (Evaluation a2 w2) = Evaluation (a1 <> a2) (w1 <> w2)
  58
  59 instance Monoid Evaluation where
  60   mempty = Evaluation [] []
  61
  62 main :: IO ()
  63 main =
  64   do  dataFileName <- getDataFileName
  65       text <- TIO.readFile dataFileName
  66       let (rules, parts) = successfulParse text
  67       -- print rules
  68       -- print parts
  69       -- print $ fmap (applyWorkflow "in" rules) parts
  70       -- print $ filter ((== Accept) . applyWorkflow "in" rules) parts
  71       print $ part1 rules parts
  72       print $ part2 rules
  73       -- print $ part2 text
  74
  75 part1 :: RuleBase -> [Part Int] -> Int
  76 part1 rules parts = sum $ fmap sumRegisters acceptedParts
  77   where acceptedParts = filter ((== Accept) . applyWorkflow "in" rules) parts
  78
  79 part2 rules = sum $ fmap registerRange accepted
  80   where accepted = evaluateRules rules
  81                                  (Evaluation [] [(WaitingPart "in" initialPart)])
  82
  83
  84 applyWorkflow :: String -> RuleBase -> Part Int -> Destination
  85 applyWorkflow name rules part =
  86   case applyRule part (rules ! name) of
  87     Rule name' -> applyWorkflow name' rules part
  88     dest -> dest
  89
  90
  91 applyRule :: Part Int -> [RuleElement] -> Destination
  92 applyRule _ [] = Reject
  93 applyRule part (x:xs) =
  94   case applyElement part x of
  95     Continue -> applyRule part xs
  96     dest -> dest
  97
  98 applyElement :: Part Int -> RuleElement -> Destination
  99 applyElement _ (WithoutTest dest) = dest
 100 applyElement part (WithTest test dest)
 101   | (test ^. comparator == Lt) =
 102       if (regValue part (test ^. register) < (test ^. threshold))
 103         then dest
 104         else Continue
 105   | otherwise =
 106       if (regValue part (test ^. register) > (test ^. threshold))
 107         then dest
 108         else Continue
 109
 110 regValue :: Part a -> Register -> a
 111 regValue part X = part ^. x
 112 regValue part M = part ^. m
 113 regValue part A = part ^. a
 114 regValue part S = part ^. s
 115
 116
 117 -- lensOfR :: Register -> Lens' (Part a) a
 118 -- lensOfR X = x
 119 -- lensOfR M = m
 120 -- lensOfR A = a
 121 -- lensOfR S = s
 122
 123 sumRegisters :: Part Int -> Int
 124 sumRegisters part = (part ^. x) + (part ^. m) + (part ^. a) + (part ^. s)
 125
 126 registerRange :: Part Interval -> Int
 127 registerRange part = ((part ^. x . high) - (part ^. x . low) + 1) *
 128                      ((part ^. m . high) - (part ^. m . low) + 1) *
 129                      ((part ^. a . high) - (part ^. a . low) + 1) *
 130                      ((part ^. s . high) - (part ^. s . low) + 1)
 131
 132 initialPart :: Part Interval
 133 initialPart = Part (Interval 1 4000) (Interval 1 4000)
 134                    (Interval 1 4000) (Interval 1 4000)
 135
 136
 137 evaluateRules :: RuleBase -> Evaluation -> [Part Interval]
 138 evaluateRules rules (Evaluation accepted []) = accepted
 139 evaluateRules rules (Evaluation accepted ((WaitingPart rulename part):waiting)) =
 140   evaluateRules rules ((Evaluation accepted waiting) <> newEvaluation)
 141   where rulebody = rules ! rulename
 142         newEvaluation = applyRuleI part rulebody
 143
 144 applyRuleI :: (Part Interval) -> [RuleElement] -> Evaluation
 145 applyRuleI part [] = mempty
 146 applyRuleI part (x:xs) =
 147   case inProgress of
 148     Nothing -> evaluation
 149     Just p -> evaluation <> (applyRuleI p xs)
 150   where (evaluation, inProgress) = applyElementI part x
 151
 152 applyElementI :: (Part Interval) -> RuleElement -> (Evaluation, Maybe (Part Interval))
 153 applyElementI part (WithoutTest Accept) = (mempty & accepted .~ [part], Nothing)
 154 applyElementI part (WithoutTest Reject) = (mempty, Nothing)
 155 applyElementI part (WithoutTest (Rule rule)) = (mempty & waiting .~ [WaitingPart rule part], Nothing)
 156 applyElementI part (WithTest test dest) = (evaluation, failing)
 157   where (passing, failing) = splitPart part test
 158         evaluation = case passing of
 159                       Nothing -> mempty
 160                       Just p -> fst $ applyElementI p (WithoutTest dest)
 161
 162 splitPart :: Part Interval -> Test -> (Maybe (Part Interval), Maybe (Part Interval))
 163 splitPart part test = (passingPart, failingPart)
 164   where -- regLens = lensOfR $ test ^. register
 165         (passingInterval, failingInterval) =
 166           -- splitInterval (part ^. regLens) (test ^. comparator) (test ^. threshold)
 167           splitInterval (regValue part (test ^. register)) (test ^. comparator) (test ^. threshold)
 168         passingPart = case passingInterval of
 169                         Nothing -> Nothing
 170                         -- Just interval -> Just (part & regLens .~ interval)
 171                         Just interval -> Just (setRegister part test interval)
 172         failingPart = case failingInterval of
 173                         Nothing -> Nothing
 174                         -- Just interval -> Just (part & regLens .~ interval)
 175                         Just interval -> Just (setRegister part test interval)
 176
 177 splitInterval :: Interval -> Comparator -> Int -> (Maybe Interval, Maybe Interval)
 178 splitInterval interval Lt threshold
 179   | (interval ^. high) < threshold = (Just interval, Nothing)
 180   | (interval ^. low) > threshold = (Nothing, Just interval)
 181   | otherwise = ( Just (Interval (interval ^. low) (threshold - 1))
 182                 , Just (Interval threshold (interval ^. high))
 183                 )
 184 splitInterval interval Gt threshold
 185   | (interval ^. low) > threshold = (Just interval, Nothing)
 186   | (interval ^. high) < threshold = (Nothing, Just interval)
 187   | otherwise = ( Just (Interval (threshold + 1) (interval ^. high))
 188                 , Just (Interval (interval ^. low) threshold)
 189                 )
 190
 191
 192 setRegister :: Part Interval -> Test -> Interval -> Part Interval
 193 setRegister part (Test X _ _) val = part & x .~ val
 194 setRegister part (Test M _ _) val = part & m .~ val
 195 setRegister part (Test A _ _) val = part & a .~ val
 196 setRegister part (Test S _ _) val = part & s .~ val
 197
 198
 199 -- Parse the input file
 200
 201 rulePartP :: Parser (RuleBase, [Part Int])
 202 rulesP :: Parser RuleBase
 203 ruleP :: Parser (String, [RuleElement])
 204 nameP :: Parser String
 205 ruleBodyP :: Parser [RuleElement]
 206 ruleElementP, withTestP, withoutTestP :: Parser RuleElement
 207 testP :: Parser Test
 208 registerP :: Parser Register
 209 destinationP :: Parser Destination
 210 comparatorP :: Parser Comparator
 211 partsP :: Parser [Part Int]
 212 partP :: Parser (Part Int)
 213
 214 rulePartP = (,) <$> (rulesP <* endOfLine <* endOfLine) <*> partsP
 215
 216 rulesP = M.fromList <$> ruleP `sepBy` endOfLine
 217 ruleP = (,) <$> (nameP <* "{") <*> (ruleBodyP <* "}")
 218
 219 nameP = unpack <$> AT.takeWhile (inClass "a-z") -- many1 letter
 220 ruleBodyP = ruleElementP `sepBy` ","
 221 ruleElementP = withTestP <|> withoutTestP
 222
 223 withTestP = WithTest <$> (testP <* ":") <*> destinationP
 224 withoutTestP = WithoutTest <$> destinationP
 225
 226 testP = Test <$> registerP <*> comparatorP <*> decimal
 227
 228 registerP = choice [ X <$ "x"
 229                    , M <$ "m"
 230                    , A <$ "a"
 231                    , S <$ "s"
 232                    ]
 233
 234 destinationP = choice [ Accept <$ "A"
 235                       , Reject <$ "R"
 236                       , Rule <$> nameP
 237                       ]
 238
 239 comparatorP = choice [ Lt <$ "<"
 240                      , Gt <$ ">"
 241                      ]
 242
 243 partsP = partP `sepBy` endOfLine
 244 partP = Part <$> ("{x=" *> decimal) <*> (",m=" *> decimal) <*> (",a=" *> decimal) <*> (",s=" *> decimal <* "}")
 245
 246 successfulParse :: Text -> (RuleBase, [Part Int])
 247 successfulParse input =
 248   case parseOnly rulePartP input of
 249     Left  _err -> (M.empty, []) -- TIO.putStr $ T.pack $ parseErrorPretty err
 250     Right matches -> matches