[Rust] MessageRegistry::from_descriptors for sibling/imported types (PR #461 feedback)

rust/NEXT_CHANGELOG.md

@@ -6,6 +6,8 @@
 
 ### New Features and Improvements
 
+- Added `MessageRegistry::from_descriptors` (behind the `zeroparser` feature) for building a registry from a parse root plus sibling/imported top-level messages, each keyed by its package-qualified FQN. This lets the zeroparser decode messages whose root fields reference imported well-known types (e.g. `google.protobuf.Timestamp`, `Duration`, and the wrapper types) that `from_descriptor` alone does not register. `others` takes borrowed `&[(&str, &DescriptorProto)]` pairs.
+
 ### Bug Fixes
 
 ### Documentation

rust/sdk/src/zeroparser/registry.rs

@@ -185,15 +185,58 @@ impl MessageRegistry {
         }
     }
 
+    /// Build a registry from the parse root plus additional top-level messages.
+    ///
+    /// Use this when root fields reference sibling/imported messages that are
+    /// not nested under `root`.
+    #[inline(always)]
+    pub fn from_descriptors(root: &DescriptorProto, others: &[(&str, &DescriptorProto)]) -> Self {
+        let mut messages = HashMap::new();
+
+        // Register the root (and its nested types) at the root prefix, exactly
+        // as `from_descriptor` does, so root-relative lookups are unchanged.
+        Self::collect_messages(root, &mut messages);
+
+        // Register every other top-level message under its package-qualified FQN.
+        for (package, desc) in others {
+            let prefix = if package.is_empty() {
+                ROOT_PREFIX.to_string()
+            } else {
+                format!(".{package}")
+            };
+            Self::collect_messages_with_prefix(desc, prefix, &mut messages);
+        }
+
+        let root_descriptor = DescriptorWithFieldCache::from_descriptor(root);
+
+        MessageRegistry {
+            messages,
+            root_descriptor,
+        }
+    }
+
     /// Iteratively collect all message descriptors using an explicit stack.
     /// This avoids stack overflow for deeply nested message types.
     #[inline(always)]
     fn collect_messages(
         root: &DescriptorProto,
         acc: &mut HashMap<String, DescriptorWithFieldCache>,
+    ) {
+        Self::collect_messages_with_prefix(root, ROOT_PREFIX.to_string(), acc);
+    }
+
+    /// Collect `root` and all of its nested types into `acc`, threading
+    /// `initial_prefix` as the enclosing namespace. An empty prefix yields
+    /// `.Name` (root behaviour); a package prefix like `.google.protobuf`
+    /// yields `.google.protobuf.Name`.
+    #[inline(always)]
+    fn collect_messages_with_prefix(
+        root: &DescriptorProto,
+        initial_prefix: String,
+        acc: &mut HashMap<String, DescriptorWithFieldCache>,
     ) {
         // Use an explicit stack: (descriptor, prefix)
-        let mut stack: Vec<(&DescriptorProto, String)> = vec![(root, ROOT_PREFIX.to_string())];
+        let mut stack: Vec<(&DescriptorProto, String)> = vec![(root, initial_prefix)];
 
         while let Some((desc, current_prefix)) = stack.pop() {
             let name = desc.name.as_deref().unwrap_or("");
@@ -355,4 +398,87 @@ pub mod tests {
         assert_eq!(registry.get_field_type_name(2), Some(".Level1.Level2"));
         assert_eq!(registry.get_field_type_name(99), None);
     }
+
+    #[test]
+    fn message_registry_from_descriptors_registers_siblings() {
+        // Root message references an imported well-known type by its FQN.
+        let root = make_descriptor(
+            "Bet",
+            vec![
+                make_field(1, "id", field_descriptor_proto::Type::String, false, None),
+                make_field(
+                    2,
+                    "note",
+                    field_descriptor_proto::Type::Message,
+                    false,
+                    Some(".google.protobuf.StringValue"),
+                ),
+            ],
+        );
+
+        // A sibling/imported top-level message from another file + package.
+        let string_value = make_descriptor(
+            "StringValue",
+            vec![make_field(
+                1,
+                "value",
+                field_descriptor_proto::Type::String,
+                false,
+                None,
+            )],
+        );
+
+        // `from_descriptor` alone never registers the sibling: lookup misses.
+        let root_only = MessageRegistry::from_descriptor(&root);
+        assert!(root_only.get(".google.protobuf.StringValue").is_none());
+
+        // `from_descriptors` registers it under its true package-qualified FQN.
+        let registry =
+            MessageRegistry::from_descriptors(&root, &[("google.protobuf", &string_value)]);
+
+        assert!(registry.get(".Bet").is_some());
+        let wrapper = registry
+            .get(".google.protobuf.StringValue")
+            .expect("sibling well-known type must be registered");
+        assert_eq!(&*wrapper.get_field(1).unwrap().name, "value");
+
+        // Root remains the parse root, and the field still resolves to the FQN.
+        assert_eq!(
+            registry.get_field_type_name(2),
+            Some(".google.protobuf.StringValue")
+        );
+    }
+
+    #[test]
+    fn message_registry_from_descriptors_empty_package_and_nesting() {
+        // Empty-package sibling registers at the root prefix (`.Name`), and its
+        // nested types are collected too.
+        let mut sibling = make_descriptor(
+            "Sibling",
+            vec![make_field(
+                1,
+                "inner",
+                field_descriptor_proto::Type::Message,
+                false,
+                Some(".Sibling.Inner"),
+            )],
+        );
+        sibling.nested_type.push(make_descriptor(
+            "Inner",
+            vec![make_field(
+                1,
+                "x",
+                field_descriptor_proto::Type::Int32,
+                false,
+                None,
+            )],
+        ));
+
+        let root = make_descriptor("Root", vec![]);
+        let registry = MessageRegistry::from_descriptors(&root, &[("", &sibling)]);
+
+        assert!(registry.get(".Root").is_some());
+        assert!(registry.get(".Sibling").is_some());
+        assert!(registry.get(".Sibling.Inner").is_some());
+    }
 }

rust/sdk/src/zeroparser/tests/common/mod.rs

@@ -60,14 +60,25 @@ pub fn get_message_descriptor(
 }
 
 pub fn create_registry_for_version(version: ProtoVersion, message_name: &str) -> MessageRegistry {
-    let file_desc_set = load_descriptor_set(version);
+    let file_desc_set =
+        FileDescriptorSet::decode(E2E_DESCRIPTOR_SET).expect("decode descriptor file");
     let (msg_desc, file) = find_message_and_file(&file_desc_set, version.package(), message_name);
     let mut descriptor = msg_desc.clone();
     descriptor.name = Some(format!(
         "{}.{message_name}",
         file.package.as_deref().unwrap_or("")
     ));
-    MessageRegistry::from_descriptor(&descriptor)
+
+    let others: Vec<(&str, &DescriptorProto)> = file_desc_set
+        .file
+        .iter()
+        .flat_map(|f| {
+            let package = f.package.as_deref().unwrap_or("");
+            f.message_type.iter().map(move |msg| (package, msg))
+        })
+        .collect();
+
+    MessageRegistry::from_descriptors(&descriptor, &others)
 }
 
 pub fn encode_message_for_version<M: Message>(

rust/sdk/src/zeroparser/tests/e2e.rs

@@ -2925,13 +2925,10 @@ fn test_invalid_utf8_various_errors(#[case] version: ProtoVersion) {
 #[rstest]
 #[case(ProtoVersion::Proto2)]
 #[case(ProtoVersion::Proto3)]
-fn test_google_protobuf_types_not_in_registry(#[case] version: ProtoVersion) {
+fn test_google_protobuf_types_decode_success(#[case] version: ProtoVersion) {
     use prost_types::{Duration, Timestamp};
 
-    let registry = create_registry_for_version(version, "GoogleProtobufTypesMessage");
-
-    // Test google.protobuf types (Timestamp/Duration)
-    let (buf_timestamp, _) = match version {
+    let (buf, registry) = match version {
         ProtoVersion::Proto2 => {
             let msg = proto2::GoogleProtobufTypesMessage {
                 timestamp: Some(Timestamp {
@@ -2942,6 +2939,15 @@ fn test_google_protobuf_types_not_in_registry(#[case] version: ProtoVersion) {
                     seconds: 3600,
                     nanos: 500000000,
                 }),
+                string_value: Some("wrapped_string".to_string()),
+                int32_value: Some(-42),
+                int64_value: Some(-4242),
+                uint32_value: Some(42),
+                uint64_value: Some(4242),
+                float_value: Some(3.5),
+                double_value: Some(7.25),
+                bool_value: Some(true),
+                bytes_value: Some(b"wrapped_bytes".to_vec()),
                 timestamps: vec![
                     Timestamp {
                         seconds: 1000000000,
@@ -2952,7 +2958,6 @@ fn test_google_protobuf_types_not_in_registry(#[case] version: ProtoVersion) {
                         nanos: 999999999,
                     },
                 ],
-                ..Default::default()
             };
             encode_message_for_version(version, &msg, "GoogleProtobufTypesMessage")
         }
@@ -2966,6 +2971,15 @@ fn test_google_protobuf_types_not_in_registry(#[case] version: ProtoVersion) {
                     seconds: 3600,
                     nanos: 500000000,
                 }),
+                string_value: Some("wrapped_string".to_string()),
+                int32_value: Some(-42),
+                int64_value: Some(-4242),
+                uint32_value: Some(42),
+                uint64_value: Some(4242),
+                float_value: Some(3.5),
+                double_value: Some(7.25),
+                bool_value: Some(true),
+                bytes_value: Some(b"wrapped_bytes".to_vec()),
                 timestamps: vec![
                     Timestamp {
                         seconds: 1000000000,
@@ -2976,37 +2990,81 @@ fn test_google_protobuf_types_not_in_registry(#[case] version: ProtoVersion) {
                         nanos: 999999999,
                     },
                 ],
-                ..Default::default()
             };
             encode_message_for_version(version, &msg, "GoogleProtobufTypesMessage")
         }
     };
 
-    let result = ParsedMessage::parse(&buf_timestamp, &registry);
-    assert!(
-        matches!(result, Err(ParseError::UnknownTypeName { .. })),
-        "Expected UnknownTypeName error for google.protobuf.Timestamp, got {:?}",
-        result
-    );
+    let parsed = ParsedMessage::parse(&buf, &registry).unwrap();
 
-    // Test wrapper types (manually encoded)
-    let mut buf_wrappers = Vec::new();
+    let timestamp = parsed.get_message(1).unwrap();
+    assert_eq!(
+        timestamp.get_scalar(1),
+        Some(&FieldValueRef::Int64(1234567890))
+    );
+    assert_eq!(
+        timestamp.get_scalar(2),
+        Some(&FieldValueRef::Int32(123456789))
+    );
 
-    // StringValue (field 3) - message with field 1 = string
-    let mut string_value_buf = Vec::new();
-    prost::encoding::string::encode(1, &"wrapped_string".to_string(), &mut string_value_buf);
-    prost::encoding::message::encode(3, &string_value_buf, &mut buf_wrappers);
+    let duration = parsed.get_message(2).unwrap();
+    assert_eq!(duration.get_scalar(1), Some(&FieldValueRef::Int64(3600)));
+    assert_eq!(
+        duration.get_scalar(2),
+        Some(&FieldValueRef::Int32(500000000))
+    );
 
-    // Int32Value (field 4) - message with field 1 = int32
-    let mut int32_value_buf = Vec::new();
-    prost::encoding::int32::encode(1, &-42, &mut int32_value_buf);
-    prost::encoding::message::encode(4, &int32_value_buf, &mut buf_wrappers);
+    assert_eq!(
+        parsed.get_message(3).unwrap().get_scalar(1),
+        Some(&FieldValueRef::String("wrapped_string"))
+    );
+    assert_eq!(
+        parsed.get_message(4).unwrap().get_scalar(1),
+        Some(&FieldValueRef::Int32(-42))
+    );
+    assert_eq!(
+        parsed.get_message(5).unwrap().get_scalar(1),
+        Some(&FieldValueRef::Int64(-4242))
+    );
+    assert_eq!(
+        parsed.get_message(6).unwrap().get_scalar(1),
+        Some(&FieldValueRef::UInt32(42))
+    );
+    assert_eq!(
+        parsed.get_message(7).unwrap().get_scalar(1),
+        Some(&FieldValueRef::UInt64(4242))
+    );
+    assert_eq!(
+        parsed.get_message(8).unwrap().get_scalar(1),
+        Some(&FieldValueRef::Float(3.5))
+    );
+    assert_eq!(
+        parsed.get_message(9).unwrap().get_scalar(1),
+        Some(&FieldValueRef::Double(7.25))
+    );
+    assert_eq!(
+        parsed.get_message(10).unwrap().get_scalar(1),
+        Some(&FieldValueRef::Bool(true))
+    );
+    assert_eq!(
+        parsed.get_message(11).unwrap().get_scalar(1),
+        Some(&FieldValueRef::Bytes(b"wrapped_bytes"))
+    );
 
-    let result = ParsedMessage::parse(&buf_wrappers, &registry);
-    assert!(
-        matches!(result, Err(ParseError::UnknownTypeName { .. })),
-        "Expected UnknownTypeName error for google.protobuf wrapper types, got {:?}",
-        result
+    let timestamps = parsed.get_repeated_messages(12);
+    assert_eq!(timestamps.len(), 2);
+    assert_eq!(
+        timestamps[0].get_scalar(1),
+        Some(&FieldValueRef::Int64(1000000000))
+    );
+    assert_eq!(timestamps[0].get_scalar(2), None);
+    assert_eq!(
+        timestamps[1].get_scalar(1),
+        Some(&FieldValueRef::Int64(2000000000))
+    );
+    assert_eq!(
+        timestamps[1].get_scalar(2),
+        Some(&FieldValueRef::Int32(999999999))
     );
 }