Skip to main content

statsig_rust/evaluation/
rkyv_value.rs

1use std::collections::HashMap;
2
3use rkyv::Archive;
4use serde::ser::{SerializeMap, SerializeSeq};
5
6// A bridging layer between Serde and Rkyv.
7// Based on Rkyv Examples: https://github.com/rkyv/rkyv/blob/main/rkyv/examples/json_like_schema.rs
8#[derive(
9    Archive,
10    Debug,
11    rkyv::Deserialize,
12    rkyv::Serialize,
13    Clone,
14    serde::Serialize,
15    serde::Deserialize,
16    PartialEq,
17)]
18#[rkyv(serialize_bounds(
19    __S: rkyv::ser::Writer + rkyv::ser::Allocator,
20    __S::Error: rkyv::rancor::Source,
21))]
22#[rkyv(deserialize_bounds(__D::Error: rkyv::rancor::Source))]
23#[rkyv(bytecheck(
24    bounds(
25        __C: rkyv::validation::ArchiveContext,
26    )
27))]
28#[rkyv(derive(Debug, PartialEq))]
29#[serde(untagged)]
30pub enum RkyvValue {
31    Null,
32    Bool(bool),
33    Number(RkyvNumber),
34    String(String),
35    Array(#[rkyv(omit_bounds)] Vec<RkyvValue>),
36    Object(#[rkyv(omit_bounds)] HashMap<String, RkyvValue>),
37}
38
39impl serde::Serialize for ArchivedRkyvValue {
40    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
41    where
42        S: serde::Serializer,
43    {
44        match self {
45            ArchivedRkyvValue::Null => serializer.serialize_none(),
46            ArchivedRkyvValue::Bool(b) => serializer.serialize_bool(*b),
47            ArchivedRkyvValue::Number(n) => n.serialize(serializer),
48            ArchivedRkyvValue::String(s) => serializer.serialize_str(s),
49            ArchivedRkyvValue::Array(a) => {
50                let mut seq = serializer.serialize_seq(Some(a.len()))?;
51                for element in a.iter() {
52                    seq.serialize_element(&element)?;
53                }
54                seq.end()
55            }
56            ArchivedRkyvValue::Object(o) => {
57                let mut map = serializer.serialize_map(Some(o.len()))?;
58
59                for (k, v) in o.iter() {
60                    map.serialize_entry(k.as_str(), v)?;
61                }
62
63                map.end()
64            }
65        }
66    }
67}
68
69impl PartialEq<ArchivedRkyvValue> for RkyvValue {
70    fn eq(&self, other: &ArchivedRkyvValue) -> bool {
71        match (self, other) {
72            (RkyvValue::Null, ArchivedRkyvValue::Null) => true,
73            (RkyvValue::Bool(a), ArchivedRkyvValue::Bool(b)) => a == b,
74            (RkyvValue::Number(a), ArchivedRkyvValue::Number(b)) => a == b,
75            (RkyvValue::String(a), ArchivedRkyvValue::String(b)) => a == b.as_str(),
76            (RkyvValue::Array(a), ArchivedRkyvValue::Array(b)) => {
77                a.len() == b.len() && a.iter().zip(b.iter()).all(|(a, b)| a == b)
78            }
79            (RkyvValue::Object(a), ArchivedRkyvValue::Object(b)) => {
80                a.len() == b.len()
81                    && b.iter()
82                        .all(|(key, value)| a.get(key.as_str()).is_some_and(|v| v == value))
83            }
84            _ => false,
85        }
86    }
87}
88
89impl PartialEq<RkyvValue> for ArchivedRkyvValue {
90    fn eq(&self, other: &RkyvValue) -> bool {
91        other == self
92    }
93}
94
95// ------------------------------------------------------------------------------- [ RkyvNumber ]
96
97#[derive(
98    Archive,
99    Debug,
100    rkyv::Deserialize,
101    rkyv::Serialize,
102    Clone,
103    serde::Serialize,
104    serde::Deserialize,
105    PartialEq,
106)]
107#[rkyv(derive(Debug, PartialEq))]
108#[serde(untagged)]
109pub enum RkyvNumber {
110    PosInt(u64),
111    NegInt(i64),
112    Float(f64),
113}
114
115impl serde::Serialize for ArchivedRkyvNumber {
116    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
117    where
118        S: serde::Serializer,
119    {
120        match self {
121            ArchivedRkyvNumber::PosInt(n) => serializer.serialize_u64(n.to_native()),
122            ArchivedRkyvNumber::NegInt(n) => serializer.serialize_i64(n.to_native()),
123            ArchivedRkyvNumber::Float(n) => serializer.serialize_f64(n.to_native()),
124        }
125    }
126}
127
128impl PartialEq<ArchivedRkyvNumber> for RkyvNumber {
129    fn eq(&self, other: &ArchivedRkyvNumber) -> bool {
130        match (self, other) {
131            (RkyvNumber::PosInt(a), ArchivedRkyvNumber::PosInt(b)) => a == b,
132            (RkyvNumber::NegInt(a), ArchivedRkyvNumber::NegInt(b)) => a == b,
133            (RkyvNumber::Float(a), ArchivedRkyvNumber::Float(b)) => a == b,
134            _ => false,
135        }
136    }
137}
138
139impl PartialEq<RkyvNumber> for ArchivedRkyvNumber {
140    fn eq(&self, other: &RkyvNumber) -> bool {
141        other == self
142    }
143}