2021: day16: part 1

aoc2021/src/day16.rs

@@ -1,4 +1,5 @@
 use std::fmt::Write;
+use std::ops::Range;
 
 use anyhow::{Context, Result};
 use bitvec::prelude::*;
@@ -29,14 +30,18 @@ enum PacketType {
 struct Packet {
     version: u8,
     type_id: u8,
+    bits_used: usize,
     packet: PacketType,
 }
 
 impl Packet {
     fn version_sum(&self) -> u64 {
+        let version = self.version as u64;
         match &self.packet {
-            PacketType::Litteral(_) => self.version as u64,
+            PacketType::Litteral(_) => version,
-            PacketType::Operator(op) => op.sub_packets.iter().map(Packet::version_sum).sum(),
+            PacketType::Operator(op) => {
+                version + op.sub_packets.iter().map(Packet::version_sum).sum::<u64>()
+            }
         }
     }
 }
@@ -69,6 +74,20 @@ impl std::str::FromStr for Packet {
     }
 }
 
+const VERSION_RANGE: Range<usize> = 0..3;
+const TYPE_ID_RANGE: Range<usize> = 3..6;
+
+const TYPE_LENGTH_ID_INDEX: usize = 6;
+const LENGTH_INDEX: usize = 6 + 1;
+
+const TLID0_SUBPACKET_LENGTH_BITS: usize = 15;
+const SUBPACKET_START_INDEX_TLID0: usize = LENGTH_INDEX + TLID0_SUBPACKET_LENGTH_BITS;
+const SUBPACKETS_BIT_LENGTH_RANGE: Range<usize> = LENGTH_INDEX..SUBPACKET_START_INDEX_TLID0;
+
+const TLID1_SUBPACKET_NUMBER_BITS: usize = 11;
+const SUBPACKET_START_INDEX_TLID1: usize = LENGTH_INDEX + TLID1_SUBPACKET_NUMBER_BITS;
+const SUBPACKETS_NUMBER_RANGE: Range<usize> = LENGTH_INDEX..SUBPACKET_START_INDEX_TLID1;
+
 impl<'bits, Store> TryFrom<&'bits BitSlice<Msb0, Store>> for Packet
 where
     Store: BitStore,
@@ -76,17 +95,20 @@ where
     type Error = anyhow::Error;
 
     fn try_from(bits: &'bits BitSlice<Msb0, Store>) -> Result<Self> {
-        let version: u8 = bits[0..3].load();
+        let version: u8 = bits[VERSION_RANGE].load_be();
-        let type_id: u8 = bits[3..6].load();
+        let type_id: u8 = bits[TYPE_ID_RANGE].load_be();
 
         match type_id {
+            // LitteralPacket
             4 => {
                 let mut value = 0;
+                let mut end = None;
                 for i in (6..).step_by(5) {
-                    let val = bits[(i + 1)..(i + 5)].load::<u64>();
+                    let val = bits[(i + 1)..(i + 5)].load_be::<u64>();
                     value = (value << 4) + val;
 
                     if !bits[i] {
+                        end = Some(i + 5);
                         break;
                     }
                 }
@@ -94,10 +116,45 @@ where
                 Ok(Packet {
                     version,
                     type_id,
+                    bits_used: end.unwrap(),
                     packet: PacketType::Litteral(LitteralPacket { value }),
                 })
             }
-            _ => unimplemented!("Operator packets aren't implemented yet"),
+            // OperatorPacket
+            _ => {
+                let length_type_id = bits[TYPE_LENGTH_ID_INDEX];
+                let mut sub_packets = Vec::new();
+                let len = if length_type_id {
+                    let subpacket_num: usize = bits[SUBPACKETS_NUMBER_RANGE].load_be();
+                    let mut start_index = SUBPACKET_START_INDEX_TLID1;
+
+                    while sub_packets.len() < subpacket_num {
+                        let packet: Packet = bits[start_index..].try_into()?;
+                        start_index += packet.bits_used;
+                        sub_packets.push(packet);
+                    }
+
+                    start_index
+                } else {
+                    let length: usize = bits[SUBPACKETS_BIT_LENGTH_RANGE].load_be();
+
+                    let mut start_index = SUBPACKET_START_INDEX_TLID0;
+                    while start_index != SUBPACKET_START_INDEX_TLID0 + length {
+                        let packet: Packet = bits[start_index..].try_into()?;
+                        start_index += packet.bits_used;
+                        sub_packets.push(packet);
+                    }
+
+                    SUBPACKET_START_INDEX_TLID0 + length
+                };
+
+                Ok(Packet {
+                    version,
+                    type_id,
+                    bits_used: len,
+                    packet: PacketType::Operator(OperatorPacket { sub_packets }),
+                })
+            }
         }
     }
 }
@@ -122,8 +179,75 @@ mod tests {
             Packet {
                 version: 6,
                 type_id: 4,
+                bits_used: 21,
                 packet: PacketType::Litteral(LitteralPacket { value: 2021 }),
             }
         );
+
+        assert_eq!(
+            PROVIDED2.parse::<Packet>().unwrap(),
+            Packet {
+                version: 1,
+                type_id: 6,
+                bits_used: 49,
+                packet: PacketType::Operator(OperatorPacket {
+                    sub_packets: vec![
+                        Packet {
+                            version: 6,
+                            type_id: 4,
+                            bits_used: 11,
+                            packet: PacketType::Litteral(LitteralPacket { value: 10 })
+                        },
+                        Packet {
+                            version: 2,
+                            type_id: 4,
+                            bits_used: 16,
+                            packet: PacketType::Litteral(LitteralPacket { value: 20 })
+                        },
+                    ],
+                }),
+            }
+        );
+
+        assert_eq!(
+            PROVIDED3.parse::<Packet>().unwrap(),
+            Packet {
+                version: 7,
+                type_id: 3,
+                bits_used: 51,
+                packet: PacketType::Operator(OperatorPacket {
+                    sub_packets: vec![
+                        Packet {
+                            version: 2,
+                            type_id: 4,
+                            bits_used: 11,
+                            packet: PacketType::Litteral(LitteralPacket { value: 1 })
+                        },
+                        Packet {
+                            version: 4,
+                            type_id: 4,
+                            bits_used: 11,
+                            packet: PacketType::Litteral(LitteralPacket { value: 2 })
+                        },
+                        Packet {
+                            version: 1,
+                            type_id: 4,
+                            bits_used: 11,
+                            packet: PacketType::Litteral(LitteralPacket { value: 3 })
+                        },
+                    ],
+                }),
+            }
+        );
+
+        assert_eq!(PROVIDED4.parse::<Packet>().unwrap().version_sum(), 16);
+        assert_eq!(PROVIDED5.parse::<Packet>().unwrap().version_sum(), 12);
+        assert_eq!(PROVIDED6.parse::<Packet>().unwrap().version_sum(), 23);
+        assert_eq!(PROVIDED7.parse::<Packet>().unwrap().version_sum(), 31);
+    }
+
+    #[test]
+    fn part1_real() {
+        assert_eq!(part1(INPUT).unwrap(), 925);
     }
 }