substrait_validator/input/
traits.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Module providing introspection traits for [`prost`]-generated protobuf
4//! types.
5
6use crate::output::primitive_data;
7use crate::output::tree;
8use crate::parse::context;
9
10/// Trait for all Rust types that represent input tree node types.
11pub trait InputNode {
12    /// Creates an empty output node for a protobuf datum of this type.
13    ///
14    /// For primitive types, this fills the value with protobuf's default.
15    fn type_to_node() -> tree::Node;
16
17    /// Creates an empty output node for a protobuf datum with this value.
18    fn data_to_node(&self) -> tree::Node;
19
20    /// Returns the name of the selected variant of a oneof field, if this
21    /// is a rust enum used to represent a oneof field.
22    fn oneof_variant(&self) -> Option<&'static str>;
23
24    /// Complete the subtrees of this datum in output that have not already
25    /// been parsed using UnknownField nodes. Returns whether any such nodes
26    /// were added.
27    fn parse_unknown(&self, context: &mut context::Context<'_>) -> bool;
28}
29
30/// Trait for all Rust types that represent protobuf messages. These are
31/// always structs for which all fields implement InputNode.
32pub trait ProtoMessage: InputNode {
33    /// Returns the protobuf type name for messages of this type.
34    fn proto_message_type() -> &'static str;
35}
36
37/// Trait for all Rust types that represent protobuf's oneof abstraction.
38/// In the world of protobuf, these aren't really a thing of their own, but
39/// in Rust, they are defined as enums, each variant containing a one-tuple
40/// of some type implementing InputNode.
41pub trait ProtoOneOf: InputNode {
42    /// Returns the name of the selected variant of a oneof field.
43    fn proto_oneof_variant(&self) -> &'static str;
44}
45
46/// Trait for Rust types that map to the protobuf primitive types.
47pub trait ProtoPrimitive: InputNode {
48    /// Returns the protobuf type name for primitives of this type.
49    fn proto_primitive_type() -> &'static str;
50
51    /// Returns the protobuf-specified default value for this primitive
52    /// data type.
53    fn proto_primitive_default() -> primitive_data::PrimitiveData;
54
55    /// Returns the actual value for this primitive data type as a
56    /// ProtoPrimitiveData variant.
57    fn proto_primitive_data(&self) -> primitive_data::PrimitiveData;
58
59    /// Returns whether this is the default value of the primitive.
60    fn proto_primitive_is_default(&self) -> bool;
61}
62
63/// Trait for all Rust types that represent protobuf enums. These are
64/// always represented as a Rust enum with no contained values for any of
65/// the variants.
66pub trait ProtoEnum: ProtoPrimitive {
67    /// Returns the protobuf type name for enums of this type.
68    fn proto_enum_type() -> &'static str;
69
70    /// Returns the name of the default variant of an enum.
71    fn proto_enum_default_variant() -> &'static str;
72
73    /// Returns the name of the selected variant of an enum.
74    fn proto_enum_variant(&self) -> &'static str;
75
76    /// Returns the enumeration entry corresponding to the given integer
77    /// value, if any.
78    fn proto_enum_from_i32(x: i32) -> Option<Self>
79    where
80        Self: Sized;
81}
82
83/// Blanket implementation to make all protobuf enums behave like
84/// primitives as well.
85impl<T: ProtoEnum> ProtoPrimitive for T {
86    fn proto_primitive_type() -> &'static str {
87        T::proto_enum_type()
88    }
89
90    fn proto_primitive_default() -> primitive_data::PrimitiveData {
91        primitive_data::PrimitiveData::Enum(T::proto_enum_default_variant())
92    }
93
94    fn proto_primitive_data(&self) -> primitive_data::PrimitiveData {
95        primitive_data::PrimitiveData::Enum(self.proto_enum_variant())
96    }
97
98    fn proto_primitive_is_default(&self) -> bool {
99        self.proto_enum_variant() == T::proto_enum_default_variant()
100    }
101}
102
103/// Blanket implementation to make all protobuf primitives behave like
104/// generic protobuf datums.
105///
106/// Note: if Rust would allow it, we could define blanket implementations
107/// for ProtoMessage and ProtoOneOf as well, since they're always the same.
108/// Unfortunately, we can only define a single blanket implementation, so
109/// we opt for the one that isn't already generated via derive macros.
110impl<T: ProtoPrimitive> InputNode for T {
111    fn type_to_node() -> tree::Node {
112        tree::NodeType::ProtoPrimitive(T::proto_primitive_type(), T::proto_primitive_default())
113            .into()
114    }
115
116    fn data_to_node(&self) -> tree::Node {
117        tree::NodeType::ProtoPrimitive(T::proto_primitive_type(), self.proto_primitive_data())
118            .into()
119    }
120
121    fn oneof_variant(&self) -> Option<&'static str> {
122        None
123    }
124
125    fn parse_unknown(&self, _context: &mut context::Context<'_>) -> bool {
126        false
127    }
128}
129
130impl InputNode for () {
131    fn type_to_node() -> tree::Node {
132        tree::NodeType::ProtoMessage("google.protobuf.Empty").into()
133    }
134
135    fn data_to_node(&self) -> tree::Node {
136        tree::NodeType::ProtoMessage("google.protobuf.Empty").into()
137    }
138
139    fn oneof_variant(&self) -> Option<&'static str> {
140        None
141    }
142
143    fn parse_unknown(&self, _context: &mut context::Context<'_>) -> bool {
144        false
145    }
146}
147
148impl ProtoMessage for () {
149    fn proto_message_type() -> &'static str {
150        "google.protobuf.Empty"
151    }
152}
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157    use crate::input::proto::substrait;
158    use crate::output::primitive_data;
159    use crate::output::tree;
160
161    #[test]
162    fn message() {
163        assert_eq!(substrait::Plan::proto_message_type(), "substrait.Plan");
164        assert_eq!(
165            substrait::Plan::type_to_node(),
166            tree::Node {
167                class: tree::Class::Misc,
168                brief: None,
169                summary: None,
170                node_type: tree::NodeType::ProtoMessage("substrait.Plan"),
171                data_type: None,
172                data: vec![],
173            }
174        );
175
176        let msg = substrait::Plan::default();
177        assert_eq!(
178            msg.data_to_node(),
179            tree::Node {
180                class: tree::Class::Misc,
181                brief: None,
182                summary: None,
183                node_type: tree::NodeType::ProtoMessage("substrait.Plan"),
184                data_type: None,
185                data: vec![],
186            }
187        );
188        assert_eq!(msg.oneof_variant(), None);
189    }
190
191    #[test]
192    fn oneof() {
193        assert_eq!(
194            substrait::plan_rel::RelType::type_to_node(),
195            tree::Node {
196                class: tree::Class::Misc,
197                brief: None,
198                summary: None,
199                node_type: tree::NodeType::ProtoMissingOneOf,
200                data_type: None,
201                data: vec![],
202            }
203        );
204
205        let oneof = substrait::plan_rel::RelType::Rel(substrait::Rel::default());
206        assert_eq!(oneof.proto_oneof_variant(), "rel");
207        assert_eq!(
208            oneof.data_to_node(),
209            tree::Node {
210                class: tree::Class::Misc,
211                brief: None,
212                summary: None,
213                node_type: tree::NodeType::ProtoMessage("substrait.Rel"),
214                data_type: None,
215                data: vec![],
216            }
217        );
218        assert_eq!(oneof.oneof_variant(), Some("rel"));
219    }
220
221    #[test]
222    fn enumeration() {
223        assert_eq!(
224            substrait::AggregationPhase::proto_enum_type(),
225            "substrait.AggregationPhase"
226        );
227        assert_eq!(
228            substrait::AggregationPhase::proto_enum_default_variant(),
229            "AGGREGATION_PHASE_UNSPECIFIED"
230        );
231        assert_eq!(
232            substrait::AggregationPhase::Unspecified.proto_enum_variant(),
233            "AGGREGATION_PHASE_UNSPECIFIED"
234        );
235
236        assert_eq!(
237            substrait::AggregationPhase::proto_primitive_type(),
238            "substrait.AggregationPhase"
239        );
240        assert_eq!(
241            substrait::AggregationPhase::proto_primitive_default(),
242            primitive_data::PrimitiveData::Enum("AGGREGATION_PHASE_UNSPECIFIED")
243        );
244        assert_eq!(
245            substrait::AggregationPhase::Unspecified.proto_primitive_data(),
246            primitive_data::PrimitiveData::Enum("AGGREGATION_PHASE_UNSPECIFIED")
247        );
248
249        assert_eq!(
250            substrait::AggregationPhase::type_to_node(),
251            tree::Node {
252                class: tree::Class::Misc,
253                brief: None,
254                summary: None,
255                node_type: tree::NodeType::ProtoPrimitive(
256                    "substrait.AggregationPhase",
257                    primitive_data::PrimitiveData::Enum("AGGREGATION_PHASE_UNSPECIFIED")
258                ),
259                data_type: None,
260                data: vec![],
261            }
262        );
263        assert_eq!(
264            substrait::AggregationPhase::Unspecified.data_to_node(),
265            tree::Node {
266                class: tree::Class::Misc,
267                brief: None,
268                summary: None,
269                node_type: tree::NodeType::ProtoPrimitive(
270                    "substrait.AggregationPhase",
271                    primitive_data::PrimitiveData::Enum("AGGREGATION_PHASE_UNSPECIFIED")
272                ),
273                data_type: None,
274                data: vec![],
275            }
276        );
277        assert_eq!(
278            substrait::AggregationPhase::Unspecified.oneof_variant(),
279            None
280        );
281    }
282
283    #[test]
284    fn primitive() {
285        assert_eq!(u32::proto_primitive_type(), "uint32");
286        assert_eq!(
287            u32::proto_primitive_default(),
288            primitive_data::PrimitiveData::Unsigned(0)
289        );
290        assert_eq!(
291            42u32.proto_primitive_data(),
292            primitive_data::PrimitiveData::Unsigned(42)
293        );
294
295        assert_eq!(
296            u32::type_to_node(),
297            tree::Node {
298                class: tree::Class::Misc,
299                brief: None,
300                summary: None,
301                node_type: tree::NodeType::ProtoPrimitive(
302                    "uint32",
303                    primitive_data::PrimitiveData::Unsigned(0)
304                ),
305                data_type: None,
306                data: vec![],
307            }
308        );
309        assert_eq!(
310            42u32.data_to_node(),
311            tree::Node {
312                class: tree::Class::Misc,
313                brief: None,
314                summary: None,
315                node_type: tree::NodeType::ProtoPrimitive(
316                    "uint32",
317                    primitive_data::PrimitiveData::Unsigned(42)
318                ),
319                data_type: None,
320                data: vec![],
321            }
322        );
323        assert_eq!(42u32.oneof_variant(), None);
324    }
325}
substrait_validator/input/traits.rs

substrait_validator/input/
traits.rs
