Skip to main content

mz_expr_derive_impl/
shape.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9//! Per-arity descriptions of the three scalar function shapes.
10//!
11//! A [`Shape`] answers the questions that differ between `EagerUnaryFunc`,
12//! `EagerBinaryFunc`, and `EagerVariadicFunc`. Adding a modifier to an arity means
13//! adding a row to that arity's table here, not writing expansion code.
14
15use proc_macro2::{Ident, Span, TokenStream};
16use quote::quote;
17
18/// The three scalar function arities the macro generates for.
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub(crate) enum Shape {
21    Unary,
22    Binary,
23    Variadic,
24}
25
26/// A modifier that maps directly onto one optional trait method.
27///
28/// Modifiers that do not produce a trait method, such as `sqlname`, `output_type`,
29/// `output_type_expr`, `test`, and `skip_display`, are absent: `crate::generate::generate`
30/// handles those explicitly because they feed `Display`, the output-type body, or the
31/// expansion decision instead.
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub(crate) enum Modifier {
34    CouldError,
35    IntroducesNulls,
36    IsMonotone,
37    IsInfixOp,
38    PropagatesNulls,
39    Inverse,
40    PreservesUniqueness,
41    IsEliminableCast,
42    Negate,
43    IsInfinityMonotone,
44    IsAssociative,
45}
46
47impl Modifier {
48    /// Every method-producing modifier.
49    pub(crate) const ALL: [Modifier; 11] = [
50        Modifier::CouldError,
51        Modifier::IntroducesNulls,
52        Modifier::IsMonotone,
53        Modifier::IsInfixOp,
54        Modifier::PropagatesNulls,
55        Modifier::Inverse,
56        Modifier::PreservesUniqueness,
57        Modifier::IsEliminableCast,
58        Modifier::Negate,
59        Modifier::IsInfinityMonotone,
60        Modifier::IsAssociative,
61    ];
62
63    /// The attribute key, which is also the generated method name.
64    pub(crate) fn name(&self) -> &'static str {
65        match self {
66            Modifier::CouldError => "could_error",
67            Modifier::IntroducesNulls => "introduces_nulls",
68            Modifier::IsMonotone => "is_monotone",
69            Modifier::IsInfixOp => "is_infix_op",
70            Modifier::PropagatesNulls => "propagates_nulls",
71            Modifier::Inverse => "inverse",
72            Modifier::PreservesUniqueness => "preserves_uniqueness",
73            Modifier::IsEliminableCast => "is_eliminable_cast",
74            Modifier::Negate => "negate",
75            Modifier::IsInfinityMonotone => "is_infinity_monotone",
76            Modifier::IsAssociative => "is_associative",
77        }
78    }
79}
80
81/// The return type of a generated override method.
82#[derive(Debug, Clone, Copy, PartialEq, Eq)]
83pub(crate) enum ReturnTy {
84    Bool,
85    BoolPair,
86    OptUnaryFunc,
87    OptBinaryFunc,
88}
89
90impl ReturnTy {
91    pub(crate) fn tokens(&self) -> TokenStream {
92        match self {
93            ReturnTy::Bool => quote! { bool },
94            ReturnTy::BoolPair => quote! { (bool, bool) },
95            ReturnTy::OptUnaryFunc => quote! { Option<crate::UnaryFunc> },
96            ReturnTy::OptBinaryFunc => quote! { Option<crate::BinaryFunc> },
97        }
98    }
99
100    /// The method body for a modifier whose value is `expr`. A function-valued
101    /// modifier names the function, and the body converts it into the enum.
102    pub(crate) fn body(&self, expr: &syn::Expr) -> TokenStream {
103        match self {
104            ReturnTy::Bool | ReturnTy::BoolPair => quote! { #expr },
105            ReturnTy::OptUnaryFunc => quote! { Some(crate::UnaryFunc::from(#expr)) },
106            ReturnTy::OptBinaryFunc => quote! { Some(crate::BinaryFunc::from(#expr)) },
107        }
108    }
109}
110
111// `crate::generate::override_methods` walks these tables in order, so they also fix
112// the method order in the generated code, which the snapshots in `crate::test` pin.
113const UNARY_MODIFIERS: &[(Modifier, ReturnTy)] = &[
114    (Modifier::CouldError, ReturnTy::Bool),
115    (Modifier::IntroducesNulls, ReturnTy::Bool),
116    (Modifier::Inverse, ReturnTy::OptUnaryFunc),
117    (Modifier::IsMonotone, ReturnTy::Bool),
118    (Modifier::PreservesUniqueness, ReturnTy::Bool),
119    (Modifier::IsEliminableCast, ReturnTy::Bool),
120];
121
122const BINARY_MODIFIERS: &[(Modifier, ReturnTy)] = &[
123    (Modifier::CouldError, ReturnTy::Bool),
124    (Modifier::IntroducesNulls, ReturnTy::Bool),
125    (Modifier::IsInfixOp, ReturnTy::Bool),
126    (Modifier::IsMonotone, ReturnTy::BoolPair),
127    (Modifier::IsInfinityMonotone, ReturnTy::Bool),
128    (Modifier::Negate, ReturnTy::OptBinaryFunc),
129    (Modifier::PropagatesNulls, ReturnTy::Bool),
130];
131
132const VARIADIC_MODIFIERS: &[(Modifier, ReturnTy)] = &[
133    (Modifier::CouldError, ReturnTy::Bool),
134    (Modifier::IntroducesNulls, ReturnTy::Bool),
135    (Modifier::IsInfixOp, ReturnTy::Bool),
136    (Modifier::IsMonotone, ReturnTy::Bool),
137    (Modifier::IsAssociative, ReturnTy::Bool),
138    (Modifier::PropagatesNulls, ReturnTy::Bool),
139];
140
141impl Shape {
142    /// The modifiers this arity accepts, with the return type of each generated
143    /// method. A modifier absent from this table is rejected.
144    pub(crate) fn modifiers(&self) -> &'static [(Modifier, ReturnTy)] {
145        match self {
146            Shape::Unary => UNARY_MODIFIERS,
147            Shape::Binary => BINARY_MODIFIERS,
148            Shape::Variadic => VARIADIC_MODIFIERS,
149        }
150    }
151
152    /// Whether this arity's table lists `modifier`.
153    pub(crate) fn accepts(&self, modifier: Modifier) -> bool {
154        self.modifiers()
155            .iter()
156            .any(|(candidate, _)| *candidate == modifier)
157    }
158
159    pub(crate) fn label(&self) -> &'static str {
160        match self {
161            Shape::Unary => "unary",
162            Shape::Binary => "binary",
163            Shape::Variadic => "variadic",
164        }
165    }
166
167    /// The trait the generated impl implements.
168    pub(crate) fn trait_path(&self) -> TokenStream {
169        match self {
170            Shape::Unary => quote! { crate::func::EagerUnaryFunc },
171            Shape::Binary => quote! { crate::func::binary::EagerBinaryFunc },
172            Shape::Variadic => quote! { crate::func::variadic::EagerVariadicFunc },
173        }
174    }
175
176    /// The name and sole parameter of the trait method that computes the output
177    /// `SqlColumnType`.
178    ///
179    /// `EagerVariadicFunc` declares no `output_sql_type`. Its core method is named
180    /// `output_type` and takes the SQL column types directly, where the unary and
181    /// binary traits declare `output_sql_type` and keep `output_type` as a wrapper
182    /// over `ReprColumnType`.
183    pub(crate) fn output_method(&self) -> (Ident, TokenStream) {
184        let span = Span::call_site();
185        match self {
186            Shape::Unary => (
187                Ident::new("output_sql_type", span),
188                quote! { input_type: mz_repr::SqlColumnType },
189            ),
190            Shape::Binary => (
191                Ident::new("output_sql_type", span),
192                quote! { input_types: &[mz_repr::SqlColumnType] },
193            ),
194            Shape::Variadic => (
195                Ident::new("output_type", span),
196                quote! { input_types: &[mz_repr::SqlColumnType] },
197            ),
198        }
199    }
200
201    /// The tail of the output-type method, which sets the nullability of the
202    /// `output` column type the caller computed.
203    ///
204    /// The caller binds `output`, `nullable`, and `propagates_nulls`, and names the
205    /// input parameter as [`Shape::output_method`] spells it. `checks` are
206    /// per-position nullability checks over `input_types`. Unary receives a single
207    /// column type, so it has no position to check and ignores them.
208    ///
209    /// The result is nullable when the function itself introduces nulls, when a
210    /// parameter that rejects NULL is handed a nullable input, because the runtime
211    /// conversion then rejects the NULL, or when the function propagates nulls and
212    /// some input is nullable, because the optimizer short-circuits an all-NULL call.
213    pub(crate) fn nullability(&self, checks: &[TokenStream]) -> TokenStream {
214        match self {
215            // NOTE: unary omits the non-nullable-position term the other two arities
216            // carry. That is sound only because `PropagatesNulls` is absent from
217            // `UNARY_MODIFIERS`, so a unary `propagates_nulls` is always
218            // `!Input::nullable()`, which is what the position check tests for every
219            // parameter type in the tree.
220            // Adding `PropagatesNulls` to the unary table without also adding the term
221            // here lets `propagates_nulls = false` report a non-nullable output for an
222            // input the evaluation layer short-circuits to NULL.
223            Shape::Unary => quote! {
224                output.nullable(nullable || (propagates_nulls && input_type.nullable))
225            },
226            Shape::Binary | Shape::Variadic => quote! {
227                let non_nullable_input_is_nullable = false #(#checks)*;
228                let inputs_nullable = input_types.iter().any(|it| it.nullable);
229                let is_null = nullable
230                    || non_nullable_input_is_nullable
231                    || (propagates_nulls && inputs_nullable);
232                output.nullable(is_null)
233            },
234        }
235    }
236}
237
238#[cfg(test)]
239mod tests {
240    use super::{Modifier, Shape};
241
242    #[mz_ore::test]
243    fn unary_accepts_inverse_and_rejects_negate() {
244        assert!(Shape::Unary.accepts(Modifier::Inverse));
245        assert!(Shape::Unary.accepts(Modifier::PreservesUniqueness));
246        assert!(Shape::Unary.accepts(Modifier::IsEliminableCast));
247        assert!(!Shape::Unary.accepts(Modifier::Negate));
248        assert!(!Shape::Unary.accepts(Modifier::IsInfixOp));
249        assert!(
250            !Shape::Unary.accepts(Modifier::PropagatesNulls),
251            "Shape::nullability omits unary's non-nullable-position term on the strength of \
252             this absence"
253        );
254        assert!(!Shape::Unary.accepts(Modifier::IsAssociative));
255    }
256}