1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157

use crate::database::DatabaseExt;
use crate::query::QueryMacroInput;
use either::Either;
use proc_macro2::TokenStream;
use quote::{format_ident, quote, quote_spanned};
use sqlx_core::describe::Describe;
use syn::spanned::Spanned;
use syn::{Expr, ExprCast, ExprGroup, ExprType, Type};

/// Returns a tokenstream which typechecks the arguments passed to the macro
/// and binds them to `DB::Arguments` with the ident `query_args`.
pub fn quote_args<DB: DatabaseExt>(
    input: &QueryMacroInput,
    info: &Describe<DB>,
) -> crate::Result<TokenStream> {
    let db_path = DB::db_path();

    if input.arg_exprs.is_empty() {
        return Ok(quote! {
            let query_args = <#db_path as ::sqlx::database::HasArguments>::Arguments::default();
        });
    }

    let arg_names = (0..input.arg_exprs.len())
        .map(|i| format_ident!("arg{}", i))
        .collect::<Vec<_>>();

    let arg_name = &arg_names;
    let arg_expr = input.arg_exprs.iter().cloned().map(strip_wildcard);

    let arg_bindings = quote! {
        #(let #arg_name = &(#arg_expr);)*
    };

    let args_check = match info.parameters() {
        None | Some(Either::Right(_)) => {
            // all we can do is check arity which we did
            TokenStream::new()
        }

        Some(Either::Left(_)) if !input.checked => {
            // this is an `*_unchecked!()` macro invocation
            TokenStream::new()
        }

        Some(Either::Left(params)) => {
            params
                .iter()
                .zip(arg_names.iter().zip(&input.arg_exprs))
                .enumerate()
                .map(|(i, (param_ty, (name, expr)))| -> crate::Result<_> {
                    let param_ty = match get_type_override(expr) {
                        // cast or type ascription will fail to compile if the type does not match
                        // and we strip casts to wildcard
                        Some(_) => return Ok(quote!()),
                        None => {
                            DB::param_type_for_id(&param_ty)
                                .ok_or_else(|| {
                                    if let Some(feature_gate) = <DB as DatabaseExt>::get_feature_gate(&param_ty) {
                                        format!(
                                            "optional feature `{}` required for type {} of param #{}",
                                            feature_gate,
                                            param_ty,
                                            i + 1,
                                        )
                                    } else {
                                        format!("unsupported type {} for param #{}", param_ty, i + 1)
                                    }
                                })?
                                .parse::<TokenStream>()
                                .map_err(|_| format!("Rust type mapping for {} not parsable", param_ty))?

                        }
                    };

                    Ok(quote_spanned!(expr.span() =>
                        // this shouldn't actually run
                        if false {
                            use ::sqlx::ty_match::{WrapSameExt as _, MatchBorrowExt as _};

                            // evaluate the expression only once in case it contains moves
                            let expr = ::sqlx::ty_match::dupe_value(#name);

                            // if `expr` is `Option<T>`, get `Option<$ty>`, otherwise `$ty`
                            let ty_check = ::sqlx::ty_match::WrapSame::<#param_ty, _>::new(&expr).wrap_same();

                            // if `expr` is `&str`, convert `String` to `&str`
                            let (mut _ty_check, match_borrow) = ::sqlx::ty_match::MatchBorrow::new(ty_check, &expr);

                            _ty_check = match_borrow.match_borrow();

                            // this causes move-analysis to effectively ignore this block
                            ::std::panic!();
                        }
                    ))
                })
                .collect::<crate::Result<TokenStream>>()?
        }
    };

    let args_count = input.arg_exprs.len();

    Ok(quote! {
        #arg_bindings

        #args_check

        let mut query_args = <#db_path as ::sqlx::database::HasArguments>::Arguments::default();
        query_args.reserve(
            #args_count,
            0 #(+ ::sqlx::encode::Encode::<#db_path>::size_hint(#arg_name))*
        );
        #(query_args.add(#arg_name);)*
    })
}

fn get_type_override(expr: &Expr) -> Option<&Type> {
    match expr {
        Expr::Group(group) => get_type_override(&group.expr),
        Expr::Cast(cast) => Some(&cast.ty),
        Expr::Type(ascription) => Some(&ascription.ty),
        _ => None,
    }
}

fn strip_wildcard(expr: Expr) -> Expr {
    match expr {
        Expr::Group(ExprGroup {
            attrs,
            group_token,
            expr,
        }) => Expr::Group(ExprGroup {
            attrs,
            group_token,
            expr: Box::new(strip_wildcard(*expr)),
        }),
        // type ascription syntax is experimental so we always strip it
        Expr::Type(ExprType { expr, .. }) => *expr,
        // we want to retain casts if they semantically matter
        Expr::Cast(ExprCast {
            attrs,
            expr,
            as_token,
            ty,
        }) => match *ty {
            // cast to wildcard `_` will produce weird errors; we interpret it as taking the value as-is
            Type::Infer(_) => *expr,
            _ => Expr::Cast(ExprCast {
                attrs,
                expr,
                as_token,
                ty,
            }),
        },
        _ => expr,
    }
}