mysql_roaring_macros/
lib.rs

1use proc_macro::TokenStream;
2
3use quote::{format_ident, quote};
4
5fn uc_first(s: &str) -> String {
6    let mut c = s.chars();
7    match c.next() {
8        None => String::new(),
9        Some(f) => f.to_uppercase().collect::<String>() + c.as_str(),
10    }
11}
12
13
14#[proc_macro]
15pub fn create(_input: TokenStream) -> TokenStream {
16    let mut output = quote!();
17    for bit_size in vec!["32", "64"] {
18        for nullsafe in vec![true, false] {
19            let struct_name = format_ident!("Roaring{}{}Create",  bit_size,  if nullsafe {"Nullsafe"} else {""});
20            let map_type_name = format_ident!("Map{}", bit_size);
21            let return_type = quote!(&'a [u8]);
22            let return_type = if nullsafe { quote!(#return_type) } else { quote!(Option<#return_type>) };
23            let maybe_null_cnf = if nullsafe { quote!(_cfg) } else { quote! {cfg} };
24            let maybe_null = if nullsafe { quote!() } else { quote! {cfg.set_maybe_null(true);} };
25            let operation = if bit_size == "64" {
26                quote! {
27                    let value = value as u64;
28                    map.insert(value);
29                }
30            } else {
31                quote! {
32                    if let Some(value) = value.to_u32() {
33                        map.insert(value);
34                    } else {
35                        return Err(ProcessError);
36                    }
37                }
38            };
39            let null_result = if nullsafe {
40                quote! {
41                    let map = <#map_type_name>::new();
42                    self.vec.reserve(map.serialized_size());
43                    if let Ok(_) = map.serialize_into(&mut self.vec) {
44                        Ok(&self.vec[0..self.vec.len()])
45                    } else {
46                        Err(ProcessError)
47                    }
48                }
49            } else {
50                quote!(Ok(None))
51            };
52            let result = quote!(&self.vec[0..self.vec.len()]);
53            let result = if nullsafe { quote!(Ok(#result)) } else { quote!(Ok(Some(#result))) };
54            output.extend(quote! {
55                #[derive(Default)]
56                struct #struct_name {
57                    vec: Vec<u8>
58                }
59
60                #[register]
61                impl BasicUdf for #struct_name {
62                    type Returns<'a> = #return_type;
63
64                    fn init(#maybe_null_cnf: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
65                        #maybe_null
66                        if args.len() < 1 {
67                            return Err(format!("Expected one or more arguments; Got {} arguments.", args.len()));
68                        }
69                        for index in 0..args.len() {
70                            if !args.get(index).unwrap().value().is_int() {
71                                return Err(format!("{} argument mast be INTEGER or NULL.", index));
72                            }
73                        }
74                        Ok(Self::default())
75                    }
76
77                    fn process<'a>(&'a mut self, _cfg: &UdfCfg<Process>, args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<Self::Returns<'a>, ProcessError> {
78                        let mut map = <#map_type_name>::new();
79                        for arg in args.iter() {
80                            if let Some(value) = arg.value().as_int() {
81                                #operation
82                            } else {
83                                return {#null_result};
84                            }
85                        }
86                        self.vec.reserve(map.serialized_size());
87                        if let Ok(_) = map.serialize_into(&mut self.vec) {
88                            #result
89                        } else {
90                            Err(ProcessError)
91                        }
92                    }
93                }
94            });
95        }
96    }
97    output.into()
98}
99
100#[proc_macro]
101pub fn map_int_op(_input: TokenStream) -> TokenStream {
102    let mut output = quote!();
103    for bit_size in vec!["32", "64"] {
104        for nullsafe in vec![true, false] {
105            for op in vec!["insert", "remove"] {
106                let struct_name = format_ident!("Roaring{}{}{}",  bit_size,  if nullsafe {"Nullsafe"} else {""}, uc_first(op));
107                let map_type_name = format_ident!("Map{}", bit_size);
108                let op_fn_name = format_ident!("{}", op);
109                let return_type = quote!(&'a [u8]);
110                let return_type = if nullsafe { quote!(#return_type) } else { quote!(Option<#return_type>) };
111                let maybe_null_cnf = if nullsafe { quote!(_cfg) } else { quote! {cfg} };
112                let maybe_null = if nullsafe { quote!() } else { quote! {cfg.set_maybe_null(true);} };
113                let operation = quote! {map.#op_fn_name(value)};
114                let operation = if bit_size == "64" {
115                    quote! {
116                        let value = value as u64;
117                        #operation;
118                    }
119                } else {
120                    quote! {
121                        if let Some(value) = value.to_u32() {
122                            #operation;
123                        } else {
124                            return Err(ProcessError);
125                        }
126                    }
127                };
128                let null_result = if nullsafe {
129                    quote! {
130                        let map = <#map_type_name>::new();
131                        self.vec.reserve(map.serialized_size());
132                        if let Ok(_) = map.serialize_into(&mut self.vec) {
133                            Ok(&self.vec[0..self.vec.len()])
134                        } else {
135                            Err(ProcessError)
136                        }
137                    }
138                } else {
139                    quote!(Ok(None))
140                };
141                let result = quote!(&self.vec[0..self.vec.len()]);
142                let result = if nullsafe { quote!(Ok(#result)) } else { quote!(Ok(Some(#result))) };
143                output.extend(quote! {
144                    #[derive(Default)]
145                    struct #struct_name {
146                        vec: Vec<u8>
147                    }
148
149                    #[register]
150                    impl BasicUdf for #struct_name {
151                        type Returns<'a> = #return_type;
152
153                        fn init(#maybe_null_cnf: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
154                            #maybe_null
155                            if args.len() < 2 {
156                                return Err(format!("Expected two or more arguments; Got {} arguments.", args.len()));
157                            }
158                            if !args.get(0).unwrap().value().is_string() {
159                                return Err(format!("First argument mast be CLOB (bitmap) or NULL."));
160                            }
161                            for index in 1..args.len() {
162                                if !args.get(index).unwrap().value().is_int() {
163                                    return Err(format!("{} argument mast be INTEGER or NULL.", index));
164                                }
165                            }
166                            Ok(Self::default())
167                        }
168
169                        fn process<'a>(&'a mut self, _cfg: &UdfCfg<Process>, args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<Self::Returns<'a>, ProcessError> {
170                            let arg = args.get(0).unwrap().value();
171                            let bytes = arg.as_bytes();
172                            if let Some(bytes) = bytes {
173                                if let Ok(mut map) = <#map_type_name>::deserialize_from(bytes) {
174                                    for index in 1..args.len() {
175                                        let arg = args.get(index).unwrap();
176                                        if let Some(value) = arg.value().as_int() {
177                                            {#operation};
178                                        } else {
179                                            return {#null_result};
180                                        }
181                                    }
182                                    self.vec.reserve(map.serialized_size());
183                                    if let Ok(_) = map.serialize_into(&mut self.vec) {
184                                        return {#result};
185                                    } else {
186                                        return Err(ProcessError);
187                                    }
188                                } else {
189                                    return Err(ProcessError);
190                                }
191                            } else {
192                                return {#null_result};
193                            }
194                        }
195                    }
196                });
197            }
198        }
199    }
200    output.into()
201}
202
203#[proc_macro]
204pub fn map_op(_input: TokenStream) -> TokenStream {
205    let mut output = quote!();
206    for bit_size in vec!["32", "64"] {
207        for nullsafe in vec![true, false] {
208            for count in vec![true, false] {
209                for op in vec!["or", "xor", "and"] {
210                    let struct_name = format_ident!("Roaring{}{}{}{}",  bit_size,  if nullsafe {"Nullsafe"} else {""}, uc_first(op), if count {"Count"} else {""});
211                    let struct_fields = if count { quote!() } else { quote!(vec: Vec<u8>) };
212                    let map_type_name = format_ident!("Map{}", bit_size);
213                    let op_fn_name = format_ident!("bit{}_assign", op);
214                    let return_type = if count { quote!(i64) } else { quote!(&'a [u8]) };
215                    let return_type = if nullsafe { quote!(#return_type) } else { quote!(Option<#return_type>) };
216                    let maybe_null_cnf = if nullsafe { quote!(_cfg) } else { quote! {cfg} };
217                    let maybe_null = if nullsafe { quote!() } else { quote! {cfg.set_maybe_null(true);} };
218                    let on_null_arg = if nullsafe { quote!(()) } else { quote!(return Ok(None)) };
219                    let on_null_agg = if nullsafe { quote!(<#map_type_name>::new()) } else { quote!(return Ok(None)) };
220                    let bytes_result = quote!(&self.vec[0..self.vec.len()]);
221                    let bytes_result = if nullsafe { quote!(#bytes_result) } else { quote!(Some(#bytes_result)) };
222                    let count_result = quote!(agg.len() as i64);
223                    let count_result = if nullsafe { quote!(#count_result) } else { quote!(Some(#count_result)) };
224                    let count_null_result = if nullsafe { quote!(0i64) } else { quote!(None) };
225                    let result = if count {
226                        quote! {
227                        match maybe_agg {
228                            Some(agg) => Ok(#count_result),
229                            None => Ok(#count_null_result)
230                        }
231                      }
232                    } else {
233                        quote! {
234                        let agg = match maybe_agg {
235                            Some(agg) => agg,
236                            None => #on_null_agg,
237                        };
238                        self.vec.reserve(agg.serialized_size());
239                        if let Ok(_) = agg.serialize_into(&mut self.vec) {
240                            Ok(#bytes_result)
241                        } else {
242                            Err(ProcessError)
243                        }
244                      }
245                    };
246                    // let null_result = if nullsafe {quote!(bytes)} else {quote!(None)};
247                    // let null_result = if nullsafe {quote!(bytes)} else {quote!(None)};
248                    output.extend(quote! {
249                        #[derive(Default)]
250                        struct #struct_name {
251                            #struct_fields
252                        }
253
254                        #[register]
255                        impl BasicUdf for #struct_name {
256                            type Returns<'a> = #return_type;
257
258                            fn init(#maybe_null_cnf: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
259                                #maybe_null
260                                if args.len() < 2 {
261                                    return Err(format!("Expected two or more CLOB (bitmap) or NULL arguments; Got {} arguments.", args.len()));
262                                }
263                                for arg in args.iter() {
264                                    if let SqlResult::String(_) = arg.value() {} else {
265                                        return Err(format!("Only CLOB (bitmap) or NULL arguments are allowed."));
266                                    }
267                                }
268                                Ok(Self::default())
269                            }
270
271                            fn process<'a>(&'a mut self, _cfg: &UdfCfg<Process>, args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<Self::Returns<'a>, ProcessError> {
272                                let mut maybe_agg: Option<#map_type_name> = None;
273                                for arg in args.iter() {
274                                    match arg.value().as_bytes() {
275                                        None => #on_null_arg,
276                                        Some(bytes) => {
277                                            if let Ok(map) = <#map_type_name>::deserialize_from(bytes) {
278                                                match maybe_agg {
279                                                    None => maybe_agg = Some(map),
280                                                    Some(ref mut agg) => agg.#op_fn_name(map),
281                                                }
282                                            } else {
283                                                return Err(ProcessError);
284                                            }
285                                        }
286                                    }
287                                }
288                                #result
289                            }
290                        }
291                    })
292                }
293            }
294        }
295    }
296    output.into()
297}
298
299#[proc_macro]
300pub fn contains(_input: TokenStream) -> TokenStream {
301    let mut output = quote!();
302    for bit_size in vec!["32", "64"] {
303        for nullsafe in vec![true, false] {
304            let struct_name = format_ident!("Roaring{}{}Contains",  bit_size,  if nullsafe {"Nullsafe"} else {""});
305            let cast_function_name = format_ident!("to_u{}", bit_size);
306            let map_type_name = format_ident!("Map{}", bit_size);
307            let return_type = if nullsafe { quote!(i64) } else { quote!(Option<i64>) };
308            let null_result = if nullsafe { quote!(0i64) } else { quote!(None) };
309            let result = quote!(map.contains(value).into());
310            let result = if nullsafe { quote!(#result) } else { quote!(Some(#result)) };
311            let maybe_null_cnf = if nullsafe { quote!(_cfg) } else { quote! {cfg} };
312            let maybe_null = if nullsafe { quote!() } else { quote! {cfg.set_maybe_null(true);} };
313            output.extend(quote! {
314                #[derive(Default)]
315                struct #struct_name;
316
317                #[register]
318                impl BasicUdf for #struct_name {
319                    type Returns<'a> = #return_type;
320
321                    fn init(#maybe_null_cnf: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
322                        if args.len() < 2 {
323                            return Err(format!("Expected two arguments; Got {} arguments.", args.len()));
324                        }
325                        if !args.get(0).unwrap().value().is_string() {
326                            return Err(format!("First argument mast be CLOB (bitmap) or NULL."));
327                        }
328                        if !args.get(1).unwrap().value().is_int() {
329                            return Err(format!("Second argument mast be INT or NULL."));
330                        }
331                        #maybe_null
332                        Ok(Self::default())
333                    }
334
335                    fn process<'a>(&'a mut self, _cfg: &UdfCfg<Process>, args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<Self::Returns<'a>, ProcessError> {
336                        let arg0 = args.get(0).unwrap().value();
337                        let arg1 = args.get(1).unwrap().value();
338                        let bytes = arg0.as_bytes();
339                        let value = arg1.as_int();
340                        if let (Some(bytes), Some(value)) = (bytes, value) {
341                            if let (Ok(map), Some(value)) = (#map_type_name::deserialize_from(bytes), value.#cast_function_name()) {
342                                Ok(#result)
343                            } else {
344                                Err(ProcessError)
345                            }
346                        } else {
347                            Ok(#null_result)
348                        }
349                    }
350                }
351            });
352        }
353    }
354    output.into()
355}
356
357#[proc_macro]
358pub fn count(_input: TokenStream) -> TokenStream {
359    let mut output = quote!();
360    for bit_size in vec!["32", "64"] {
361        for nullsafe in vec![true, false] {
362            let struct_name = format_ident!("Roaring{}{}Count",  bit_size,  if nullsafe {"Nullsafe"} else {""});
363            let map_type_name = format_ident!("Map{}", bit_size);
364            let return_type = if nullsafe { quote!(i64) } else { quote!(Option<i64>) };
365            let null_result = if nullsafe { quote!(0i64) } else { quote!(None) };
366            let result = quote!(map.len() as i64);
367            let result = if nullsafe { quote!(#result) } else { quote!(Some(#result)) };
368            let maybe_null_cnf = if nullsafe { quote!(_cfg) } else { quote! {cfg} };
369            let maybe_null = if nullsafe { quote!() } else { quote! {cfg.set_maybe_null(true);} };
370            output.extend(quote! {
371                #[derive(Default)]
372                struct #struct_name;
373
374                #[register]
375                impl BasicUdf for #struct_name {
376                    type Returns<'a> = #return_type;
377
378                    fn init(#maybe_null_cnf: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
379                        #maybe_null
380                        if args.len() != 1 {
381                            return Err(format!("Expected one arguments; Got {} arguments.", args.len()));
382                        }
383                        if !args.get(0).unwrap().value().is_string() {
384                            return Err(format!("First argument mast be CLOB (bitmap) or NULL."));
385                        }
386                        Ok(Self::default())
387                    }
388
389                    fn process<'a>(&'a mut self, _cfg: &UdfCfg<Process>, args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<Self::Returns<'a>, ProcessError> {
390                        let arg = args.get(0).unwrap().value();
391                        let bytes = arg.as_bytes();
392                        if let Some(bytes) = bytes {
393                            if let Ok(map) = <#map_type_name>::deserialize_from(bytes) {
394                                Ok(#result)
395                            } else {
396                                Err(ProcessError)
397                            }
398                        } else {
399                            Ok(#null_result)
400                        }
401                    }
402                }
403            });
404        }
405    }
406    output.into()
407}
408
409#[proc_macro]
410pub fn json(_input: TokenStream) -> TokenStream {
411    let mut output = quote!();
412    for bit_size in vec!["32", "64"] {
413        for nullsafe in vec![true, false] {
414            let struct_name = format_ident!("Roaring{}{}Json",  bit_size,  if nullsafe {"Nullsafe"} else {""});
415            let map_type_name = format_ident!("Map{}", bit_size);
416            let return_type = if nullsafe { quote!(&'a [u8]) } else { quote!(Option<&'a [u8]>) };
417            let null_result = if nullsafe {
418                quote! {
419                    self.vec.extend(b"[]");
420                    Ok(&self.vec[0..self.vec.len()])
421                }
422            } else {
423                quote!(Ok(None))
424            };
425            let result = quote!(&self.vec[0..self.vec.len()]);
426            let result = if nullsafe { quote!(#result) } else { quote!(Some(#result)) };
427            let maybe_null_cnf = if nullsafe { quote!(_cfg) } else { quote! {cfg} };
428            let maybe_null = if nullsafe { quote!() } else { quote! {cfg.set_maybe_null(true);} };
429            output.extend(quote! {
430                #[derive(Default)]
431                struct #struct_name {
432                    vec: Vec<u8>
433                }
434
435                #[register]
436                impl BasicUdf for #struct_name {
437                    type Returns<'a> = #return_type;
438
439                    fn init(#maybe_null_cnf: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
440                        #maybe_null
441                        if args.len() != 1 {
442                            return Err(format!("Expected one arguments; Got {} arguments.", args.len()));
443                        }
444                        if !args.get(0).unwrap().value().is_string() {
445                            return Err(format!("First argument mast be CLOB (bitmap) or NULL."));
446                        }
447                        Ok(Self::default())
448                    }
449
450                    fn process<'a>(&'a mut self, _cfg: &UdfCfg<Process>, args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<Self::Returns<'a>, ProcessError> {
451                        let arg = args.get(0).unwrap().value();
452                        let bytes = arg.as_bytes();
453                        if let Some(bytes) = bytes {
454                            if let Ok(map) = <#map_type_name>::deserialize_from(bytes) {
455                                let mut count = map.len() as usize;
456                                let capacity = ((10usize+1)*count)+2; // "[" + (log10(2^32) + ",") * map.len() + "]";
457                                self.vec.reserve(capacity);
458                                const COMA: &[u8] = b",";
459                                self.vec.extend(b"[");
460                                for item in map.iter() {
461                                    self.vec.extend(item.to_string().as_bytes());
462                                    count-=1;
463                                    if count != 0 {
464                                        self.vec.extend(COMA);
465                                    }
466                                }
467                                self.vec.extend(b"]");
468                                Ok(#result)
469                            } else {
470                                Err(ProcessError)
471                            }
472                        } else {
473                            #null_result
474                        }
475                    }
476                }
477            });
478        }
479    }
480    output.into()
481}
482
483
484#[proc_macro]
485pub fn group_map_op(_input: TokenStream) -> TokenStream {
486    let mut output = quote!();
487    for bit_size in vec!["32", "64"] {
488        for op in vec!["or", "xor", "and"] {
489            let ucf_op = uc_first(op);
490            let struct_name = format_ident!("Roaring{}Group{}",  bit_size,  ucf_op);
491            let struct_nullsafe_name = format_ident!("Roaring{}NullsafeGroup{}",  bit_size,  ucf_op);
492            let struct_count_name = format_ident!("Roaring{}Group{}Count",  bit_size,  ucf_op);
493            let struct_nullsafe_count_name = format_ident!("Roaring{}NullsafeGroup{}Count",  bit_size,  ucf_op);
494            let map_type_name = format_ident!("Map{}", bit_size);
495            let op_fn_name = format_ident!("bit{}_assign", op);
496            let add = quote! {
497                fn add(&mut self, _cfg: &UdfCfg<Process>, args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<(), NonZeroU8> {
498                    if let Some(bytes) = args.get(0).unwrap().value().as_bytes() {
499                        if let Ok(map) = <#map_type_name>::deserialize_from(bytes) {
500                            if let Some(ref mut acc) = self.map {
501                                acc.#op_fn_name(map);
502                                Ok(())
503                            } else {
504                                self.map = Some(map);
505                                Ok(())
506                            }
507                        } else {
508                            Err(NonZeroU8::new(1u8).unwrap())
509                        }
510                    } else {
511                        Ok(())
512                    }
513                }
514            };
515            let init = quote! {
516                fn init(_cfg: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
517                    if args.len() != 1 {
518                        return Err(format!("Expected one arguments; Got {} arguments.", args.len()));
519                    }
520                    if !args.get(0).unwrap().value().is_string() {
521                        return Err(format!("{} argument mast be BLOB or NULL.", 0));
522                    }
523                    Ok(Self::default())
524                }
525            };
526            let init_nullable = quote! {
527                fn init(cfg: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
528                    cfg.set_maybe_null(true);
529                    if args.len() != 1 {
530                        return Err(format!("Expected one arguments; Got {} arguments.", args.len()));
531                    }
532                    if !args.get(0).unwrap().value().is_string() {
533                        return Err(format!("{} argument mast be BLOB or NULL.", 0));
534                    }
535                    Ok(Self::default())
536                }
537            };
538            let clear = quote! {
539                fn clear(&mut self, _cfg: &UdfCfg<Process>, _error: Option<NonZeroU8>) -> Result<(), NonZeroU8> {
540                    self.map = None;
541                    self.vec.clear();
542                    Ok(())
543                }
544            };
545            let clear_count = quote! {
546                fn clear(&mut self, _cfg: &UdfCfg<Process>, _error: Option<NonZeroU8>) -> Result<(), NonZeroU8> {
547                    self.map = None;
548                    Ok(())
549                }
550            };
551            let fields = quote! {
552                map: Option<#map_type_name>,
553                vec: Vec<u8>,
554            };
555            let fields_count = quote! {
556                map: Option<#map_type_name>,
557            };
558            output.extend(quote! {
559                #[derive(Default)]
560                struct #struct_name {
561                    #fields
562                }
563
564                #[register]
565                impl AggregateUdf for #struct_name {
566                    #clear
567                    #add
568                }
569
570                #[register]
571                impl BasicUdf for #struct_name {
572                    type Returns<'a> = Option<&'a [u8]>;
573
574                    #init_nullable
575
576                    fn process<'a>(&'a mut self, _cfg: &UdfCfg<Process>, _args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<Self::Returns<'a>, ProcessError> {
577                        if let Some(ref agg) = self.map {
578                            let capacity = agg.serialized_size();
579                            if self.vec.capacity() < capacity {
580                                self.vec.reserve(capacity - self.vec.capacity());
581                            }
582                            if let Ok(_) = agg.serialize_into(&mut self.vec) {
583                                Ok(Some(&self.vec[0..self.vec.len()]))
584                            } else {
585                                Err(ProcessError)
586                            }
587                        } else {
588                            Ok(None)
589                        }
590                    }
591                }
592
593                #[derive(Default)]
594                struct #struct_count_name {
595                    #fields_count
596                }
597
598                #[register]
599                impl AggregateUdf for #struct_count_name {
600                    #clear_count
601                    #add
602                }
603
604                #[register]
605                impl BasicUdf for #struct_count_name {
606                    type Returns<'a> = Option<i64>;
607
608                    #init_nullable
609
610                    fn process<'a>(&'a mut self, _cfg: &UdfCfg<Process>, _args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<Self::Returns<'a>, ProcessError> {
611                        if let Some(ref agg) = self.map {
612                            Ok(Some(agg.len() as i64))
613                        } else {
614                            Ok(None)
615                        }
616                    }
617                }
618
619                #[derive(Default)]
620                struct #struct_nullsafe_name {
621                    #fields
622                }
623
624                #[register]
625                impl AggregateUdf for #struct_nullsafe_name {
626                    #clear
627                    #add
628                }
629
630                #[register]
631                impl BasicUdf for #struct_nullsafe_name {
632                    type Returns<'a> = &'a [u8];
633
634                    #init
635
636                    fn process<'a>(&'a mut self, _cfg: &UdfCfg<Process>, _args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<Self::Returns<'a>, ProcessError> {
637                        let map;
638                        let null;
639                        if let Some(ref agg) = self.map {
640                            map = agg;
641                        } else {
642                            null = <#map_type_name>::new();
643                            map = &null;
644                        }
645                        let capacity = map.serialized_size();
646                        if self.vec.capacity() < capacity {
647                            self.vec.reserve(capacity - self.vec.capacity());
648                        }
649                        if let Ok(_) = map.serialize_into(&mut self.vec) {
650                            Ok(&self.vec[0..self.vec.len()])
651                        } else {
652                            Err(ProcessError)
653                        }
654                    }
655                }
656
657                #[derive(Default)]
658                struct #struct_nullsafe_count_name {
659                    #fields_count
660                }
661
662                #[register]
663                impl AggregateUdf for #struct_nullsafe_count_name {
664                    #clear_count
665                    #add
666                }
667
668                #[register]
669                impl BasicUdf for #struct_nullsafe_count_name {
670                    type Returns<'a> = i64;
671
672                    #init
673
674                    fn process<'a>(&'a mut self, _cfg: &UdfCfg<Process>, _args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<Self::Returns<'a>, ProcessError> {
675                        if let Some(ref agg) = self.map {
676                            Ok(agg.len() as i64)
677                        } else {
678                            Ok(0i64)
679                        }
680                    }
681                }
682            })
683        }
684    }
685    output.into()
686}
687
688#[proc_macro]
689pub fn group_create(_input: TokenStream) -> TokenStream {
690    let mut output = quote!();
691    for bit_size in vec!["32", "64"] {
692        for nullsafe in vec![true, false] {
693            let struct_name = format_ident!("Roaring{}{}GroupCreate",  bit_size,  if nullsafe {"Nullsafe"} else {""});
694            let map_type_name = format_ident!("Map{}", bit_size);
695            let return_type = quote!(&'a [u8]);
696            let return_type = if nullsafe { quote!(#return_type) } else { quote!(Option<#return_type>) };
697            let maybe_null_cnf = if nullsafe { quote!(_cfg) } else { quote! {cfg} };
698            let maybe_null = if nullsafe { quote!() } else { quote! {cfg.set_maybe_null(true);} };
699            let bytes_result = quote!(&self.vec[0..self.vec.len()]);
700            let bytes_result = if nullsafe { quote!(#bytes_result) } else { quote!(Some(#bytes_result)) };
701            let extend_capacity = quote! {
702                let capacity = self.map.serialized_size();
703                if self.vec.capacity() < capacity {
704                    self.vec.reserve(capacity - self.vec.capacity());
705                }
706            };
707            let null_bytes_result = if nullsafe { quote!(&self.vec[0..self.vec.len()]) } else { quote!(None) };
708            let null_result = quote! {
709                #extend_capacity
710                if let Ok(_) = self.map.serialize_into(&mut self.vec) {
711                    Ok(#null_bytes_result)
712                } else {
713                    Err(ProcessError)
714                }
715            };
716            let result = quote! {
717                #extend_capacity
718                if let Ok(_) = self.map.serialize_into(&mut self.vec) {
719                    Ok(#bytes_result)
720                } else {
721                    Err(ProcessError)
722                }
723            };
724
725            let mut operations = quote!();
726            for (func, op) in vec![("add", "insert"), ("remove", "remove")].iter() {
727                let func = format_ident!("{}", func);
728                let op = format_ident!("{}", op);
729                let operation = if bit_size == "64" {
730                    quote! {
731                        self.map.#op(value as u64);
732                        Ok(())
733                    }
734                } else {
735                    quote! {
736                        if let Some(value) = value.to_u32() {
737                            self.map.#op(value);
738                            Ok(())
739                        } else {
740                            Err(NonZeroU8::new(1u8).unwrap())
741                        }
742                    }
743                };
744                operations.extend(quote! {
745                    fn #func(&mut self, _cfg: &UdfCfg<Process>, args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<(), NonZeroU8> {
746                        if let Some(value) = args.get(0).unwrap().value().as_int() {
747                            #operation
748                        } else {
749                            Ok(())
750                        }
751                    }
752                });
753            }
754            output.extend(quote! {
755                #[derive(Default)]
756                struct #struct_name {
757                    map: #map_type_name,
758                    vec: Vec<u8>,
759                }
760
761                #[register]
762                impl AggregateUdf for #struct_name {
763                    fn clear(&mut self, _cfg: &UdfCfg<Process>, _error: Option<NonZeroU8>) -> Result<(), NonZeroU8> {
764                        self.map.clear();
765                        self.vec.clear();
766                        Ok(())
767                    }
768                    #operations
769                }
770
771                #[register]
772                impl BasicUdf for #struct_name {
773                    type Returns<'a> = #return_type;
774
775                    fn init(#maybe_null_cnf: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
776                        #maybe_null
777                        if args.len() != 1 {
778                            return Err(format!("Expected one arguments; Got {} arguments.", args.len()));
779                        }
780                        if !args.get(0).unwrap().value().is_int() {
781                            return Err(format!("{} argument mast be INTEGER or NULL.", 0));
782                        }
783                        Ok(Self::default())
784                    }
785
786                    fn process<'a>(&'a mut self, _cfg: &UdfCfg<Process>, _args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<Self::Returns<'a>, ProcessError> {
787                        if(self.map.len() > 0) {
788                            #result
789                        } else {
790                            #null_result
791                        }
792                    }
793                }
794            })
795        }
796    }
797    output.into()
798}