3929938005

Build Type

push

github

Committed by GitHub

Commit Message

Merge #713

Run Details

84930 of 96741 relevant lines covered (87.79%)

79640.1 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

89.13

/operators/src/processing/expression/codegen.rs

use std::{collections::BTreeSet, fmt::Debug};

use proc_macro2::TokenStream;
use quote::{format_ident, quote, ToTokens};
use snafu::ensure;

use super::{
    error::{self, ExpressionError},
    functions::FUNCTIONS,
};

type Result<T, E = ExpressionError> = std::result::Result<T, E>;

// TODO: prefix for variables and functions

/// An expression as an abstract syntax tree.
/// Allows genering Rust code.
#[derive(Debug, Clone)]
pub struct ExpressionAst {
    /// This name is the generated function name after generating code.
    name: Identifier,
    root: AstNode,
    parameters: Vec<Parameter>,
    functions: BTreeSet<AstFunction>,
    // TODO: dtype Float or Int
}

impl ExpressionAst {
    pub fn new(
        name: Identifier,
        parameters: Vec<Parameter>,
        functions: BTreeSet<AstFunction>,
        root: AstNode,
    ) -> Result<ExpressionAst> {
        ensure!(!name.as_ref().is_empty(), error::EmptyExpressionName);

        Ok(Self {
            name,
            root,
            parameters,
            functions,
        })
    }

    pub fn code(&self) -> String {
        self.to_token_stream().to_string()
    }

    pub fn name(&self) -> &str {
        self.name.as_ref()
    }
}

impl ToTokens for ExpressionAst {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let dtype = format_ident!("{}", "f64");

        for function in &self.functions {
            function.to_tokens(tokens);
        }

        let fn_name = &self.name;
        let params: Vec<TokenStream> = self
            .parameters
            .iter()
            .map(|p| match p {
                Parameter::Number(param) => quote! { #param: Option<#dtype> },
            })
            .collect();
        let content = &self.root;

        tokens.extend(quote! {
            #[inline]
            fn apply(a: Option<#dtype>, b: Option<#dtype>, f: fn(#dtype, #dtype) -> #dtype) -> Option<#dtype> {
                match (a, b) {
                    (Some(a), Some(b)) => Some(f(a, b)),
                    _ => None,
                }
            }

            #[no_mangle]
            pub extern "Rust" fn #fn_name (#(#params),*) -> Option<#dtype> {
                #content
            }
        });
    }
}

#[derive(Debug, Clone)]
pub enum AstNode {
    Constant(f64),
    NoData,
    Variable(Identifier),
    Operation {
        left: Box<AstNode>,
        op: AstOperator,
        right: Box<AstNode>,
    },
    Function {
        name: Identifier,
        args: Vec<AstNode>,
    },
    Branch {
        condition_branches: Vec<Branch>,
        else_branch: Box<AstNode>,
    },
    AssignmentsAndExpression {
        assignments: Vec<Assignment>,
        expression: Box<AstNode>,
    },
}

impl ToTokens for AstNode {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let new_tokens = match self {
            Self::Constant(n) => quote! { Some(#n) },
            Self::NoData => quote! { None },
            Self::Variable(v) => quote! { #v },
            Self::Operation { left, op, right } => {
                quote! { apply(#left, #right, #op) }
            }
            Self::Function { name, args } => {
                let fn_name = format_ident!("import_{}__{}", name.as_ref(), args.len());
                quote! { #fn_name(#(#args),*) }
            }
            AstNode::Branch {
                condition_branches,
                else_branch: default_branch,
            } => {
                let mut new_tokens = TokenStream::new();
                for (i, branch) in condition_branches.iter().enumerate() {
                    let condition = &branch.condition;
                    let body = &branch.body;

                    new_tokens.extend(if i == 0 {
                        // first
                        quote! {
                            if #condition {
                                #body
                            }
                        }
                    } else {
                        // middle
                        quote! {
                            else if #condition {
                                #body
                            }
                        }
                    });
                }

                new_tokens.extend(quote! {
                    else {
                        #default_branch
                    }
                });

                new_tokens
            }
            Self::AssignmentsAndExpression {
                assignments,
                expression,
            } => {
                quote! {
                    #(#assignments)*
                    #expression
                }
            }
        };

        tokens.extend(new_tokens);
    }
}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(transparent)]
pub struct Identifier(String);

impl ToTokens for Identifier {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let identifier = format_ident!("{}", self.0);
        tokens.extend(quote! { #identifier });
    }
}

impl From<String> for Identifier {
    fn from(s: String) -> Self {
        Self(s)
    }
}

impl From<&str> for Identifier {
    fn from(s: &str) -> Self {
        Self(s.to_string())
    }
}

impl From<&String> for Identifier {
    fn from(s: &String) -> Self {
        Self(s.to_string())
    }
}

impl AsRef<str> for Identifier {
    fn as_ref(&self) -> &str {
        &self.0
    }
}

impl std::fmt::Display for Identifier {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
        std::fmt::Display::fmt(&self.0, f)
    }
}

#[derive(Debug, Clone)]
pub enum AstOperator {
    Add,
    Subtract,
    Multiply,
    Divide,
}

impl ToTokens for AstOperator {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let new_tokens = match self {
            AstOperator::Add => quote! { std::ops::Add::add },
            AstOperator::Subtract => quote! { std::ops::Sub::sub },
            AstOperator::Multiply => quote! { std::ops::Mul::mul },
            AstOperator::Divide => quote! { std::ops::Div::div },
        };

        tokens.extend(new_tokens);
    }
}

#[derive(Debug, Clone)]
pub struct Branch {
    pub condition: BooleanExpression,
    pub body: AstNode,
}

#[derive(Debug, Clone)]
pub enum BooleanExpression {
    Constant(bool),
    Comparison {
        left: Box<AstNode>,
        op: BooleanComparator,
        right: Box<AstNode>,
    },
    Operation {
        left: Box<BooleanExpression>,
        op: BooleanOperator,
        right: Box<BooleanExpression>,
    },
}

impl ToTokens for BooleanExpression {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let new_tokens = match self {
            Self::Constant(b) => quote! { #b },
            Self::Comparison { left, op, right } => quote! { ((#left) #op (#right)) },
            Self::Operation { left, op, right } => quote! { ( (#left) #op (#right) ) },
        };

        tokens.extend(new_tokens);
    }
}

#[derive(Debug, Clone)]
pub enum BooleanComparator {
    Equal,
    NotEqual,
    LessThan,
    LessThanOrEqual,
    GreaterThan,
    GreaterThanOrEqual,
}

impl ToTokens for BooleanComparator {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let new_tokens = match self {
            Self::Equal => quote! { == },
            Self::NotEqual => quote! { != },
            Self::LessThan => quote! { < },
            Self::LessThanOrEqual => quote! { <= },
            Self::GreaterThan => quote! { > },
            Self::GreaterThanOrEqual => quote! { >= },
        };

        tokens.extend(new_tokens);
    }
}

#[derive(Debug, Clone)]
pub enum BooleanOperator {
    And,
    Or,
}

impl ToTokens for BooleanOperator {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let new_tokens = match self {
            Self::And => quote! { && },
            Self::Or => quote! { || },
        };

        tokens.extend(new_tokens);
    }
}

#[derive(Debug, Clone)]
pub struct Assignment {
    pub identifier: Identifier,
    pub expression: AstNode,
}

impl ToTokens for Assignment {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let Self {
            identifier,
            expression,
        } = self;
        let new_tokens = quote! {
            let #identifier = #expression;
        };

        tokens.extend(new_tokens);
    }
}

#[derive(Debug, Clone)]
pub enum Parameter {
    Number(Identifier),
}

impl AsRef<str> for Parameter {
    fn as_ref(&self) -> &str {
        match self {
            Self::Number(identifier) => identifier.as_ref(),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct AstFunction {
    pub name: Identifier,
    pub num_parameters: usize,
}

impl ToTokens for AstFunction {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let Some(function) = FUNCTIONS.get(self.name.as_ref()) else {
            return; // do nothing if, for some reason, the function doesn't exist
        };

        let prefixed_fn_name =
            format_ident!("import_{}__{}", self.name.as_ref(), self.num_parameters);

        tokens.extend(quote! {
            #[inline]
        });

        tokens.extend((function.token_fn)(self.num_parameters, prefixed_fn_name));
    }
}

1	use std::{collections::BTreeSet, fmt::Debug};
2
3	use proc_macro2::TokenStream;
4	use quote::{format_ident, quote, ToTokens};
5	use snafu::ensure;
6
7	use super::{
8	error::{self, ExpressionError},
9	functions::FUNCTIONS,
10	};
11
12	type Result<T, E = ExpressionError> = std::result::Result<T, E>;
13
14	// TODO: prefix for variables and functions
15
16	/// An expression as an abstract syntax tree.
17	/// Allows genering Rust code.
18	#[derive(Debug, Clone)]	×
19	pub struct ExpressionAst {
20	/// This name is the generated function name after generating code.
21	name: Identifier,
22	root: AstNode,
23	parameters: Vec<Parameter>,
24	functions: BTreeSet<AstFunction>,
25	// TODO: dtype Float or Int
26	}
27
28	impl ExpressionAst {
29	pub fn new(	33✔
30	name: Identifier,	33✔
31	parameters: Vec<Parameter>,	33✔
32	functions: BTreeSet<AstFunction>,	33✔
33	root: AstNode,	33✔
34	) -> Result<ExpressionAst> {	33✔
35	ensure!(!name.as_ref().is_empty(), error::EmptyExpressionName);	33✔
36
37	Ok(Self {	33✔
38	name,	33✔
39	root,	33✔
40	parameters,	33✔
41	functions,	33✔
42	})	33✔
43	}	33✔
44
45	pub fn code(&self) -> String {	10✔
46	self.to_token_stream().to_string()	10✔
47	}	10✔
48
49	pub fn name(&self) -> &str {	10✔
50	self.name.as_ref()	10✔
51	}	10✔
52	}
53
54	impl ToTokens for ExpressionAst {
55	fn to_tokens(&self, tokens: &mut TokenStream) {	33✔
56	let dtype = format_ident!("{}", "f64");	33✔
57
58	for function in &self.functions {	58✔
59	function.to_tokens(tokens);	25✔
60	}	25✔
61
62	let fn_name = &self.name;	33✔
63	let params: Vec<TokenStream> = self	33✔
64	.parameters	33✔
65	.iter()	33✔
66	.map(\|p\| match p {	33✔
67	Parameter::Number(param) => quote! { #param: Option<#dtype> },	31✔
68	})	33✔
69	.collect();	33✔
70	let content = &self.root;	33✔
71		33✔
72	tokens.extend(quote! {	33✔
73	#[inline]	33✔
74	fn apply(a: Option<#dtype>, b: Option<#dtype>, f: fn(#dtype, #dtype) -> #dtype) -> Option<#dtype> {	33✔
75	match (a, b) {	33✔
76	(Some(a), Some(b)) => Some(f(a, b)),	33✔
77	_ => None,	33✔
78	}	33✔
79	}	33✔
80		33✔
81	#[no_mangle]	33✔
82	pub extern "Rust" fn #fn_name (#(#params),*) -> Option<#dtype> {	33✔
83	#content	33✔
84	}	33✔
85	});	33✔
86	}	33✔
87	}
88
89	#[derive(Debug, Clone)]	×
90	pub enum AstNode {
91	Constant(f64),
92	NoData,
93	Variable(Identifier),
94	Operation {
95	left: Box<AstNode>,
96	op: AstOperator,
97	right: Box<AstNode>,
98	},
99	Function {
100	name: Identifier,
101	args: Vec<AstNode>,
102	},
103	Branch {
104	condition_branches: Vec<Branch>,
105	else_branch: Box<AstNode>,
106	},
107	AssignmentsAndExpression {
108	assignments: Vec<Assignment>,
109	expression: Box<AstNode>,
110	},
111	}
112
113	impl ToTokens for AstNode {
114	fn to_tokens(&self, tokens: &mut TokenStream) {	205✔
115	let new_tokens = match self {	205✔
116	Self::Constant(n) => quote! { Some(#n) },	61✔
117	Self::NoData => quote! { None },	5✔
118	Self::Variable(v) => quote! { #v },	39✔
119	Self::Operation { left, op, right } => {	35✔
120	quote! { apply(#left, #right, #op) }	35✔
121	}
122	Self::Function { name, args } => {	25✔
123	let fn_name = format_ident!("import_{}__{}", name.as_ref(), args.len());	25✔
124	quote! { #fn_name(#(#args),*) }	25✔
125	}
126	AstNode::Branch {
127	condition_branches,	7✔
128	else_branch: default_branch,	7✔
129	} => {	7✔
130	let mut new_tokens = TokenStream::new();	7✔
131	for (i, branch) in condition_branches.iter().enumerate() {	12✔
132	let condition = &branch.condition;	12✔
133	let body = &branch.body;	12✔
134		12✔
135	new_tokens.extend(if i == 0 {	12✔
136	// first
137	quote! {	7✔
138	if #condition {	7✔
139	#body	7✔
140	}	7✔
141	}	7✔
142	} else {
143	// middle
144	quote! {	5✔
145	else if #condition {	5✔
146	#body	5✔
147	}	5✔
148	}	5✔
149	});
150	}
151
152	new_tokens.extend(quote! {	7✔
153	else {	7✔
154	#default_branch	7✔
155	}	7✔
156	});	7✔
157		7✔
158	new_tokens	7✔
159	}
160	Self::AssignmentsAndExpression {
161	assignments,	33✔
162	expression,	33✔
163	} => {	33✔
164	quote! {	33✔
165	#(#assignments)*	33✔
166	#expression	33✔
167	}	33✔
168	}
169	};
170
171	tokens.extend(new_tokens);	205✔
172	}	205✔
173	}
174
175	#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]	88✔
176	#[repr(transparent)]
177	pub struct Identifier(String);
178
179	impl ToTokens for Identifier {
180	fn to_tokens(&self, tokens: &mut TokenStream) {	105✔
181	let identifier = format_ident!("{}", self.0);	105✔
182	tokens.extend(quote! { #identifier });	105✔
183	}	105✔
184	}
185
186	impl From<String> for Identifier {
187	fn from(s: String) -> Self {	56✔
188	Self(s)	56✔
189	}	56✔
190	}
191
192	impl From<&str> for Identifier {
193	fn from(s: &str) -> Self {	75✔
194	Self(s.to_string())	75✔
195	}	75✔
196	}
197
198	impl From<&String> for Identifier {
199	fn from(s: &String) -> Self {	×
200	Self(s.to_string())	×
201	}	×
202	}
203
204	impl AsRef<str> for Identifier {
205	fn as_ref(&self) -> &str {	173✔
206	&self.0	173✔
207	}	173✔
208	}
209
210	impl std::fmt::Display for Identifier {
211	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {	×
212	std::fmt::Display::fmt(&self.0, f)	×
213	}	×
214	}
215
216	#[derive(Debug, Clone)]	×
217	pub enum AstOperator {
218	Add,
219	Subtract,
220	Multiply,
221	Divide,
222	}
223
224	impl ToTokens for AstOperator {
225	fn to_tokens(&self, tokens: &mut TokenStream) {	35✔
226	let new_tokens = match self {	35✔
227	AstOperator::Add => quote! { std::ops::Add::add },	23✔
228	AstOperator::Subtract => quote! { std::ops::Sub::sub },	4✔
229	AstOperator::Multiply => quote! { std::ops::Mul::mul },	6✔
230	AstOperator::Divide => quote! { std::ops::Div::div },	2✔
231	};
232
233	tokens.extend(new_tokens);	35✔
234	}	35✔
235	}
236
237	#[derive(Debug, Clone)]	×
238	pub struct Branch {
239	pub condition: BooleanExpression,
240	pub body: AstNode,
241	}
242
243	#[derive(Debug, Clone)]	×
244	pub enum BooleanExpression {
245	Constant(bool),
246	Comparison {
247	left: Box<AstNode>,
248	op: BooleanComparator,
249	right: Box<AstNode>,
250	},
251	Operation {
252	left: Box<BooleanExpression>,
253	op: BooleanOperator,
254	right: Box<BooleanExpression>,
255	},
256	}
257
258	impl ToTokens for BooleanExpression {
259	fn to_tokens(&self, tokens: &mut TokenStream) {	16✔
260	let new_tokens = match self {	16✔
261	Self::Constant(b) => quote! { #b },	6✔
262	Self::Comparison { left, op, right } => quote! { ((#left) #op (#right)) },	8✔
263	Self::Operation { left, op, right } => quote! { ( (#left) #op (#right) ) },	2✔
264	};
265
266	tokens.extend(new_tokens);	16✔
267	}	16✔
268	}
269
270	#[derive(Debug, Clone)]	×
271	pub enum BooleanComparator {
272	Equal,
273	NotEqual,
274	LessThan,
275	LessThanOrEqual,
276	GreaterThan,
277	GreaterThanOrEqual,
278	}
279
280	impl ToTokens for BooleanComparator {
281	fn to_tokens(&self, tokens: &mut TokenStream) {	8✔
282	let new_tokens = match self {	8✔
283	Self::Equal => quote! { == },	5✔
284	Self::NotEqual => quote! { != },	×
285	Self::LessThan => quote! { < },	3✔
286	Self::LessThanOrEqual => quote! { <= },	×
287	Self::GreaterThan => quote! { > },	×
288	Self::GreaterThanOrEqual => quote! { >= },	×
289	};
290
291	tokens.extend(new_tokens);	8✔
292	}	8✔
293	}
294
295	#[derive(Debug, Clone)]	×
296	pub enum BooleanOperator {
297	And,
298	Or,
299	}
300
301	impl ToTokens for BooleanOperator {
302	fn to_tokens(&self, tokens: &mut TokenStream) {	2✔
303	let new_tokens = match self {	2✔
304	Self::And => quote! { && },	2✔
305	Self::Or => quote! { \|\| },	×
306	};
307
308	tokens.extend(new_tokens);	2✔
309	}	2✔
310	}
311
312	#[derive(Debug, Clone)]	×
313	pub struct Assignment {
314	pub identifier: Identifier,
315	pub expression: AstNode,
316	}
317
318	impl ToTokens for Assignment {
319	fn to_tokens(&self, tokens: &mut TokenStream) {	2✔
320	let Self {	2✔
321	identifier,	2✔
322	expression,	2✔
323	} = self;	2✔
324	let new_tokens = quote! {	2✔
325	let #identifier = #expression;	2✔
326	};	2✔
327		2✔
328	tokens.extend(new_tokens);	2✔
329	}	2✔
330	}
331
332	#[derive(Debug, Clone)]	31✔
333	pub enum Parameter {
334	Number(Identifier),
335	}
336
337	impl AsRef<str> for Parameter {
338	fn as_ref(&self) -> &str {	31✔
339	match self {	31✔
340	Self::Number(identifier) => identifier.as_ref(),	31✔
341	}	31✔
342	}	31✔
343	}
344
345	#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]	20✔
346	pub struct AstFunction {
347	pub name: Identifier,
348	pub num_parameters: usize,
349	}
350
351	impl ToTokens for AstFunction {
352	fn to_tokens(&self, tokens: &mut TokenStream) {	25✔
353	let Some(function) = FUNCTIONS.get(self.name.as_ref()) else {	25✔
354	return; // do nothing if, for some reason, the function doesn't exist	×
355	};
356
357	let prefixed_fn_name =	25✔
358	format_ident!("import_{}__{}", self.name.as_ref(), self.num_parameters);	25✔
359		25✔
360	tokens.extend(quote! {	25✔
361	#[inline]	25✔
362	});	25✔
363		25✔
364	tokens.extend((function.token_fn)(self.num_parameters, prefixed_fn_name));	25✔
365	}	25✔
366	}

geo-engine / geoengine / 3929938005

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous