advent19/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2023/12/24/advent-of-code-2023-day-19/
   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       if part ^. l < test ^. threshold
 104         then dest
 105         else Continue
 106   | otherwise =
 107       -- if (regValue part (test ^. register) > (test ^. threshold))
 108       if part ^. l > test ^. threshold
 109         then dest
 110         else Continue
 111   where l = lensOfR (test ^. register)
 112
 113 regValue :: Part a -> Register -> a
 114 regValue part X = part ^. x
 115 regValue part M = part ^. m
 116 regValue part A = part ^. a
 117 regValue part S = part ^. s
 118
 119
 120 lensOfR :: Register -> Lens' (Part a) a
 121 lensOfR X = x
 122 lensOfR M = m
 123 lensOfR A = a
 124 lensOfR S = s
 125
 126 sumRegisters :: Part Int -> Int
 127 sumRegisters part = (part ^. x) + (part ^. m) + (part ^. a) + (part ^. s)
 128
 129 registerRange :: Part Interval -> Int
 130 registerRange part = ((part ^. x . high) - (part ^. x . low) + 1) *
 131                      ((part ^. m . high) - (part ^. m . low) + 1) *
 132                      ((part ^. a . high) - (part ^. a . low) + 1) *
 133                      ((part ^. s . high) - (part ^. s . low) + 1)
 134
 135 initialPart :: Part Interval
 136 initialPart = Part (Interval 1 4000) (Interval 1 4000)
 137                    (Interval 1 4000) (Interval 1 4000)
 138
 139
 140 evaluateRules :: RuleBase -> Evaluation -> [Part Interval]
 141 evaluateRules rules (Evaluation accepted []) = accepted
 142 evaluateRules rules (Evaluation accepted ((WaitingPart rulename part):waiting)) =
 143   evaluateRules rules ((Evaluation accepted waiting) <> newEvaluation)
 144   where rulebody = rules ! rulename
 145         newEvaluation = applyRuleI part rulebody
 146
 147 applyRuleI :: Part Interval -> [RuleElement] -> Evaluation
 148 applyRuleI _ [] = mempty
 149 applyRuleI part (e:es) =
 150   case inProgress of
 151     Nothing -> evaluation
 152     Just p -> evaluation <> (applyRuleI p es)
 153   where (evaluation, inProgress) = applyElementI part e
 154
 155 applyElementI :: Part Interval -> RuleElement -> (Evaluation, Maybe (Part Interval))
 156 applyElementI part (WithoutTest Accept) = (mempty & accepted .~ [part], Nothing)
 157 applyElementI part (WithoutTest Reject) = (mempty, Nothing)
 158 applyElementI part (WithoutTest (Rule rule)) = (mempty & waiting .~ [WaitingPart rule part], Nothing)
 159 applyElementI part (WithTest test dest) = (evaluation, failing)
 160   where (passing, failing) = splitPart part test
 161         evaluation = case passing of
 162                       Nothing -> mempty
 163                       Just p -> fst $ applyElementI p (WithoutTest dest)
 164
 165 splitPart :: Part Interval -> Test -> (Maybe (Part Interval), Maybe (Part Interval))
 166 splitPart part test = (passingPart, failingPart)
 167   where l = lensOfR (test ^. register) :: Lens' (Part Interval) Interval
 168         (passingInterval, failingInterval) =
 169           -- splitInterval (part ^. l) (test ^. comparator) (test ^. threshold)
 170           splitInterval (regValue part (test ^. register)) (test ^. comparator) (test ^. threshold)
 171         passingPart = case passingInterval of
 172                         Nothing -> Nothing
 173                         Just pi -> Just (part & l .~ pi)
 174                         -- Just interval -> Just (setRegister part test interval)
 175         failingPart = case failingInterval of
 176                         Nothing -> Nothing
 177                         Just fi -> Just (part & l .~ fi)
 178                         -- Just interval -> Just (setRegister part test interval)
 179
 180 splitInterval :: Interval -> Comparator -> Int -> (Maybe Interval, Maybe Interval)
 181 splitInterval interval Lt threshold
 182   | (interval ^. high) < threshold = (Just interval, Nothing)
 183   | (interval ^. low) > threshold = (Nothing, Just interval)
 184   | otherwise = ( Just (Interval (interval ^. low) (threshold - 1))
 185                 , Just (Interval threshold (interval ^. high))
 186                 )
 187 splitInterval interval Gt threshold
 188   | (interval ^. low) > threshold = (Just interval, Nothing)
 189   | (interval ^. high) < threshold = (Nothing, Just interval)
 190   | otherwise = ( Just (Interval (threshold + 1) (interval ^. high))
 191                 , Just (Interval (interval ^. low) threshold)
 192                 )
 193
 194
 195 -- setRegister :: Part Interval -> Test -> Interval -> Part Interval
 196 -- setRegister part (Test X _ _) val = part & x .~ val
 197 -- setRegister part (Test M _ _) val = part & m .~ val
 198 -- setRegister part (Test A _ _) val = part & a .~ val
 199 -- setRegister part (Test S _ _) val = part & s .~ val
 200
 201
 202 -- Parse the input file
 203
 204 rulePartP :: Parser (RuleBase, [Part Int])
 205 rulesP :: Parser RuleBase
 206 ruleP :: Parser (String, [RuleElement])
 207 nameP :: Parser String
 208 ruleBodyP :: Parser [RuleElement]
 209 ruleElementP, withTestP, withoutTestP :: Parser RuleElement
 210 testP :: Parser Test
 211 registerP :: Parser Register
 212 destinationP :: Parser Destination
 213 comparatorP :: Parser Comparator
 214 partsP :: Parser [Part Int]
 215 partP :: Parser (Part Int)
 216
 217 rulePartP = (,) <$> (rulesP <* endOfLine <* endOfLine) <*> partsP
 218
 219 rulesP = M.fromList <$> ruleP `sepBy` endOfLine
 220 ruleP = (,) <$> (nameP <* "{") <*> (ruleBodyP <* "}")
 221
 222 nameP = unpack <$> AT.takeWhile (inClass "a-z") -- many1 letter
 223 ruleBodyP = ruleElementP `sepBy` ","
 224 ruleElementP = withTestP <|> withoutTestP
 225
 226 withTestP = WithTest <$> (testP <* ":") <*> destinationP
 227 withoutTestP = WithoutTest <$> destinationP
 228
 229 testP = Test <$> registerP <*> comparatorP <*> decimal
 230
 231 registerP = choice [ X <$ "x"
 232                    , M <$ "m"
 233                    , A <$ "a"
 234                    , S <$ "s"
 235                    ]
 236
 237 destinationP = choice [ Accept <$ "A"
 238                       , Reject <$ "R"
 239                       , Rule <$> nameP
 240                       ]
 241
 242 comparatorP = choice [ Lt <$ "<"
 243                      , Gt <$ ">"
 244                      ]
 245
 246 partsP = partP `sepBy` endOfLine
 247 partP = Part <$> ("{x=" *> decimal) <*> (",m=" *> decimal) <*> (",a=" *> decimal) <*> (",s=" *> decimal <* "}")
 248
 249 successfulParse :: Text -> (RuleBase, [Part Int])
 250 successfulParse input =
 251   case parseOnly rulePartP input of
 252     Left  _err -> (M.empty, []) -- TIO.putStr $ T.pack $ parseErrorPretty err
 253     Right matches -> matches