mz_expr_derive_impl/
source.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//! Source text the scalar function registry records for each `#[sqlfunc]`.
10//!
11//! The registry compares a function's declaration, its types-only signature, and a
12//! fingerprint of its body across builds, so every string here depends only on the
13//! tokens the call site wrote, not on how they were spaced or commented.
14
15use proc_macro2::{Delimiter, Spacing, TokenStream, TokenTree};
16use quote::{ToTokens, quote};
17use syn::Lifetime;
18use syn::spanned::Spanned;
19
20/// Renders a token stream as compact, canonical source text.
21///
22/// The rendering walks the token trees rather than calling
23/// `TokenStream::to_string`, whose spacing differs between the compiler and
24/// the proc-macro2 fallback and has changed across rustc releases. Idents
25/// and literals are rendered verbatim, so whitespace inside a string literal
26/// is preserved and distinguishes two bodies. Spacing between tokens follows
27/// fixed rules that read like formatted code for signatures and attribute
28/// arguments, and is merely deterministic for bodies.
29fn render_tokens(tokens: &TokenStream) -> String {
30    let mut out = String::new();
31    render_into(tokens.clone(), &mut out);
32    out
33}
34
35fn render_into(tokens: TokenStream, out: &mut String) {
36    /// What the previous token was, as far as spacing cares.
37    enum Last {
38        Start,
39        Ident,
40        Punct { ch: char, joint: bool },
41        Other,
42    }
43    let mut last = Last::Start;
44    for tree in tokens {
45        let glue = match &last {
46            Last::Start => true,
47            Last::Punct { joint: true, .. } => true,
48            Last::Punct { ch, .. } if "&.!#<".contains(*ch) => true,
49            _ => out.ends_with("::"),
50        } || match &tree {
51            TokenTree::Punct(p) if ",;:.?>".contains(p.as_char()) => true,
52            // Generic parameter lists and macro invocations attach to the
53            // preceding ident.
54            TokenTree::Punct(p) if "<!".contains(p.as_char()) => matches!(last, Last::Ident),
55            // Call and index groups attach to the callee, including a closing
56            // generic angle bracket, but not to an arrow's `>`.
57            TokenTree::Group(g) => {
58                matches!(g.delimiter(), Delimiter::Parenthesis | Delimiter::Bracket)
59                    && (matches!(last, Last::Ident)
60                        || (matches!(last, Last::Punct { ch: '>', .. })
61                            && !out.ends_with("->")
62                            && !out.ends_with("=>")))
63            }
64            _ => false,
65        };
66        if !glue {
67            out.push(' ');
68        }
69        last = match tree {
70            TokenTree::Ident(ident) => {
71                out.push_str(&ident.to_string());
72                Last::Ident
73            }
74            TokenTree::Literal(lit) => {
75                out.push_str(&lit.to_string());
76                Last::Other
77            }
78            TokenTree::Punct(p) => {
79                out.push(p.as_char());
80                Last::Punct {
81                    ch: p.as_char(),
82                    joint: p.spacing() == Spacing::Joint,
83                }
84            }
85            TokenTree::Group(g) => {
86                let (open, close) = match g.delimiter() {
87                    Delimiter::Parenthesis => ("(", ")"),
88                    Delimiter::Bracket => ("[", "]"),
89                    Delimiter::Brace => ("{", "}"),
90                    Delimiter::None => ("", ""),
91                };
92                out.push_str(open);
93                render_into(g.stream(), out);
94                out.push_str(close);
95                Last::Other
96            }
97        };
98    }
99}
100
101/// FNV-1a, 64 bit. A fingerprint, not a secure hash: it only needs to be
102/// stable across builds and sensitive to any change in its input.
103fn fnv1a64(bytes: &[u8]) -> u64 {
104    bytes.iter().fold(0xcbf29ce484222325, |hash, byte| {
105        (hash ^ u64::from(*byte)).wrapping_mul(0x100000001b3)
106    })
107}
108
109/// Emits the source-derived members of the generated `FuncName` impl: the
110/// `SQLFUNC` const (declaration text, types-only signature, body
111/// fingerprint) and `sqlfunc_input_types`, which yields the column types the
112/// function naturally consumes when every parameter type has one. Both exist
113/// only under mz-expr's `func-registry` feature, like the trait members they
114/// implement.
115///
116/// Text comes from the token trees via [`render_tokens`], so formatting and
117/// comments do not affect it. `input_tys_raw` are the parameter types
118/// before generic erasure, with `'a` added to elided references, for the signature. `input_tys` are the erased types the trait
119/// impl uses, for the probes.
120pub(crate) fn sqlfunc_source(
121    attr: &TokenStream,
122    func: &syn::ItemFn,
123    input_tys_raw: &[syn::Type],
124    output_ty_raw: &syn::Type,
125    input_tys: &[syn::Type],
126) -> TokenStream {
127    let attr = render_tokens(attr);
128    let attr = if attr.is_empty() {
129        String::new()
130    } else {
131        format!("({attr})")
132    };
133    let decl = format!(
134        "#[sqlfunc{attr}] {}",
135        render_tokens(&func.sig.to_token_stream())
136    );
137    let render_type = |ty: &syn::Type| render_tokens(&ty.to_token_stream());
138    let signature = format!(
139        "fn({}) -> {}",
140        input_tys_raw
141            .iter()
142            .map(render_type)
143            .collect::<Vec<_>>()
144            .join(", "),
145        render_type(output_ty_raw)
146    );
147    let body_fingerprint = fnv1a64(render_tokens(&func.block.to_token_stream()).as_bytes());
148    let probes: Vec<TokenStream> = input_tys.iter().flat_map(probe_column_types).collect();
149    quote! {
150        #[cfg(feature = "func-registry")]
151        const SQLFUNC: Option<crate::func::SqlFuncSource> = Some(crate::func::SqlFuncSource {
152            decl: #decl,
153            signature: #signature,
154            body_fingerprint: #body_fingerprint,
155        });
156
157        #[cfg(feature = "func-registry")]
158        fn sqlfunc_input_types() -> Option<Vec<mz_repr::SqlColumnType>> {
159            use crate::func::registry::{ProbeColumnType as _, ProbeColumnTypeFallback as _};
160            [#(#probes),*].into_iter().collect()
161        }
162    }
163}
164
165/// One probe expression per datum a parameter consumes. `Variadic<T>` stands
166/// for two `T` arguments and `OptionalArg<T>` for one present `T`.
167fn probe_column_types(ty: &syn::Type) -> Vec<TokenStream> {
168    if let Some((wrapper, inner)) = single_generic_arg(ty) {
169        match wrapper.as_str() {
170            "Variadic" => return vec![probe_column_type(inner), probe_column_type(inner)],
171            "OptionalArg" => return vec![probe_column_type(inner)],
172            _ => {}
173        }
174    }
175    vec![probe_column_type(ty)]
176}
177
178/// An expression of type `Option<SqlColumnType>`: the column type of `ty` if
179/// it implements `AsColumnType`, else `None`. Resolved by autoref
180/// specialization on `ColumnTypeProbe`, so it needs no trait bound the macro
181/// cannot check.
182fn probe_column_type(ty: &syn::Type) -> TokenStream {
183    let ty = staticize_lifetimes(ty);
184    quote! {
185        (&crate::func::registry::ColumnTypeProbe::<#ty>(::std::marker::PhantomData)).column_type()
186    }
187}
188
189/// The last path segment's name and its single type argument, for types
190/// shaped like `Wrapper<T>`.
191fn single_generic_arg(ty: &syn::Type) -> Option<(String, &syn::Type)> {
192    let syn::Type::Path(path) = ty else {
193        return None;
194    };
195    let segment = path.path.segments.last()?;
196    let syn::PathArguments::AngleBracketed(args) = &segment.arguments else {
197        return None;
198    };
199    let mut types = args.args.iter().filter_map(|arg| match arg {
200        syn::GenericArgument::Type(ty) => Some(ty),
201        _ => None,
202    });
203    let inner = types.next()?;
204    if types.next().is_some() {
205        return None;
206    }
207    Some((segment.ident.to_string(), inner))
208}
209
210/// Replaces every lifetime in `ty` with `'static`, so a type written against
211/// the trait impl's `'a` can be named in a function body.
212fn staticize_lifetimes(ty: &syn::Type) -> syn::Type {
213    let mut ty = ty.clone();
214    fn walk(ty: &mut syn::Type) {
215        match ty {
216            syn::Type::Reference(r) => {
217                r.lifetime = Some(Lifetime::new("'static", r.span()));
218                walk(&mut r.elem);
219            }
220            syn::Type::Tuple(t) => t.elems.iter_mut().for_each(walk),
221            syn::Type::Path(p) => {
222                for segment in &mut p.path.segments {
223                    if let syn::PathArguments::AngleBracketed(args) = &mut segment.arguments {
224                        for arg in &mut args.args {
225                            match arg {
226                                syn::GenericArgument::Lifetime(lt) => {
227                                    *lt = Lifetime::new("'static", lt.span());
228                                }
229                                syn::GenericArgument::Type(ty) => walk(ty),
230                                _ => {}
231                            }
232                        }
233                    }
234                }
235            }
236            _ => {}
237        }
238    }
239    walk(&mut ty);
240    ty
241}
242
243#[cfg(test)]
244mod render_tests {
245    use super::render_tokens;
246
247    #[mz_ore::test]
248    fn signature_reads_like_formatted_code() {
249        let sig: proc_macro2::TokenStream = syn::parse_quote! {
250            fn f<'a>(mut a: &'a str, b: Option<i32>) -> Result<Cow<'a, str>, E>
251        };
252        assert_eq!(
253            render_tokens(&sig),
254            "fn f<'a>(mut a: &'a str, b: Option<i32>) -> Result<Cow<'a, str>, E>"
255        );
256        let attr: proc_macro2::TokenStream = syn::parse_quote! {
257            is_monotone = "(true, true)", sqlname = "+", could_error = false
258        };
259        assert_eq!(
260            render_tokens(&attr),
261            "is_monotone = \"(true, true)\", sqlname = \"+\", could_error = false"
262        );
263    }
264
265    #[mz_ore::test]
266    fn literal_contents_are_verbatim() {
267        let a: proc_macro2::TokenStream = syn::parse_quote! { format!("({x})") };
268        let b: proc_macro2::TokenStream = syn::parse_quote! { format!("( {x})") };
269        assert_ne!(render_tokens(&a), render_tokens(&b));
270        assert_eq!(render_tokens(&a), "format!(\"({x})\")");
271    }
272
273    #[mz_ore::test]
274    fn formatting_is_invisible() {
275        let a: proc_macro2::TokenStream = "fn f ( a : i32 ) -> i32 { a + 1 }".parse().unwrap();
276        let b: proc_macro2::TokenStream = "fn f(a:i32)->i32{a+1}".parse().unwrap();
277        assert_eq!(render_tokens(&a), render_tokens(&b));
278    }
279}
mz_expr_derive_impl/source.rs

mz_expr_derive_impl/
source.rs
