1#![deny(rust_2018_idioms)]
2#![deny(clippy::all)]
3
4use educe::Educe;
5use indexmap::IndexSet;
6use itertools::Itertools;
7use miette::Diagnostic;
8use partiql_common::node::{AutoNodeIdGenerator, NodeId, NodeIdGenerator, NullIdGenerator};
9use std::fmt::{Debug, Display, Formatter};
10use std::hash::{Hash, Hasher};
11use std::sync::OnceLock;
12use thiserror::Error;
13
14#[derive(Debug, Clone, Eq, PartialEq, Hash, Error, Diagnostic)]
15#[error("ShapeResult Error")]
16#[non_exhaustive]
17pub enum ShapeResultError {
18 #[error("Unexpected type `{0:?}` for static type bool")]
19 UnexpectedType(String),
20}
21
22pub type ShapeResult<T> = Result<T, ShapeResultError>;
24
25pub trait Type {}
26
27impl Type for StaticType {}
28
29fn indexset_hash<H, T>(set: &IndexSet<T>, state: &mut H)
30where
31 H: Hasher,
32 T: Hash,
33{
34 for val in set {
35 val.hash(state)
36 }
37}
38
39#[macro_export]
40macro_rules! type_dynamic {
41 ($bld:expr) => {
42 $bld.new_dynamic()
43 };
44}
45
46#[macro_export]
47macro_rules! type_int {
48 ($bld:expr) => {
49 $bld.new_static($crate::Static::Int)
50 };
51}
52
53#[macro_export]
54macro_rules! type_int8 {
55 ($bld:expr) => {
56 $bld.new_static($crate::Static::Int8)
57 };
58}
59
60#[macro_export]
61macro_rules! type_int16 {
62 ($bld:expr) => {
63 $bld.new_static($crate::Static::Int16)
64 };
65}
66
67#[macro_export]
68macro_rules! type_int32 {
69 ($bld:expr) => {
70 $bld.new_static($crate::Static::Int32)
71 };
72}
73
74#[macro_export]
75macro_rules! type_int64 {
76 ($bld:expr) => {
77 $bld.new_static($crate::Static::Int64)
78 };
79}
80
81#[macro_export]
82macro_rules! type_decimal {
83 ($bld:expr) => {
84 $bld.new_static($crate::Static::Decimal)
85 };
86}
87
88#[macro_export]
91macro_rules! type_float32 {
92 ($bld:expr) => {
93 $bld.new_static($crate::Static::Float32)
94 };
95}
96
97#[macro_export]
98macro_rules! type_float64 {
99 ($bld:expr) => {
100 $bld.new_static($crate::Static::Float64)
101 };
102}
103
104#[macro_export]
105macro_rules! type_string {
106 ($bld:expr) => {
107 $bld.new_static($crate::Static::String)
108 };
109}
110
111#[macro_export]
112macro_rules! type_bool {
113 ($bld:expr) => {
114 $bld.new_static($crate::Static::Bool)
115 };
116}
117
118#[macro_export]
119macro_rules! type_numeric {
120 ($bld:expr) => {
121 [
122 $crate::type_int!($bld),
123 $crate::type_float32!($bld),
124 $crate::type_float64!($bld),
125 $crate::type_decimal!($bld),
126 ]
127 .into_any_of($bld)
128 };
129}
130
131#[macro_export]
132macro_rules! type_datetime {
133 ($bld:expr) => {
134 $bld.new_static($crate::Static::DateTime)
135 };
136}
137
138#[macro_export]
139macro_rules! type_struct {
140 ($bld:expr) => {
141 $bld.new_struct(StructType::new_any())
142 };
143 ($bld:expr, $elem:expr) => {{
144 let elem = $elem;
145 $bld.new_struct(StructType::new(elem))
146 }};
147}
148
149#[macro_export]
150macro_rules! type_graph {
151 ($bld:expr) => {
152 $bld.new_graph()
153 };
154}
155
156#[macro_export]
157macro_rules! struct_fields {
158 ($(($x:expr, $y:expr)),+ $(,)?) => (
159 $crate::StructConstraint::Fields([$(($x, $y).into()),+].into())
160 );
161}
162
163#[macro_export]
164macro_rules! type_bag {
165 ($bld:expr) => {
166 $bld.new_bag(BagType::new_any())
167 };
168 ($bld:expr, $elem:expr) => {{
169 let elem = $elem;
170 $bld.new_bag($crate::BagType::new(Box::new(elem)))
171 }};
172}
173
174#[macro_export]
175macro_rules! type_array {
176 ($bld:expr) => {
177 $bld.new_array(ArrayType::new_any())
178 };
179 ($bld:expr, $elem:expr) => {{
180 let elem = $elem;
181 $bld.new_array($crate::ArrayType::new(Box::new(elem)))
182 }};
183}
184
185#[derive(Debug, Clone, Eq, PartialEq, Hash)]
187#[non_exhaustive]
191pub enum PartiqlShape {
192 Dynamic,
193 AnyOf(AnyOf),
194 Static(StaticType),
195 Undefined,
196}
197
198#[allow(dead_code)]
199impl PartiqlShape {
200 #[inline]
201 pub fn static_type_id(&self) -> Option<NodeId> {
202 if let PartiqlShape::Static(StaticType { id, .. }) = self {
203 Some(*id)
204 } else {
205 None
206 }
207 }
208
209 #[inline]
210 pub fn is_string(&self) -> bool {
211 matches!(
212 &self,
213 PartiqlShape::Static(StaticType {
214 ty: Static::String,
215 nullable: true,
216 ..
217 })
218 )
219 }
220
221 #[inline]
222 pub fn is_struct(&self) -> bool {
223 matches!(
224 *self,
225 PartiqlShape::Static(StaticType {
226 ty: Static::Struct(_),
227 nullable: true,
228 ..
229 })
230 )
231 }
232
233 #[inline]
234 pub fn is_collection(&self) -> bool {
235 matches!(
236 *self,
237 PartiqlShape::Static(StaticType {
238 ty: Static::Bag(_),
239 nullable: true,
240 ..
241 })
242 ) || matches!(
243 *self,
244 PartiqlShape::Static(StaticType {
245 ty: Static::Array(_),
246 nullable: true,
247 ..
248 })
249 )
250 }
251
252 #[inline]
253 pub fn is_unordered_collection(&self) -> bool {
254 !self.is_ordered_collection()
255 }
256
257 #[inline]
258 pub fn is_ordered_collection(&self) -> bool {
259 matches!(
261 *self,
262 PartiqlShape::Static(StaticType {
263 ty: Static::Array(_),
264 nullable: true,
265 ..
266 })
267 )
268 }
269
270 #[inline]
271 pub fn is_graph(&self) -> bool {
272 matches!(
273 *self,
274 PartiqlShape::Static(StaticType {
275 ty: Static::Graph(),
276 nullable: true,
277 ..
278 })
279 )
280 }
281
282 #[inline]
283 pub fn is_bag(&self) -> bool {
284 matches!(
285 *self,
286 PartiqlShape::Static(StaticType {
287 ty: Static::Bag(_),
288 nullable: true,
289 ..
290 })
291 )
292 }
293
294 #[inline]
295 pub fn is_array(&self) -> bool {
296 matches!(
297 *self,
298 PartiqlShape::Static(StaticType {
299 ty: Static::Array(_),
300 nullable: true,
301 ..
302 })
303 )
304 }
305
306 #[inline]
307 pub fn is_dynamic(&self) -> bool {
308 matches!(*self, PartiqlShape::Dynamic)
309 }
310
311 #[inline]
312 pub fn is_undefined(&self) -> bool {
313 matches!(*self, PartiqlShape::Undefined)
314 }
315
316 pub fn expect_bool(&self) -> ShapeResult<StaticType> {
317 if let PartiqlShape::Static(StaticType {
318 id,
319 ty: Static::Bool,
320 nullable: n,
321 }) = self
322 {
323 Ok(StaticType {
324 id: *id,
325 ty: Static::Bool,
326 nullable: *n,
327 })
328 } else {
329 Err(ShapeResultError::UnexpectedType(format!("{self}")))
330 }
331 }
332
333 pub fn expect_bag(&self) -> ShapeResult<BagType> {
334 if let PartiqlShape::Static(StaticType {
335 ty: Static::Bag(bag),
336 ..
337 }) = self
338 {
339 Ok(bag.clone())
340 } else {
341 Err(ShapeResultError::UnexpectedType(format!("{self}")))
342 }
343 }
344
345 pub fn expect_struct(&self) -> ShapeResult<StructType> {
346 if let PartiqlShape::Static(StaticType {
347 ty: Static::Struct(stct),
348 ..
349 }) = self
350 {
351 Ok(stct.clone())
352 } else {
353 Err(ShapeResultError::UnexpectedType(format!("{self}")))
354 }
355 }
356
357 pub fn expect_static(&self) -> ShapeResult<StaticType> {
358 if let PartiqlShape::Static(s) = self {
359 Ok(s.clone())
360 } else {
361 Err(ShapeResultError::UnexpectedType(format!("{self}")))
362 }
363 }
364
365 pub fn expect_dynamic_type(&self) -> ShapeResult<PartiqlShape> {
366 if let PartiqlShape::Dynamic = self {
367 Ok(PartiqlShape::Dynamic)
368 } else {
369 Err(ShapeResultError::UnexpectedType(format!("{self}")))
370 }
371 }
372
373 pub fn expect_undefined(&self) -> ShapeResult<PartiqlShape> {
374 if let PartiqlShape::Undefined = self {
375 Ok(PartiqlShape::Undefined)
376 } else {
377 Err(ShapeResultError::UnexpectedType(format!("{self}")))
378 }
379 }
380}
381
382impl Display for PartiqlShape {
383 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
384 let x = match self {
385 PartiqlShape::Dynamic => "Dynamic".to_string(),
386 PartiqlShape::AnyOf(anyof) => {
387 format!("AnyOf({})", anyof.types.iter().cloned().join(", "))
388 }
389 PartiqlShape::Static(s) => format!("{s}"),
390 PartiqlShape::Undefined => "Undefined".to_string(),
391 };
392 write!(f, "{x}")
393 }
394}
395
396#[derive(Default, Debug)]
397pub struct ShapeBuilder<Id: NodeIdGenerator> {
398 id_gen: Id,
399}
400
401pub type PartiqlShapeBuilder = ShapeBuilder<AutoNodeIdGenerator>;
402pub type PartiqlNoIdShapeBuilder = ShapeBuilder<NullIdGenerator>;
403
404impl<Id: NodeIdGenerator + Default> ShapeBuilder<Id> {
405 #[track_caller]
408 pub fn singleton() -> &'static PartiqlShapeBuilder {
409 static SHAPE_BUILDER: OnceLock<PartiqlShapeBuilder> = OnceLock::new();
410 SHAPE_BUILDER.get_or_init(PartiqlShapeBuilder::default)
411 }
412
413 #[track_caller]
414 pub fn dummy_singleton() -> &'static PartiqlNoIdShapeBuilder {
415 static SHAPE_BUILDER: OnceLock<PartiqlNoIdShapeBuilder> = OnceLock::new();
416 SHAPE_BUILDER.get_or_init(PartiqlNoIdShapeBuilder::default)
417 }
418}
419
420impl<Id: NodeIdGenerator> ShapeBuilder<Id> {
421 #[inline]
422 pub fn new_static(&mut self, ty: Static) -> PartiqlShape {
423 let id = self.id_gen.next_id();
424 PartiqlShape::Static(StaticType {
425 id,
426 ty,
427 nullable: true,
428 })
429 }
430
431 #[inline]
432 pub fn new_non_nullable_static(&mut self, ty: Static) -> PartiqlShape {
433 let id = self.id_gen.next_id();
434 PartiqlShape::Static(StaticType {
435 id,
436 ty,
437 nullable: false,
438 })
439 }
440
441 #[inline]
442 pub fn new_dynamic(&mut self) -> PartiqlShape {
443 PartiqlShape::Dynamic
444 }
445
446 #[inline]
447 pub fn new_undefined(&mut self) -> PartiqlShape {
448 PartiqlShape::Undefined
449 }
450
451 #[inline]
452 pub fn new_struct(&mut self, s: StructType) -> PartiqlShape {
453 self.new_static(Static::Struct(s))
454 }
455
456 #[inline]
457 pub fn new_graph(&mut self) -> PartiqlShape {
458 self.new_static(Static::Graph())
459 }
460
461 #[inline]
462 pub fn new_struct_of_dyn(&mut self) -> PartiqlShape {
463 self.new_struct(StructType::new_any())
464 }
465
466 #[inline]
467 pub fn new_bag(&mut self, b: BagType) -> PartiqlShape {
468 self.new_static(Static::Bag(b))
469 }
470
471 #[inline]
472 pub fn new_bag_of<E>(&mut self, element_type: E) -> PartiqlShape
473 where
474 E: Into<PartiqlShape>,
475 {
476 self.new_bag(BagType::new_of(element_type))
477 }
478
479 #[inline]
480 pub fn new_bag_of_dyn(&mut self) -> PartiqlShape {
481 let element_type = self.new_dynamic();
482 self.new_bag_of(element_type)
483 }
484
485 #[inline]
486 pub fn new_bag_of_static(&mut self, ty: Static) -> PartiqlShape {
487 let element_type = self.new_static(ty);
488 self.new_bag_of(element_type)
489 }
490
491 #[inline]
492 pub fn new_bag_any_of<I>(&mut self, types: I) -> PartiqlShape
493 where
494 I: IntoIterator<Item = PartiqlShape>,
495 {
496 let shape = self.any_of(types);
497 let bag_type = BagType::new(Box::new(shape));
498 self.new_bag(bag_type)
499 }
500
501 #[inline]
502 pub fn new_array(&mut self, a: ArrayType) -> PartiqlShape {
503 self.new_static(Static::Array(a))
504 }
505
506 #[inline]
507 pub fn new_array_of<E>(&mut self, element_type: E) -> PartiqlShape
508 where
509 E: Into<PartiqlShape>,
510 {
511 self.new_array(ArrayType::new_of(element_type))
512 }
513
514 #[inline]
515 pub fn new_array_of_dyn(&mut self) -> PartiqlShape {
516 let element_type = self.new_dynamic();
517 self.new_array_of(element_type)
518 }
519
520 #[inline]
521 pub fn new_array_of_static(&mut self, ty: Static) -> PartiqlShape {
522 let element_type = self.new_static(ty);
523 self.new_array_of(element_type)
524 }
525
526 #[inline]
527 pub fn new_array_any_of<I>(&mut self, types: I) -> PartiqlShape
528 where
529 I: IntoIterator<Item = PartiqlShape>,
530 {
531 let shape = self.any_of(types);
532 let array_type = ArrayType::new(Box::new(shape));
533 self.new_array(array_type)
534 }
535
536 pub fn any_of<I>(&mut self, types: I) -> PartiqlShape
544 where
545 I: IntoIterator<Item = PartiqlShape>,
546 {
547 let any_of = AnyOf::from_iter(types);
548 match any_of.types.len() {
549 0 => self.new_dynamic(),
550 1 => {
551 let AnyOf { types } = any_of;
552 types.into_iter().next().unwrap()
553 }
554 _ => PartiqlShape::AnyOf(any_of),
556 }
557 }
558
559 #[inline]
560 pub fn union(&mut self, lhs: PartiqlShape, rhs: PartiqlShape) -> PartiqlShape {
561 match (lhs, rhs) {
562 (PartiqlShape::Dynamic, _) | (_, PartiqlShape::Dynamic) => PartiqlShape::Dynamic,
563 (PartiqlShape::AnyOf(lhs), PartiqlShape::AnyOf(rhs)) => {
564 self.any_of(lhs.types.into_iter().chain(rhs.types))
565 }
566 (PartiqlShape::AnyOf(anyof), other) | (other, PartiqlShape::AnyOf(anyof)) => {
567 let mut types = anyof.types;
568 types.insert(other);
569 self.any_of(types)
570 }
571 (l, r) => {
572 let types = [l, r];
573 self.any_of(types)
574 }
575 }
576 }
577
578 pub fn union_of<I>(&mut self, types: I) -> PartiqlShape
579 where
580 I: IntoIterator<Item = PartiqlShape>,
581 {
582 let types: Vec<_> = types.into_iter().collect();
583 match types.len() {
584 0 => self.new_undefined(),
585 1 => types.into_iter().next().unwrap(),
586 _ => types.into_iter().reduce(|l, r| self.union(l, r)).unwrap(),
587 }
588 }
589
590 #[inline]
591 pub fn as_non_nullable(&mut self, shape: &PartiqlShape) -> Option<PartiqlShape> {
592 if let PartiqlShape::Static(stype) = shape {
593 Some(self.new_non_nullable_static(stype.ty.clone()))
594 } else {
595 None
596 }
597 }
598}
599
600pub trait ShapeBuilderExtensions<Id: NodeIdGenerator> {
601 fn into_union(self, bld: &mut ShapeBuilder<Id>) -> PartiqlShape;
602
603 fn into_any_of(self, bld: &mut ShapeBuilder<Id>) -> PartiqlShape;
604
605 fn into_array(self, bld: &mut ShapeBuilder<Id>) -> PartiqlShape;
606
607 fn into_bag(self, bld: &mut ShapeBuilder<Id>) -> PartiqlShape;
608}
609
610impl<const N: usize, Id: NodeIdGenerator, E: Into<PartiqlShape>> ShapeBuilderExtensions<Id>
611 for [E; N]
612{
613 #[inline]
614 fn into_union(self, bld: &mut ShapeBuilder<Id>) -> PartiqlShape {
615 bld.union_of(self.into_iter().map(|e| e.into()))
616 }
617
618 #[inline]
619 fn into_any_of(self, bld: &mut ShapeBuilder<Id>) -> PartiqlShape {
620 bld.any_of(self.into_iter().map(|e| e.into()))
621 }
622
623 #[inline]
624 fn into_array(self, bld: &mut ShapeBuilder<Id>) -> PartiqlShape {
625 let ty = self.into_any_of(bld);
626 bld.new_array_of(ty)
627 }
628
629 #[inline]
630 fn into_bag(self, bld: &mut ShapeBuilder<Id>) -> PartiqlShape {
631 let ty = self.into_any_of(bld);
632 bld.new_bag_of(ty)
633 }
634}
635
636impl<Id: NodeIdGenerator, E: Into<PartiqlShape>> ShapeBuilderExtensions<Id> for E {
637 #[inline]
638 fn into_union(self, bld: &mut ShapeBuilder<Id>) -> PartiqlShape {
639 [self].into_union(bld)
640 }
641
642 #[inline]
643 fn into_any_of(self, bld: &mut ShapeBuilder<Id>) -> PartiqlShape {
644 [self].into_any_of(bld)
645 }
646
647 #[inline]
648 fn into_array(self, bld: &mut ShapeBuilder<Id>) -> PartiqlShape {
649 bld.new_array_of(self)
650 }
651
652 #[inline]
653 fn into_bag(self, bld: &mut ShapeBuilder<Id>) -> PartiqlShape {
654 bld.new_bag_of(self)
655 }
656}
657
658#[derive(Educe, Eq, Debug, Clone)]
659#[educe(PartialEq, Hash)]
660#[allow(dead_code)]
661pub struct AnyOf {
662 #[educe(Hash(method(indexset_hash)))]
663 types: IndexSet<PartiqlShape>,
664}
665
666impl AnyOf {
667 #[inline]
668 pub const fn new(types: IndexSet<PartiqlShape>) -> Self {
669 AnyOf { types }
670 }
671
672 pub fn types(&self) -> impl Iterator<Item = &PartiqlShape> {
673 self.types.iter()
674 }
675}
676
677impl FromIterator<PartiqlShape> for AnyOf {
678 fn from_iter<T: IntoIterator<Item = PartiqlShape>>(iter: T) -> Self {
679 AnyOf {
680 types: iter.into_iter().collect(),
681 }
682 }
683}
684
685#[derive(Debug, Clone, Eq, PartialEq, Hash)]
686pub struct StaticType {
687 id: NodeId,
688 ty: Static,
689 nullable: bool,
690}
691
692impl StaticType {
693 #[inline]
694 pub fn ty(&self) -> &Static {
695 &self.ty
696 }
697
698 pub fn ty_id(&self) -> &NodeId {
699 &self.id
700 }
701
702 #[inline]
703 pub fn is_nullable(&self) -> bool {
704 self.nullable
705 }
706
707 #[inline]
708 pub fn is_not_nullable(&self) -> bool {
709 !self.nullable
710 }
711
712 pub fn is_scalar(&self) -> bool {
713 self.ty.is_scalar()
714 }
715
716 pub fn is_sequence(&self) -> bool {
717 self.ty.is_sequence()
718 }
719
720 pub fn is_struct(&self) -> bool {
721 self.ty.is_struct()
722 }
723}
724
725impl From<StaticType> for PartiqlShape {
726 fn from(value: StaticType) -> Self {
727 PartiqlShape::Static(value)
728 }
729}
730
731impl Display for StaticType {
732 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
733 let ty = &self.ty;
734 if self.nullable {
735 write!(f, "{ty}")
736 } else {
737 write!(f, "NOT NULL {ty}")
738 }
739 }
740}
741
742#[derive(Debug, Clone, Eq, PartialEq, Hash)]
743#[non_exhaustive]
744pub enum Static {
745 Int,
747 Int8,
748 Int16,
749 Int32,
750 Int64,
751 Bool,
752 Decimal,
753 DecimalP(usize, usize),
754
755 Float32,
756 Float64,
757
758 String,
759 StringFixed(usize),
760 StringVarying(usize),
761
762 DateTime,
763
764 Struct(StructType),
766 Bag(BagType),
767 Array(ArrayType),
768
769 Graph(
770 ),
772 }
774
775#[non_exhaustive]
776pub enum StaticCategory<'a> {
777 Graph(),
778 Tuple(),
779 Sequence(&'a PartiqlShape),
780 Scalar(&'a Static),
781}
782
783impl Static {
784 pub fn is_scalar(&self) -> bool {
785 !matches!(self, Static::Struct(_) | Static::Bag(_) | Static::Array(_))
786 }
787
788 pub fn is_sequence(&self) -> bool {
789 matches!(self, Static::Bag(_) | Static::Array(_))
790 }
791
792 pub fn is_struct(&self) -> bool {
793 matches!(self, Static::Struct(_))
794 }
795
796 pub fn category(&self) -> StaticCategory<'_> {
797 match self {
798 Static::Struct(_) => StaticCategory::Tuple(),
799 Static::Bag(b) => StaticCategory::Sequence(b.element_type()),
800 Static::Array(b) => StaticCategory::Sequence(b.element_type()),
801 Static::Graph() => StaticCategory::Graph(),
802 _ => StaticCategory::Scalar(self),
803 }
804 }
805}
806
807impl Display for Static {
809 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
810 match self {
811 Static::Int => write!(f, "Int"),
812 Static::Int8 => write!(f, "Int8"),
813 Static::Int16 => write!(f, "Int16"),
814 Static::Int32 => write!(f, "Int32"),
815 Static::Int64 => write!(f, "Int64"),
816 Static::Bool => write!(f, "Bool"),
817 Static::Decimal => write!(f, "Decimal"),
818 Static::DecimalP(p, s) => {
819 write!(f, "Decimal({p},{s})")
820 }
821 Static::Float32 => write!(f, "Float32"),
822 Static::Float64 => write!(f, "Float64"),
823 Static::String => write!(f, "String"),
824 Static::StringFixed(_) => {
825 todo!()
826 }
827 Static::StringVarying(_) => {
828 todo!()
829 }
830 Static::DateTime => write!(f, "DateTime"),
831 Static::Struct(inner) => std::fmt::Display::fmt(inner, f),
832 Static::Bag(inner) => std::fmt::Display::fmt(inner, f),
833 Static::Array(inner) => std::fmt::Display::fmt(inner, f),
834 Static::Graph() => write!(f, "Graph"),
835 }
836 }
837}
838
839pub const TYPE_DYNAMIC: PartiqlShape = PartiqlShape::Dynamic;
840
841#[derive(Educe, Eq, Debug, Clone)]
842#[educe(PartialEq, Hash)]
843#[allow(dead_code)]
844pub struct StructType {
845 #[educe(Hash(method(indexset_hash)))]
846 constraints: IndexSet<StructConstraint>,
847}
848
849impl StructType {
850 #[inline]
851 pub fn new(constraints: IndexSet<StructConstraint>) -> Self {
852 StructType { constraints }
853 }
854
855 #[inline]
856 pub fn new_any() -> Self {
857 StructType {
858 constraints: Default::default(),
859 }
860 }
861
862 pub fn fields_set(&self) -> IndexSet<StructField> {
863 self.constraints
864 .iter()
865 .flat_map(|c| {
866 if let StructConstraint::Fields(fields) = c.clone() {
867 fields
868 } else {
869 Default::default()
870 }
871 })
872 .collect()
873 }
874
875 pub fn fields(&self) -> impl Iterator<Item = &StructField> {
876 self.constraints
877 .iter()
878 .filter_map(|c| {
879 if let StructConstraint::Fields(fields) = c {
880 Some(fields)
881 } else {
882 None
883 }
884 })
885 .flat_map(|f| f.iter())
886 }
887
888 #[inline]
889 pub fn is_partial(&self) -> bool {
890 !self.is_closed()
891 }
892
893 #[inline]
894 pub fn is_closed(&self) -> bool {
895 self.constraints.contains(&StructConstraint::Open(false))
896 }
897}
898
899impl Display for StructType {
900 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
901 let partial = self.is_partial();
902 write!(f, "{{")?;
903 let mut first = true;
904 for StructField { name, ty, optional } in self.fields() {
905 if !first {
906 write!(f, ", ")?;
907 }
908 if *optional {
909 write!(f, "{name}?: {ty}")?;
910 } else {
911 write!(f, "{name}: {ty}")?;
912 }
913 first = false
914 }
915 if partial {
916 if !first {
917 write!(f, ", ")?;
918 }
919 write!(f, "...")?;
920 }
921 write!(f, "}}")
922 }
923}
924
925#[derive(Educe, Eq, Debug, Clone)]
926#[educe(PartialEq, Hash)]
927#[allow(dead_code)]
928#[non_exhaustive]
929pub enum StructConstraint {
930 Open(bool),
931 Ordered(bool),
932 DuplicateAttrs(bool),
933 Fields(#[educe(Hash(method(indexset_hash)))] IndexSet<StructField>),
934}
935
936#[derive(Debug, Clone, Hash, Eq, PartialEq)]
937#[allow(dead_code)]
938pub struct StructField {
939 optional: bool,
940 name: String,
941 ty: PartiqlShape,
942}
943
944impl StructField {
945 #[inline]
946 pub fn new(name: &str, ty: PartiqlShape) -> Self {
947 StructField {
948 name: name.to_string(),
949 ty,
950 optional: false,
951 }
952 }
953
954 #[inline]
955 pub fn new_optional(name: &str, ty: PartiqlShape) -> Self {
956 StructField {
957 name: name.to_string(),
958 ty,
959 optional: true,
960 }
961 }
962
963 #[inline]
964 pub fn name(&self) -> &str {
965 self.name.as_str()
966 }
967
968 #[inline]
969 pub fn ty(&self) -> &PartiqlShape {
970 &self.ty
971 }
972
973 #[inline]
974 pub fn is_optional(&self) -> bool {
975 self.optional
976 }
977}
978
979impl From<(&str, PartiqlShape)> for StructField {
980 fn from(value: (&str, PartiqlShape)) -> Self {
981 StructField {
982 name: value.0.to_string(),
983 ty: value.1,
984 optional: false,
985 }
986 }
987}
988
989impl From<(&str, PartiqlShape, bool)> for StructField {
990 fn from(value: (&str, PartiqlShape, bool)) -> Self {
991 StructField {
992 name: value.0.to_string(),
993 ty: value.1,
994 optional: value.2,
995 }
996 }
997}
998
999#[derive(Educe, Eq, Debug, Clone)]
1000#[educe(PartialEq, Hash)]
1001#[allow(dead_code)]
1002pub struct BagType {
1003 element_type: Box<PartiqlShape>,
1004}
1005
1006impl BagType {
1007 #[inline]
1008 pub fn new(typ: Box<PartiqlShape>) -> Self {
1009 BagType { element_type: typ }
1010 }
1011
1012 #[inline]
1013 pub fn new_of<E>(element_type: E) -> Self
1014 where
1015 E: Into<PartiqlShape>,
1016 {
1017 Self::new(Box::new(element_type.into()))
1018 }
1019
1020 #[inline]
1021 pub fn new_any() -> Self {
1022 Self::new_of(PartiqlShape::Dynamic)
1023 }
1024
1025 pub fn element_type(&self) -> &PartiqlShape {
1026 &self.element_type
1027 }
1028}
1029
1030impl Display for BagType {
1031 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
1032 let ty = &self.element_type;
1033 write!(f, "<<{ty}>>")
1034 }
1035}
1036
1037#[derive(Educe, Eq, Debug, Clone)]
1038#[educe(PartialEq, Hash)]
1039#[allow(dead_code)]
1040pub struct ArrayType {
1041 element_type: Box<PartiqlShape>,
1042 }
1045
1046impl ArrayType {
1047 #[inline]
1048 pub fn new(typ: Box<PartiqlShape>) -> Self {
1049 ArrayType { element_type: typ }
1050 }
1051
1052 #[inline]
1053 pub fn new_of<E>(element_type: E) -> Self
1054 where
1055 E: Into<PartiqlShape>,
1056 {
1057 Self::new(Box::new(element_type.into()))
1058 }
1059
1060 #[inline]
1061 pub fn new_any() -> Self {
1062 Self::new_of(PartiqlShape::Dynamic)
1063 }
1064
1065 #[inline]
1066 pub fn element_type(&self) -> &PartiqlShape {
1067 &self.element_type
1068 }
1069}
1070
1071impl Display for ArrayType {
1072 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
1073 let ty = &self.element_type;
1074 write!(f, "[{ty}]")
1075 }
1076}
1077
1078#[cfg(test)]
1079mod tests {
1080 use crate::{
1081 PartiqlNoIdShapeBuilder, PartiqlShape, PartiqlShapeBuilder, ShapeBuilderExtensions, Static,
1082 StructConstraint, StructField, StructType,
1083 };
1084 use indexmap::IndexSet;
1085
1086 #[test]
1087 fn union() {
1088 let mut bld = PartiqlNoIdShapeBuilder::default();
1089
1090 let expect_int = bld.new_static(Static::Int);
1091 assert_eq!(
1092 expect_int,
1093 [type_int!(bld), type_int!(bld)].into_union(&mut bld)
1094 );
1095
1096 let numbers = [bld.new_static(Static::Int), bld.new_static(Static::Float32)];
1097 let expect_nums = bld.any_of(numbers);
1098 assert_eq!(
1099 expect_nums,
1100 [type_int!(bld), type_float32!(bld)].into_union(&mut bld)
1101 );
1102 assert_eq!(
1103 expect_nums,
1104 [
1105 [type_int!(bld), type_float32!(bld)].into_union(&mut bld),
1106 [type_int!(bld), type_float32!(bld)].into_union(&mut bld),
1107 ]
1108 .into_any_of(&mut bld)
1109 );
1110 assert_eq!(
1111 expect_nums,
1112 [
1113 [type_int!(bld), type_float32!(bld)].into_union(&mut bld),
1114 [type_int!(bld), type_float32!(bld)].into_union(&mut bld),
1115 [
1116 [type_int!(bld), type_float32!(bld)].into_union(&mut bld),
1117 [type_int!(bld), type_float32!(bld)].into_union(&mut bld)
1118 ]
1119 .into_any_of(&mut bld)
1120 ]
1121 .into_any_of(&mut bld)
1122 );
1123 }
1124
1125 #[test]
1126 fn unique_node_ids() {
1127 let mut bld = PartiqlShapeBuilder::default();
1128 let age_field = struct_fields![("age", type_int!(bld))];
1129 let details = type_struct![bld, IndexSet::from([age_field])];
1130
1131 let fields = [
1132 StructField::new("id", type_int!(bld)),
1133 StructField::new("name", type_string!(bld)),
1134 StructField::new("details", details.clone()),
1135 ];
1136
1137 let row = type_struct![
1138 bld,
1139 IndexSet::from([
1140 StructConstraint::Fields(IndexSet::from(fields)),
1141 StructConstraint::Open(false)
1142 ])
1143 ];
1144
1145 let shape = type_bag![bld, row.clone()];
1146
1147 let mut ids = collect_ids(shape);
1148 ids.sort_unstable();
1149 assert!(ids.windows(2).all(|w| w[0] != w[1]));
1150 }
1151
1152 fn collect_ids(row: PartiqlShape) -> Vec<u32> {
1153 let mut out = vec![];
1154 match row {
1155 PartiqlShape::Dynamic => {}
1156 PartiqlShape::AnyOf(anyof) => {
1157 for shape in anyof.types {
1158 out.push(collect_ids(shape));
1159 }
1160 }
1161 PartiqlShape::Static(static_type) => {
1162 out.push(vec![static_type.id.0]);
1163 match static_type.ty {
1164 Static::Struct(struct_type) => {
1165 for f in struct_type.fields() {
1166 out.push(collect_ids(f.ty.clone()));
1167 }
1168 }
1169 Static::Bag(bag_type) => out.push(collect_ids(*bag_type.element_type)),
1170 Static::Array(array_type) => out.push(collect_ids(*array_type.element_type)),
1171 _ => {}
1172 }
1173 }
1174 PartiqlShape::Undefined => {}
1175 }
1176 out.into_iter().flatten().collect()
1177 }
1178}