Skip to main content

partiql_types/
lib.rs

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
22/// Result of attempts to encode to Ion.
23pub 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// TODO add macro_rule for Decimal with precision and scale
89
90#[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/// Represents a PartiQL Shape
186#[derive(Debug, Clone, Eq, PartialEq, Hash)]
187// With this implementation `Dynamic` and `AnyOf` cannot have `nullability`; this does not mean their
188// `null` value at runtime cannot belong to their domain.
189// TODO adopt the correct model Pending PartiQL Types semantics finalization: https://github.com/partiql/partiql-lang/issues/18
190#[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        // TODO Add Sexp when added
260        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    /// A thread-safe method for creating PartiQL shapes with guaranteed uniqueness over
406    /// generated `NodeId`s.
407    #[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    // The AnyOf::from_iter(types) uses an IndexSet internally to
537    // deduplicate types, thus the match on any_of.types.len() could
538    // "flatten" AnyOfs that had duplicates.
539    // With the addition of IDs, this deduplication no longer happens.
540    // TODO revisit the current implementaion and consider an implementation
541    // that allows merging of the `metas` for the same type, e.g., with a
542    // user-defined control.
543    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            // TODO figure out what does it mean for a Union to be nullable or not
555            _ => 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    // Scalar Types
746    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    // Container Types
765    Struct(StructType),
766    Bag(BagType),
767    Array(ArrayType),
768
769    Graph(
770        /* TODO: https://github.com/partiql/partiql-lang/blob/main/RFCs/0025-graph-data-model.md */
771    ),
772    // TODO Add BitString, ByteString, Blob, Clob, and Graph types
773}
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
807// TODO, this should probably be via a prettyprint...
808impl 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    // TODO Add Array constraint once we have Schema Specification:
1043    // https://github.com/partiql/partiql-spec/issues/49
1044}
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}