advent16/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2021/12/18/advent-of-code-2021-day-16/
   2
   3 import Data.Word
   4 import Data.Bits
   5 import Data.Char
   6 import Data.Int
   7
   8 import Control.Monad.State.Lazy
   9
  10 import qualified Data.ByteString as BYS
  11 import qualified Data.Bitstream as BS
  12
  13 type Transmission = BS.Bitstream BS.Right
  14
  15 type ParseTrans = State Transmission
  16
  17 data Packet = Packet Integer PacketContents
  18   deriving (Show, Eq)
  19
  20 data PacketContents
  21   = Literal Integer
  22   -- | Operator [Packet]
  23   | Sum [Packet]
  24   | Product [Packet]
  25   | Minimum [Packet]
  26   | Maximum [Packet]
  27   | GreaterThan Packet Packet
  28   | LessThan Packet Packet
  29   | EqualTo Packet Packet
  30   deriving (Show, Eq)
  31
  32 main :: IO ()
  33 main =
  34   do  text <- readFile "data/advent16.txt"
  35       let packetStream = bitify text
  36       let (packet, _remaining) = runState getPacket packetStream
  37       print $ part1 packet
  38       print $ part2 packet
  39
  40 part1 :: Packet -> Integer
  41 part1 = packetVersionSum
  42
  43 part2 :: Packet -> Integer
  44 part2 = evaluatePacket
  45
  46 bitify :: String -> Transmission
  47 bitify = BS.fromByteString . BYS.pack . hexPack . map (fromIntegral . digitToInt)
  48
  49 hexPack :: [Word8] -> [Word8]
  50 hexPack [] = []
  51 hexPack (x:[]) = hexPack [x, 0]
  52 hexPack (x:y:xs) = ((x `shift` 4) .|. y) : hexPack xs
  53
  54 getBool :: ParseTrans Bool
  55 getBool =
  56   do  bs <- get
  57       let value = head $ BS.unpack $ BS.take 1 bs
  58       put $ BS.drop 1 bs
  59       return value
  60
  61 getInt :: Int64 -> ParseTrans Integer
  62 getInt n =
  63   do  bs <- get
  64       let value = BS.toBits $ BS.take n bs
  65       put $ BS.drop n bs
  66       return value
  67
  68 getBits :: Int64 -> ParseTrans Transmission
  69 getBits n =
  70   do  bs <- get
  71       let bits = BS.take n bs
  72       put $ BS.drop n bs
  73       return bits
  74
  75 getPacket :: ParseTrans Packet
  76 getPacket =
  77   do version <- getInt 3
  78      pType <- getInt 3
  79      payload <- if pType == 4
  80                 then do val <- getLiteral
  81                         return $ Literal val
  82                 else do contents <- getOperatorContents
  83                         return $ mkOperator pType contents
  84      return $ Packet version payload
  85
  86 getLiteral :: ParseTrans Integer
  87 getLiteral = getLiteralAcc 0
  88
  89 getLiteralAcc :: Integer -> ParseTrans Integer
  90 getLiteralAcc acc =
  91   do  continues <- getBool
  92       nybble <- getInt 4
  93       let acc' = acc * 16 + nybble
  94       if continues
  95       then do getLiteralAcc acc'
  96       else return acc'
  97
  98 getOperatorContents :: ParseTrans [Packet]
  99 getOperatorContents =
 100   do isNumPackets <- getBool
 101      if isNumPackets
 102      then do numPackets <- getInt 11
 103              getPacketsByCount numPackets
 104      else do numBits <- getInt 15
 105              subString <- getBits (fromIntegral numBits)
 106              return $ getPacketsByLength subString
 107
 108 getPacketsByLength :: Transmission -> [Packet]
 109 getPacketsByLength bits
 110   | BS.null bits = []
 111   | otherwise = p : (getPacketsByLength remaining)
 112   where (p, remaining) = runState getPacket bits
 113
 114 getPacketsByCount :: Integer -> ParseTrans [Packet]
 115 getPacketsByCount 0 = return []
 116 getPacketsByCount n =
 117   do p <- getPacket
 118      ps <- getPacketsByCount (n - 1)
 119      return (p : ps)
 120
 121 mkOperator :: Integer -> [Packet] -> PacketContents
 122 mkOperator pType contents = case pType of
 123   0 -> Sum contents
 124   1 -> Product contents
 125   2 -> Minimum contents
 126   3 -> Maximum contents
 127   5 -> GreaterThan (contents!!0) (contents!!1)
 128   6 -> LessThan (contents!!0) (contents!!1)
 129   7 -> EqualTo (contents!!0) (contents!!1)
 130
 131
 132 packetVersionSum :: Packet -> Integer
 133 packetVersionSum (Packet version contents) =
 134   version + (contentsVersionSum contents)
 135
 136 contentsVersionSum :: PacketContents -> Integer
 137 contentsVersionSum (Sum packets) = sum $ map packetVersionSum packets
 138 contentsVersionSum (Product packets) = sum $ map packetVersionSum packets
 139 contentsVersionSum (Minimum packets) = sum $ map packetVersionSum packets
 140 contentsVersionSum (Maximum packets) = sum $ map packetVersionSum packets
 141 contentsVersionSum (Literal _) = 0
 142 contentsVersionSum (GreaterThan packet1 packet2) =
 143   (packetVersionSum packet1) + (packetVersionSum packet2)
 144 contentsVersionSum (LessThan packet1 packet2) =
 145   (packetVersionSum packet1) + (packetVersionSum packet2)
 146 contentsVersionSum (EqualTo packet1 packet2) =
 147   (packetVersionSum packet1) + (packetVersionSum packet2)
 148
 149 evaluatePacket :: Packet -> Integer
 150 evaluatePacket (Packet _version contents) = evaluateContents contents
 151
 152 evaluateContents :: PacketContents -> Integer
 153 evaluateContents (Sum packets) = sum $ map evaluatePacket packets
 154 evaluateContents (Product packets) = product $ map evaluatePacket packets
 155 evaluateContents (Minimum packets) = minimum $ map evaluatePacket packets
 156 evaluateContents (Maximum packets) = maximum $ map evaluatePacket packets
 157 evaluateContents (Literal n) = n
 158 evaluateContents (GreaterThan packet1 packet2) =
 159   if (evaluatePacket packet1) > (evaluatePacket packet2) then 1 else 0
 160 evaluateContents (LessThan packet1 packet2) =
 161   if (evaluatePacket packet1) < (evaluatePacket packet2) then 1 else 0
 162 evaluateContents (EqualTo packet1 packet2) =
 163   if (evaluatePacket packet1) == (evaluatePacket packet2) then 1 else 0