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

use proc_macro::TokenStream;
use quote::quote;
use syn::parse::{Parse, ParseStream};

// Convert rust code to a field in a names-only schema
// Example Input:
//  a: { b: {}, c: { d: {} } }
//
// Example Output:
//  NamesOnlySchemaNode { name: "a".to_string(), children: vec![
//    NamesOnlySchemaNode { name: "b".to_string(), children: vec![] },
//    NamesOnlySchemaNode { name: "c".to_string(), children: vec![
//      NamesOnlySchemaNode { name: "d".to_string(), children: vec![] }
//    ]}
//  ]}
fn rust_field_to_names_field(field: &Field) -> proc_macro2::TokenStream {
    let name = field.name.to_string();
    let children = rust_to_names_fields(&field.ty);
    quote! {substrait_expr::helpers::schema::NamesOnlySchemaNode {
        name: #name.to_string(),
        children: #children,
    }}
}

// Convert rust code to a vector of NamesOnlySchemaNode
//
// Example input:
//  { a: { b: {}, c: { d: {} } } }
//
// Example Output:
// vec![
//   NamesOnlySchemaNode { name: "a".to_string(), children: vec![
//     NamesOnlySchemaNode { name: "b".to_string(), children: vec![] },
//     NamesOnlySchemaNode { name: "c".to_string(), children: vec![
//       NamesOnlySchemaNode { name: "d".to_string(), children: vec![] }
//     ]}
//   ]}
// ]
fn rust_to_names_fields(schema: &NestedType) -> proc_macro2::TokenStream {
    let parsed_fields = schema
        .fields
        .iter()
        .map(|field| rust_field_to_names_field(field))
        .collect::<Vec<_>>();
    quote! {vec![#(#parsed_fields),*]}
}

// Convert rust code to a root NamesOnlySchemaNode (that has the empty string for a name)
//
// Example input:
//  { a: { b: {}, c: { d: {} } } }
//
// Example Output:
// NamesOnlySchemaNode {
//   name: "".to_string(),
//   vec![
//     NamesOnlySchemaNode { name: "a".to_string(), children: vec![
//       NamesOnlySchemaNode { name: "b".to_string(), children: vec![] },
//       NamesOnlySchemaNode { name: "c".to_string(), children: vec![
//         NamesOnlySchemaNode { name: "d".to_string(), children: vec![] }
//       ]}
//     ]}
//   ]
// }
fn rust_to_names_schema(schema: &NestedType) -> proc_macro2::TokenStream {
    let children = rust_to_names_fields(schema);
    quote! {
        substrait_expr::helpers::schema::SchemaInfo::Names(substrait_expr::helpers::schema::NamesOnlySchema::new(#children))
    }
}

// New rust syntax for a field in a names only schema
//
// Examples:
//  foo: {}
//  blah: { x: {}, y: { z: {} } }
struct Field {
    name: syn::Ident,
    _colon_token: syn::Token![:],
    ty: NestedType,
}

impl Parse for Field {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        Ok(Field {
            name: input.parse()?,
            _colon_token: input.parse()?,
            ty: input.parse()?,
        })
    }
}

// New rust syntax for a nested names-only type ({field, field, field})
//
// Examples:
//  { foo: {} }
//  { blah: { x: {}, y: { z: {} } } }
struct NestedType {
    _brace_token: syn::token::Brace,
    fields: syn::punctuated::Punctuated<Field, syn::Token![,]>,
}

impl Parse for NestedType {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let content;
        Ok(Self {
            _brace_token: syn::braced!(content in input),
            fields: content.parse_terminated(Field::parse, syn::Token![,])?,
        })
    }
}

fn names_schema2(input: proc_macro2::TokenStream) -> syn::Result<proc_macro2::TokenStream> {
    let struct_type: NestedType = syn::parse2(input)?;
    Ok(rust_to_names_schema(&struct_type))
}

/// A macro to create names-only schemas from a dictionary-like rust syntax
///
/// # Examples
/// ```ignore
/// use substrait_expr::macros::names_schema;
///
/// let schema = names_schema!({
///   vector: {},
///   metadata: {
///     caption: {},
///     user_score: {}
///  }
/// });
/// ```
#[proc_macro]
pub fn names_schema(input: TokenStream) -> TokenStream {
    let input = proc_macro2::TokenStream::from(input);
    names_schema2(input).unwrap().into()
}