bench_utils/
utils.rs

1use std::str::FromStr;
2
3use if_chain::if_chain;
4use proc_macro::TokenStream;
5use proc_macro2::{Ident, Span, TokenStream as TokenStream2};
6use quote::{ToTokens, quote};
7
8pub fn parse_expr<T: FromStr>(expr: syn::Expr) -> T {
9  expr.to_token_stream().to_string().parse().ok().unwrap()
10}
11
12pub struct ArrayAssign {
13  pub var_name: syn::Expr,
14  pub var_val: syn::Expr,
15  pub num_locations: usize,
16}
17
18/// Parses the expression:
19/// `var_name: [_; num_locations] = var_val`
20pub fn parse_arr_assign(input: TokenStream) -> ArrayAssign {
21  let expr: syn::ExprAssign = syn::parse(input).unwrap();
22
23  if_chain! {
24    if let syn::Expr::Type(expr_ty) = *expr.left;
25    if let syn::Type::Array(arr_ty) = *expr_ty.ty;
26    then {
27      return ArrayAssign {
28          var_name: *expr_ty.expr,
29          var_val: *expr.right,
30          num_locations: parse_expr(arr_ty.len),
31      };
32    }
33  }
34
35  panic!();
36}
37
38pub struct StructAssign {
39  pub var_name: syn::Expr,
40  pub struct_name: proc_macro2::Ident,
41  pub field_val: syn::Expr,
42  pub field_ty: syn::Type,
43  pub num_fields: usize,
44}
45
46/// Parses the expression:
47/// `var_name: struct_name<[field_ty; num_fields]> = field_val`
48pub fn parse_struct_assign(input: TokenStream) -> StructAssign {
49  let expr: syn::ExprAssign = syn::parse(input).unwrap();
50
51  if_chain! {
52    if let syn::Expr::Type(expr_ty) = *expr.left;
53
54    if let syn::Type::Path(struct_ty) = *expr_ty.ty;
55    let ty_segments = struct_ty.path.segments.first().unwrap();
56
57    if let syn::PathArguments::AngleBracketed(generics) = &ty_segments.arguments;
58    let arr_generic = generics.args.first().unwrap();
59
60    if let syn::GenericArgument::Type(syn::Type::Array(arr_ty)) = arr_generic;
61
62    then {
63      return StructAssign {
64        var_name: *expr_ty.expr,
65        struct_name: ty_segments.ident.clone(),
66        field_ty: *arr_ty.elem.clone(),
67        field_val: *expr.right,
68        num_fields: parse_expr(arr_ty.len.clone()),
69      };
70    }
71  }
72
73  panic!();
74}
75
76/// Represents a "level" of the nested struct
77#[derive(Clone)]
78pub struct TreeLevel {
79  ident: Ident,
80  pub def: TokenStream2,
81  instance_ident: Ident,
82  pub instantiation: TokenStream2,
83}
84
85impl TreeLevel {
86  pub fn new(
87    level: usize,
88    fields: Vec<Ident>,
89    child_level: Option<&TreeLevel>,
90    field_val: &syn::Expr,
91    field_ty: &syn::Type,
92  ) -> TreeLevel {
93    let ident = Ident::new(&format!("struct_{level}"), Span::call_site());
94    let instance_ident = Ident::new(&format!("struct_{level}_inst"), Span::call_site());
95    let type_ident =
96      Ident::new(&field_ty.to_token_stream().to_string(), Span::call_site());
97
98    // If a child level exists, fields of the current level should
99    // have the type of the child struct (if not, fall back to primitive)
100    let field_ty = if let Some(child) = child_level {
101      &child.ident
102    } else {
103      &type_ident
104    };
105
106    let def = quote! {
107      #[derive(Clone)]
108      struct #ident {
109        #(#fields: #field_ty,)*
110      }
111    };
112
113    let instantiation = match child_level {
114      Some(child) => {
115        let field_val = &child.instance_ident;
116
117        quote! {
118          let #instance_ident = #ident {
119            #(#fields: #field_val.clone(),)*
120          };
121        }
122      }
123      None => {
124        quote! {
125          let #instance_ident = #ident {
126            #(#fields: #field_val,)*
127          };
128        }
129      }
130    };
131
132    TreeLevel {
133      ident,
134      def,
135      instance_ident,
136      instantiation,
137    }
138  }
139}