Solution

  1. /// A wire type as seen on the wire.
  2. enum WireType {
  3.     /// The Varint WireType indicates the value is a single VARINT.
  4.     Varint,
  5.     /// The I64 WireType indicates that the value is precisely 8 bytes in
  6.     /// little-endian order containing a 64-bit signed integer or double type.
  7.     //I64,  -- not needed for this exercise
  8.     /// The Len WireType indicates that the value is a length represented as a
  9.     /// VARINT followed by exactly that number of bytes.
  10.     Len,
  11.     // The I32 WireType indicates that the value is precisely 4 bytes in
  12.     // little-endian order containing a 32-bit signed integer or float type.
  13.     //I32,  -- not needed for this exercise
  14. }
  16. #[derive(Debug)]
  17. /// A field's value, typed based on the wire type.
  18. enum FieldValue<'a> {
  19.     Varint(u64),
  20.     //I64(i64),  -- not needed for this exercise
  21.     Len(&'a [u8]),
  22.     //I32(i32),  -- not needed for this exercise
  23. }
  25. #[derive(Debug)]
  26. /// A field, containing the field number and its value.
  27. struct Field<'a> {
  28.     field_num: u64,
  29.     value: FieldValue<'a>,
  30. }
  32. trait ProtoMessage<'a>: Default {
  33.     fn add_field(&mut self, field: Field<'a>);
  34. }
  36. impl From<u64> for WireType {
  37.     fn from(value: u64) -> Self {
  38.         match value {
  39.             0 => WireType::Varint,
  40.             //1 => WireType::I64,  -- not needed for this exercise
  41.             2 => WireType::Len,
  42.             //5 => WireType::I32,  -- not needed for this exercise
  43.             _ => panic!("Invalid wire type: {value}"),
  44.         }
  45.     }
  46. }
  48. impl<'a> FieldValue<'a> {
  49.     fn as_str(&self) -> &'a str {
  50.         let FieldValue::Len(data) = self else {
  51.             panic!("Expected string to be a `Len` field");
  52.         };
  53.         std::str::from_utf8(data).expect("Invalid string")
  54.     }
  56.     fn as_bytes(&self) -> &'a [u8] {
  57.         let FieldValue::Len(data) = self else {
  58.             panic!("Expected bytes to be a `Len` field");
  59.         };
  60.         data
  61.     }
  63.     fn as_u64(&self) -> u64 {
  64.         let FieldValue::Varint(value) = self else {
  65.             panic!("Expected `u64` to be a `Varint` field");
  66.         };
  67.         *value
  68.     }
  69. }
  71. /// Parse a VARINT, returning the parsed value and the remaining bytes.
  72. fn parse_varint(data: &[u8]) -> (u64, &[u8]) {
  73.     for i in 0..7 {
  74.         let Some(b) = data.get(i) else {
  75.             panic!("Not enough bytes for varint");
  76.         };
  77.         if b & 0x80 == 0 {
  78.             // This is the last byte of the VARINT, so convert it to
  79.             // a u64 and return it.
  80.             let mut value = 0u64;
  81.             for b in data[..=i].iter().rev() {
  82.                 value = (value << 7) | (b & 0x7f) as u64;
  83.             }
  84.             return (value, &data[i + 1..]);
  85.         }
  86.     }
  88.     // More than 7 bytes is invalid.
  89.     panic!("Too many bytes for varint");
  90. }
  92. /// Convert a tag into a field number and a WireType.
  93. fn unpack_tag(tag: u64) -> (u64, WireType) {
  94.     let field_num = tag >> 3;
  95.     let wire_type = WireType::from(tag & 0x7);
  96.     (field_num, wire_type)
  97. }
  99. /// Parse a field, returning the remaining bytes
  100. fn parse_field(data: &[u8]) -> (Field, &[u8]) {
  101.     let (tag, remainder) = parse_varint(data);
  102.     let (field_num, wire_type) = unpack_tag(tag);
  103.     let (fieldvalue, remainder) = match wire_type {
  104.         WireType::Varint => {
  105.             let (value, remainder) = parse_varint(remainder);
  106.             (FieldValue::Varint(value), remainder)
  107.         }
  108.         WireType::Len => {
  109.             let (len, remainder) = parse_varint(remainder);
  110.             let len: usize = len.try_into().expect("len not a valid `usize`");
  111.             if remainder.len() < len {
  112.                 panic!("Unexpected EOF");
  113.             }
  114.             let (value, remainder) = remainder.split_at(len);
  115.             (FieldValue::Len(value), remainder)
  116.         }
  117.     };
  118.     (Field { field_num, value: fieldvalue }, remainder)
  119. }
  121. /// Parse a message in the given data, calling `T::add_field` for each field in
  122. /// the message.
  123. ///
  124. /// The entire input is consumed.
  125. fn parse_message<'a, T: ProtoMessage<'a>>(mut data: &'a [u8]) -> T {
  126.     let mut result = T::default();
  127.     while !data.is_empty() {
  128.         let parsed = parse_field(data);
  129.         result.add_field(parsed.0);
  130.         data = parsed.1;
  131.     }
  132.     result
  133. }
  135. #[derive(PartialEq)]
  136. #[derive(Debug, Default)]
  137. struct PhoneNumber<'a> {
  138.     number: &'a str,
  139.     type_: &'a str,
  140. }
  142. #[derive(PartialEq)]
  143. #[derive(Debug, Default)]
  144. struct Person<'a> {
  145.     name: &'a str,
  146.     id: u64,
  147.     phone: Vec<PhoneNumber<'a>>,
  148. }
  150. impl<'a> ProtoMessage<'a> for Person<'a> {
  151.     fn add_field(&mut self, field: Field<'a>) {
  152.         match field.field_num {
  153.             1 => self.name = field.value.as_str(),
  154.             2 => self.id = field.value.as_u64(),
  155.             3 => self.phone.push(parse_message(field.value.as_bytes())),
  156.             _ => {} // skip everything else
  157.         }
  158.     }
  159. }
  161. impl<'a> ProtoMessage<'a> for PhoneNumber<'a> {
  162.     fn add_field(&mut self, field: Field<'a>) {
  163.         match field.field_num {
  164.             1 => self.number = field.value.as_str(),
  165.             2 => self.type_ = field.value.as_str(),
  166.             _ => {} // skip everything else
  167.         }
  168.     }
  169. }
  171. fn main() {
  172.     let person: Person = parse_message(&[
  173.         0x0a, 0x07, 0x6d, 0x61, 0x78, 0x77, 0x65, 0x6c, 0x6c, 0x10, 0x2a, 0x1a,
  174.         0x16, 0x0a, 0x0e, 0x2b, 0x31, 0x32, 0x30, 0x32, 0x2d, 0x35, 0x35, 0x35,
  175.         0x2d, 0x31, 0x32, 0x31, 0x32, 0x12, 0x04, 0x68, 0x6f, 0x6d, 0x65, 0x1a,
  176.         0x18, 0x0a, 0x0e, 0x2b, 0x31, 0x38, 0x30, 0x30, 0x2d, 0x38, 0x36, 0x37,
  177.         0x2d, 0x35, 0x33, 0x30, 0x38, 0x12, 0x06, 0x6d, 0x6f, 0x62, 0x69, 0x6c,
  178.         0x65,
  179.     ]);
  180.     println!("{:#?}", person);
  181. }
  183. #[cfg(test)]
  184. mod tests {
  185.     use super::*;
  187.     #[test]
  188.     fn test_id() {
  189.         let person_id: Person = parse_message(&[0x10, 0x2a]);
  190.         assert_eq!(person_id, Person { name: "", id: 42, phone: vec![] });
  191.     }
  193.     #[test]
  194.     fn test_name() {
  195.         let person_name: Person = parse_message(&[
  196.             0x0a, 0x0e, 0x62, 0x65, 0x61, 0x75, 0x74, 0x69, 0x66, 0x75, 0x6c, 0x20,
  197.             0x6e, 0x61, 0x6d, 0x65,
  198.         ]);
  199.         assert_eq!(
  200.             person_name,
  201.             Person { name: "beautiful name", id: 0, phone: vec![] }
  202.         );
  203.     }
  205.     #[test]
  206.     fn test_just_person() {
  207.         let person_name_id: Person =
  208.             parse_message(&[0x0a, 0x04, 0x45, 0x76, 0x61, 0x6e, 0x10, 0x16]);
  209.         assert_eq!(person_name_id, Person { name: "Evan", id: 22, phone: vec![] });
  210.     }
  212.     #[test]
  213.     fn test_phone() {
  214.         let phone: Person = parse_message(&[
  215.             0x0a, 0x00, 0x10, 0x00, 0x1a, 0x16, 0x0a, 0x0e, 0x2b, 0x31, 0x32, 0x33,
  216.             0x34, 0x2d, 0x37, 0x37, 0x37, 0x2d, 0x39, 0x30, 0x39, 0x30, 0x12, 0x04,
  217.             0x68, 0x6f, 0x6d, 0x65,
  218.         ]);
  219.         assert_eq!(
  220.             phone,
  221.             Person {
  222.                 name: "",
  223.                 id: 0,
  224.                 phone: vec![PhoneNumber { number: "+1234-777-9090", type_: "home" },],
  225.             }
  226.         );
  227.     }
  228. }