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