deserialize_map/seq: support partial deserialization

[utkarshgupta137]

I'm using serde_json for extracting specific fields from nested JSONs (like a json pointer) & I'm trying to optimize my implementation.

Since the field I want is know ahead of time, I'm using the DeserializeSeed/Visitor API to drive the deserializer. The issue is that once I have found the field that I want, I still have to finish iterating over the rest of the JSON (basically skipping it all with de::IgnoredAny) before I can return this value since deserialize_map/seq call end_map/seq to ensure that the user has finished iterating the map/seq.
I've tried benchmarking my code with/without call to end_map/seq & I've found cases where the field extraction is 95% faster when the field is at the start of a large JSON object.

Code

struct BooleanVisitor;

impl<'de> de::Visitor<'de> for BooleanVisitor {
    type Value = Option<LhsValue<'de>>;

    fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(formatter, "boolean")
    }

    #[inline]
    fn visit_bool<E>(self, v: bool) -> Result<Self::Value, E>
    where
        E: de::Error,
    {
        Ok(Some(LhsValue::Bool(v)))
    }
}

struct IntegerVisitor;

impl<'de> de::Visitor<'de> for IntegerVisitor {
    type Value = Option<LhsValue<'de>>;

    fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(formatter, "integer")
    }

    #[inline]
    fn visit_i64<E>(self, v: i64) -> Result<Self::Value, E>
    where
        E: de::Error,
    {
        Ok(Some(LhsValue::Int(v)))
    }

    #[inline]
    fn visit_u64<E>(self, v: u64) -> Result<Self::Value, E>
    where
        E: de::Error,
    {
        Ok(i64::try_from(v).ok().map(LhsValue::Int))
    }
}

struct StringVisitor;

impl<'de> de::Visitor<'de> for StringVisitor {
    type Value = Option<LhsValue<'de>>;

    fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(formatter, "string")
    }

    #[inline]
    fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
    where
        E: de::Error,
    {
        Ok(Some(LhsValue::Bytes(v.to_owned().into())))
    }

    #[inline]
    fn visit_borrowed_str<E>(self, v: &'de str) -> Result<Self::Value, E>
    where
        E: de::Error,
    {
        Ok(Some(LhsValue::Bytes(v.into())))
    }

    #[inline]
    fn visit_string<E>(self, v: String) -> Result<Self::Value, E>
    where
        E: de::Error,
    {
        Ok(Some(LhsValue::Bytes(v.into())))
    }
}

struct KeyMatcher<'a>(&'a str);

impl<'de, 'a> de::DeserializeSeed<'de> for KeyMatcher<'a> {
    type Value = bool;

    #[inline]
    fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
    where
        D: de::Deserializer<'de>,
    {
        deserializer.deserialize_str(self)
    }
}

impl<'de, 'a> de::Visitor<'de> for KeyMatcher<'a> {
    type Value = bool;

    fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(formatter, "str")
    }

    #[inline]
    fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
    where
        E: de::Error,
    {
        Ok(v == self.0)
    }
}

struct MapVisitor<'a>(&'a str, JsonVisitor<'a>);

impl<'a, 'de> de::Visitor<'de> for MapVisitor<'a> {
    type Value = Option<LhsValue<'de>>;

    fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(formatter, "map")
    }

    #[inline]
    fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
    where
        A: de::MapAccess<'de>,
    {
        let mut ret = None;
        while let Some(matched) = map.next_key_seed(KeyMatcher(self.0))? {
            if matched {
                ret = map.next_value_seed(self.1)?;
                break;
            } else {
                map.next_value::<de::IgnoredAny>()?;
            }
        }
        // PERF: this is unavoidable
        while map
            .next_entry::<de::IgnoredAny, de::IgnoredAny>()?
            .is_some()
        {}
        Ok(ret)
    }
}

struct SeqVisitor<'a>(usize, JsonVisitor<'a>);

impl<'a, 'de> de::Visitor<'de> for SeqVisitor<'a> {
    type Value = Option<LhsValue<'de>>;

    fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(formatter, "seq")
    }

    #[inline]
    fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
    where
        A: de::SeqAccess<'de>,
    {
        for _ in 0..self.0 {
            if seq.next_element::<de::IgnoredAny>()?.is_none() {
                return Ok(None);
            }
        }
        let Some(ret) = seq.next_element_seed(self.1)? else {
            return Ok(None);
        };
        // PERF: this is unavoidable
        while seq.next_element::<de::IgnoredAny>()?.is_some() {}
        Ok(ret)
    }
}

struct JsonVisitor<'a>(&'a [JsonComponent], JsonScalarType);

impl<'a, 'de> de::DeserializeSeed<'de> for JsonVisitor<'a> {
    type Value = Option<LhsValue<'de>>;

    #[inline]
    fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
    where
        D: de::Deserializer<'de>,
    {
        if let Some((comp, rest)) = self.0.split_first() {
            let visitor = JsonVisitor(rest, self.1);
            match comp {
                JsonComponent::ObjectProperty(key) => {
                    deserializer.deserialize_map(MapVisitor(key, visitor))
                }
                JsonComponent::ArrayIndex(index) => {
                    deserializer.deserialize_seq(SeqVisitor(*index, visitor))
                }
            }
        } else {
            match self.1 {
                JsonScalarType::Boolean => deserializer.deserialize_bool(BooleanVisitor),
                JsonScalarType::Integer => deserializer.deserialize_i64(IntegerVisitor)
                JsonScalarType::String => deserializer.deserialize_str(StringVisitor),
            }
        }
    }
}

#[derive(Clone, Debug)]
enum JsonComponent {
    ObjectProperty(Box<str>),
    ArrayIndex(usize),
}

#[derive(Clone, Copy, Debug)]
pub enum JsonScalarType {
    Boolean,
    String,
    Integer,
}

I was wondering if there was a way I can avoid this. I've come up with a few solutions to this issue:

• extend serde::MapAccess/SeqAccess trait to add skip_rest to allow the Visitor to the skip rest of the input. This might be invasive to the whole serde ecosystem (the default implementation could just return an error or even panic), but it will provide a generic solution.
• add serde_json::Deserializer::enable_partial to allow opt-in access to this functionality for serde_json only. This could even sit behind a feature flag like unbounded_depth does.
• add serde_json::get functionality similar to sonic-rs::get which either always does partial deserialization or optionally allows partial deserialization via a bool.

I'm happy to iterate on these ideas & implement them as well. Personally, the 2nd option seems like the one with the least changes, but the 3rd option would probably be more useful in general.