advent24/MainLax.hs

   1 -- Writeup at https://work.njae.me.uk/2021/12/29/advent-of-code-2021-day-24/
   2 -- Based on ideas by Daniel Lin,
   3 --   taken from https://github.com/ephemient/aoc2021/blob/main/hs/src/Day24.hs
   4
   5 import Debug.Trace
   6
   7 import Data.Text (Text)
   8 import qualified Data.Text.IO as TIO
   9 import Data.Attoparsec.Text -- hiding (take, takeWhile)
  10 import Control.Applicative
  11 import qualified Data.Map as M
  12 import Data.Map ((!))
  13 import Data.List
  14 import Control.Monad
  15 import Data.Maybe
  16
  17 data Register = W | X | Y | Z deriving (Eq, Ord, Show, Enum)
  18
  19 data Interval = Interval Integer Integer
  20   deriving (Eq, Ord, Show)
  21
  22 data Argument = Reg Register | Lit Integer
  23   deriving (Eq, Ord, Show)
  24
  25 data Instruction
  26   = Inp Register
  27   | Add Register Argument
  28   | Mul Register Argument
  29   | Div Register Argument
  30   | Mod Register Argument
  31   | Eql Register Argument
  32   deriving (Eq, Ord, Show)
  33
  34 type LiteralMachine = M.Map Register Integer
  35 type IntervalMachine = M.Map Register (Maybe Interval)
  36
  37 data CodeMachine = CodeMachine
  38   { mCode :: [Integer]
  39   , mMachine :: LiteralMachine
  40   } deriving (Show)
  41
  42 -- Main
  43
  44 main :: IO ()
  45 main =
  46   do  text <- TIO.readFile "data/advent24.txt"
  47       let instrs = successfulParse text
  48       let m0 = CodeMachine {mCode = [], mMachine = emptyMachine}
  49       putStrLn $ part1 m0 instrs
  50       putStrLn $ part2 m0 instrs
  51
  52 part1 :: CodeMachine -> [Instruction] -> String
  53 part1 = findCode [9, 8..1]
  54
  55 part2 :: CodeMachine -> [Instruction] -> String
  56 part2 = findCode [1..9]
  57
  58 findCode :: [Integer] -> CodeMachine -> [Instruction] -> String
  59 findCode digits machine instrs = concatMap show $ mCode $ head $ runLit instrs digits machine
  60
  61 plausible :: [Instruction] -> LiteralMachine -> Bool
  62 plausible instrs litMachine = feasible ranMachine
  63   where intMachine = intervalify litMachine
  64         ranMachine = runInt instrs intMachine
  65
  66 feasible :: IntervalMachine -> Bool
  67 -- feasible machine | trace ("Feasible " ++ (show machine)) False = undefined
  68 feasible machine
  69   | (w && x && y && isJust z) = a <= 0 && b >= 0
  70   | otherwise = False
  71   where w = isJust $ machine ! W
  72         x = isJust $ machine ! X
  73         y = isJust $ machine ! Y
  74         z = machine ! Z
  75         Just (Interval a b) = z
  76
  77 valid :: CodeMachine -> Bool
  78 valid (CodeMachine{..}) = (mMachine ! Z) == 0
  79
  80
  81 emptyMachine :: LiteralMachine
  82 emptyMachine = M.fromList [(r, 0) | r <- [W .. Z]]
  83
  84 intervalify :: LiteralMachine -> IntervalMachine
  85 intervalify = M.map (\i -> Just (Interval i i))
  86
  87
  88 runLit :: [Instruction] -> [Integer] -> CodeMachine -> [CodeMachine]
  89 -- runLit instrs _digits m0 | trace ((show $ length instrs) ++ " " ++ (show m0)) False = undefined
  90 runLit [] _ machine = [machine]
  91 runLit (Inp reg : instrs) digits (CodeMachine {..}) =
  92   do guard (plausible (Inp reg : instrs) mMachine)
  93      i <- digits
  94      let m1 = M.insert reg i mMachine
  95      mm2 <- runLit instrs digits (CodeMachine { mCode = mCode ++ [i], mMachine = m1})
  96      guard (valid mm2)
  97      return mm2
  98 runLit (Add reg arg : instrs) digits (CodeMachine {..}) =
  99   runLit instrs digits (CodeMachine{mMachine = M.insert reg c mMachine, ..})
 100   where a = mMachine ! reg
 101         b = evaluateLit arg mMachine
 102         c = a + b
 103 runLit (Mul reg arg : instrs) digits (CodeMachine {..}) =
 104   runLit instrs digits (CodeMachine{mMachine = M.insert reg c mMachine, ..})
 105   where a = mMachine ! reg
 106         b = evaluateLit arg mMachine
 107         c = a * b
 108 runLit (Div reg arg : instrs) digits (CodeMachine {..}) =
 109   runLit instrs digits (CodeMachine{mMachine = M.insert reg c mMachine, ..})
 110   where a = mMachine ! reg
 111         b = evaluateLit arg mMachine
 112         c = a `quot` b
 113 runLit (Mod reg arg : instrs) digits (CodeMachine {..}) =
 114   runLit instrs digits (CodeMachine{mMachine = M.insert reg c mMachine, ..})
 115   where a = mMachine ! reg
 116         b = evaluateLit arg mMachine
 117         c = a `rem` b
 118 runLit (Eql reg arg : instrs) digits (CodeMachine {..}) =
 119   runLit instrs digits (CodeMachine{mMachine = M.insert reg c mMachine, ..})
 120   where a = mMachine ! reg
 121         b = evaluateLit arg mMachine
 122         c = if a == b then 1 else 0
 123
 124
 125 runInt :: [Instruction] -> IntervalMachine -> IntervalMachine
 126 runInt instrs machine = foldl' interpretInt machine instrs
 127
 128 interpretInt :: IntervalMachine -> Instruction -> IntervalMachine
 129 -- interpretInt machine instr | trace ("iInt " ++ (show instr) ++ " " ++  (show machine)) False = undefined
 130 interpretInt machine (Inp reg) = M.insert reg (Just (Interval 1 9)) machine
 131 interpretInt machine (Add reg arg) = M.insert reg c machine
 132   where a = machine ! reg
 133         b = evaluateInt arg machine
 134         c = join $ (+:) <$> a <*> b
 135         -- c = join $ (liftM2 (+:)) a b
 136 interpretInt machine (Mul reg arg) = M.insert reg c machine
 137   where a = machine ! reg
 138         b = evaluateInt arg machine
 139         c = join $ (*:) <$> a <*> b
 140 interpretInt machine (Div reg arg) = M.insert reg c machine
 141   where a = machine ! reg
 142         b = evaluateInt arg machine
 143         c = join $ (/:) <$> a <*> b
 144 interpretInt machine (Mod reg arg) = M.insert reg c machine
 145   where a = machine ! reg
 146         b = evaluateInt arg machine
 147         c = join $ (%:) <$> a <*> b
 148 interpretInt machine (Eql reg arg) = M.insert reg c machine
 149   where a = machine ! reg
 150         b = evaluateInt arg machine
 151         c = join $ (=:) <$> a <*> b
 152
 153 (+:), (*:), (/:), (%:), (=:) :: Interval -> Interval -> Maybe Interval
 154 (Interval a b) +: (Interval c d) = Just (Interval (a + c) (b + d))
 155 (Interval a b) *: (Interval c d)
 156   | a >= 0 && c >= 0 = Just ( Interval (a * c) (b * d) )
 157   | b <= 0 && d <= 0 = Just ( Interval (b * d) (a * c) )
 158   | a >= 0 && d <= 0 = Just ( Interval (a * d) (b * c) )
 159   | b <= 0 && c >= 0 = Just ( Interval (b * c) (a * d) )
 160 (Interval a b) /: (Interval c d)
 161   | c > 0 = Just ( Interval (a `quot` d) (b `quot` c) )
 162   | d < 0 = Just ( Interval (a `quot` c) (b `quot` d) )
 163   | otherwise = Nothing
 164 (Interval _a _b) %: (Interval c d)
 165   | c > 0 = Just ( Interval 0 (d - 1))
 166   | otherwise = Nothing
 167 (Interval a b) =: (Interval c d)
 168   | b < c = Just (Interval 0 0)
 169   | a > d = Just (Interval 0 0)
 170   | a == b && a == c && a == d = Just (Interval 1 1)
 171   | otherwise = Just (Interval 0 1)
 172
 173 evaluateLit :: Argument -> LiteralMachine -> Integer
 174 evaluateLit (Reg reg) machine =  machine ! reg
 175 evaluateLit (Lit n) _ = n
 176
 177 evaluateInt :: Argument -> IntervalMachine -> Maybe Interval
 178 evaluateInt (Reg reg) machine =  machine ! reg
 179 evaluateInt (Lit n) _ = Just (Interval n n)
 180
 181
 182 -- Parse the input file
 183
 184 instructionsP:: Parser [Instruction]
 185 instructionsP = instructionP `sepBy` endOfLine
 186
 187 instructionP:: Parser Instruction
 188 instructionP = choice [inpP, addP, mulP, divP, modP, eqlP]
 189
 190 inpP, addP, mulP, divP, modP, eqlP :: Parser Instruction
 191 inpP = Inp <$> ("inp " *> registerP)
 192 addP = Add <$> ("add " *> registerP) <*> (" " *> argumentP)
 193 mulP = Mul <$> ("mul " *> registerP) <*> (" " *> argumentP)
 194 divP = Div <$> ("div " *> registerP) <*> (" " *> argumentP)
 195 modP = Mod <$> ("mod " *> registerP) <*> (" " *> argumentP)
 196 eqlP = Eql <$> ("eql " *> registerP) <*> (" " *> argumentP)
 197
 198 registerP, wP, xP, yP, zP :: Parser Register
 199 registerP = choice [wP, xP, yP, zP]
 200 wP = "w" *> pure W
 201 xP = "x" *> pure X
 202 yP = "y" *> pure Y
 203 zP = "z" *> pure Z
 204
 205 argumentP :: Parser Argument
 206 argumentP = (Reg <$> registerP) <|> (Lit <$> signed decimal)
 207
 208 successfulParse :: Text -> [Instruction]
 209 successfulParse input =
 210   case parseOnly instructionsP input of
 211     Left  _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
 212     Right instrs -> instrs