eed9e357c5fc0726c9cdd47db7ce5b7ffca9fd0a
[advent-of-code-23.git] / 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