Skip to content

Commit 9a78451

Browse files
0xbrayo0xbrayo
authored
perf(query): borrow variables for read-only function calls (#675)
* perf(query): borrow variables for read-only function calls * fix(query): mark DataType non_exhaustive to guard downstream matches Adding the ReadOnlyFunction variant to the public DataType enum is a compile-time break for any consumer that matches on it exhaustively. Mark the enum #[non_exhaustive] so future variants can't silently break downstream code; all matches on DataType inside aw-query are local to the crate and are unaffected.
1 parent b834a8f commit 9a78451

4 files changed

Lines changed: 214 additions & 30 deletions

File tree

‎aw-query/benches/benchmark.rs‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -102,11 +102,26 @@ mod query_benchmarks {
102102
})
103103
});
104104
}
105+
106+
pub fn bench_read_only_references(c: &mut Criterion) {
107+
let ds = setup_datastore();
108+
create_bucket(&ds, BUCKETNAME.to_string());
109+
insert_events(&ds, BUCKETNAME, 50_000);
110+
let interval = TimeInterval::new_from_string(TIME_INTERVAL).unwrap();
111+
let code = format!(
112+
"events = query_bucket(\"testbucket\"); total = 0; {} return total;",
113+
"total = total + sum_durations(events);".repeat(20)
114+
);
115+
c.bench_function("read-only references 50000 events", |b| {
116+
b.iter(|| aw_query::query(&code, &interval, &ds).unwrap())
117+
});
118+
}
105119
}
106120

107121
criterion_group!(
108122
benches,
109123
query_benchmarks::bench_assign,
110-
query_benchmarks::bench_many_events
124+
query_benchmarks::bench_many_events,
125+
query_benchmarks::bench_read_only_references
111126
);
112127
criterion_main!(benches);

‎aw-query/src/datatype.rs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ use serde_json::Number;
1414

1515
#[derive(Clone, Serialize)]
1616
#[serde(untagged)]
17+
#[non_exhaustive]
1718
pub enum DataType {
1819
None(),
1920
Bool(bool),
@@ -24,14 +25,12 @@ pub enum DataType {
2425
Dict(HashMap<String, DataType>),
2526
#[serde(serialize_with = "serialize_function")]
2627
Function(String, functions::QueryFn),
28+
#[serde(serialize_with = "serialize_function")]
29+
ReadOnlyFunction(String, functions::ReadOnlyQueryFn),
2730
}
2831

2932
#[allow(clippy::trivially_copy_pass_by_ref)]
30-
fn serialize_function<S>(
31-
_element: &str,
32-
_fun: &functions::QueryFn,
33-
_serializer: S,
34-
) -> Result<S::Ok, S::Error>
33+
fn serialize_function<S, F>(_element: &str, _fun: &F, _serializer: S) -> Result<S::Ok, S::Error>
3534
where
3635
S: Serializer,
3736
{
@@ -53,6 +52,7 @@ impl fmt::Debug for DataType {
5352
DataType::List(l) => write!(f, "List({l:?})"),
5453
DataType::Dict(d) => write!(f, "Dict({d:?})"),
5554
DataType::Function(name, _fun) => write!(f, "Function({name})"),
55+
DataType::ReadOnlyFunction(name, _fun) => write!(f, "Function({name})"),
5656
}
5757
}
5858
}

‎aw-query/src/functions.rs‎

Lines changed: 31 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,13 @@ use aw_datastore::Datastore;
66
pub type QueryFn =
77
fn(args: Vec<DataType>, env: &VarEnv, ds: &Datastore) -> Result<DataType, QueryError>;
88

9+
pub type ReadOnlyQueryFn =
10+
fn(args: &[&DataType], env: &VarEnv, ds: &Datastore) -> Result<DataType, QueryError>;
11+
912
pub fn fill_env(env: &mut VarEnv) {
1013
env.insert(
1114
"print".to_string(),
12-
DataType::Function("print".to_string(), qfunctions::print),
15+
DataType::ReadOnlyFunction("print".to_string(), qfunctions::print),
1316
);
1417
env.insert(
1518
"query_bucket".to_string(),
@@ -35,15 +38,15 @@ pub fn fill_env(env: &mut VarEnv) {
3538
);
3639
env.insert(
3740
"sum_durations".to_string(),
38-
DataType::Function("sum_durations".to_string(), qfunctions::sum_durations),
41+
DataType::ReadOnlyFunction("sum_durations".to_string(), qfunctions::sum_durations),
3942
);
4043
env.insert(
4144
"limit_events".to_string(),
4245
DataType::Function("limit_events".to_string(), qfunctions::limit_events),
4346
);
4447
env.insert(
4548
"contains".to_string(),
46-
DataType::Function("contains".to_string(), qfunctions::contains),
49+
DataType::ReadOnlyFunction("contains".to_string(), qfunctions::contains),
4750
);
4851
env.insert(
4952
"flood".to_string(),
@@ -126,7 +129,7 @@ mod qfunctions {
126129
use crate::VarEnv;
127130

128131
pub fn print(
129-
args: Vec<DataType>,
132+
args: &[&DataType],
130133
_env: &VarEnv,
131134
_ds: &Datastore,
132135
) -> Result<DataType, QueryError> {
@@ -236,25 +239,25 @@ mod qfunctions {
236239
}
237240

238241
pub fn contains(
239-
args: Vec<DataType>,
242+
args: &[&DataType],
240243
_env: &VarEnv,
241244
_ds: &Datastore,
242245
) -> Result<DataType, QueryError> {
243246
// typecheck
244-
validate::args_length(&args, 2)?;
245-
match args.first().unwrap() {
246-
DataType::List(ref list) => Ok(DataType::Bool(list.contains(&args[1]))),
247+
validate::args_length(args, 2)?;
248+
match args[0] {
249+
DataType::List(list) => Ok(DataType::Bool(list.contains(args[1]))),
247250
DataType::Dict(ref dict) => {
248251
let s = match &args[1] {
249-
DataType::String(s) => s.to_string(),
252+
DataType::String(s) => s.as_str(),
250253
_ => {
251254
return Err(QueryError::InvalidFunctionParameters(format!(
252255
"function contains got second argument {:?}, expected type String",
253256
args[0]
254257
)))
255258
}
256259
};
257-
Ok(DataType::Bool(dict.contains_key(&s)))
260+
Ok(DataType::Bool(dict.contains_key(s)))
258261
}
259262
_ => Err(QueryError::InvalidFunctionParameters(format!(
260263
"function contains got first argument {:?}, expected type List or Dict",
@@ -381,18 +384,28 @@ mod qfunctions {
381384
}
382385

383386
pub fn sum_durations(
384-
args: Vec<DataType>,
387+
args: &[&DataType],
385388
_env: &VarEnv,
386389
_ds: &Datastore,
387390
) -> Result<DataType, QueryError> {
388391
// typecheck
389-
validate::args_length(&args, 1)?;
390-
let mut events: Vec<Event> = args.into_iter().next().unwrap().try_into()?;
391-
392-
// Sort by duration
392+
validate::args_length(args, 1)?;
393+
let events = match args[0] {
394+
DataType::List(events) => events,
395+
invalid_type => {
396+
return Err(QueryError::InvalidFunctionParameters(format!(
397+
"Expected function parameter of type List, got {invalid_type:?}"
398+
)))
399+
}
400+
};
393401
let mut sum_durations = chrono::Duration::zero();
394-
for event in events.drain(..) {
395-
sum_durations += event.duration;
402+
for event in events {
403+
match event {
404+
DataType::Event(event) => sum_durations += event.duration,
405+
invalid_type => return Err(QueryError::InvalidFunctionParameters(format!(
406+
"Expected function parameter of type List of Events, list contains {invalid_type:?}"
407+
))),
408+
}
396409
}
397410
Ok(DataType::Number(
398411
(sum_durations.num_milliseconds() as f64) / 1000.0,
@@ -601,7 +614,7 @@ mod validate {
601614
use crate::{DataType, QueryError, VarEnv};
602615
use aw_models::TimeInterval;
603616

604-
pub fn args_length(args: &[DataType], len: usize) -> Result<(), QueryError> {
617+
pub fn args_length<T>(args: &[T], len: usize) -> Result<(), QueryError> {
605618
if args.len() != len {
606619
return Err(QueryError::InvalidFunctionParameters(format!(
607620
"Expected {} parameters in function, got {}",

‎aw-query/src/interpret.rs‎

Lines changed: 162 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use std::collections::HashMap;
1+
use std::{borrow::Cow, collections::HashMap};
22

33
use crate::functions;
44

@@ -213,6 +213,41 @@ fn interpret_expr(
213213
Ok(DataType::None())
214214
}
215215
Function(fname, e) => {
216+
// Borrow variables only when every argument is an atom. Evaluating
217+
// arbitrary arguments may assign to the environment, including the
218+
// function binding itself, so those retain the owned evaluation path.
219+
if let Expr_::List(exprs) = &e.node {
220+
if exprs.iter().all(|expr| {
221+
matches!(
222+
expr.node,
223+
Expr_::Var(_) | Expr_::Bool(_) | Expr_::Number(_) | Expr_::String(_)
224+
)
225+
}) {
226+
if let Some(DataType::ReadOnlyFunction(_, fun)) = env.get(&fname) {
227+
let args: Result<Vec<Cow<'_, DataType>>, QueryError> = exprs
228+
.iter()
229+
.map(|expr| {
230+
Ok(match &expr.node {
231+
Expr_::Var(name) => {
232+
Cow::Borrowed(env.get(name).ok_or_else(|| {
233+
QueryError::VariableNotDefined(name.clone())
234+
})?)
235+
}
236+
Expr_::Bool(value) => Cow::Owned(DataType::Bool(*value)),
237+
Expr_::Number(value) => Cow::Owned(DataType::Number(*value)),
238+
Expr_::String(value) => {
239+
Cow::Owned(DataType::String(value.clone()))
240+
}
241+
_ => unreachable!(),
242+
})
243+
})
244+
.collect();
245+
let args = args?;
246+
let refs: Vec<_> = args.iter().map(|arg| arg.as_ref()).collect();
247+
return fun(&refs, env, ds);
248+
}
249+
}
250+
}
216251
let args = match interpret_expr(env, ds, *e)? {
217252
DataType::List(l) => l,
218253
_ => unreachable!(),
@@ -221,11 +256,14 @@ fn interpret_expr(
221256
Some(v) => v,
222257
None => return Err(QueryError::VariableNotDefined(fname.clone())),
223258
};
224-
let (_name, fun) = match var {
225-
DataType::Function(name, fun) => (name, fun),
226-
_data => return Err(QueryError::InvalidType(fname.to_string())),
227-
};
228-
fun(args, env, ds)
259+
match var {
260+
DataType::Function(_, fun) => fun(args, env, ds),
261+
DataType::ReadOnlyFunction(_, fun) => {
262+
let refs: Vec<_> = args.iter().collect();
263+
fun(&refs, env, ds)
264+
}
265+
_ => Err(QueryError::InvalidType(fname.to_string())),
266+
}
229267
}
230268
List(list) => {
231269
let mut l = Vec::new();
@@ -245,3 +283,121 @@ fn interpret_expr(
245283
}
246284
}
247285
}
286+
287+
#[cfg(test)]
288+
mod tests {
289+
use super::*;
290+
291+
fn run(code: &str, env: &mut VarEnv, ds: &Datastore) -> DataType {
292+
let program = crate::parser::parse(crate::lexer::Lexer::new(code)).unwrap();
293+
for expr in program.stmts {
294+
interpret_expr(env, ds, expr).unwrap();
295+
}
296+
env.remove("RETURN").unwrap_or(DataType::None())
297+
}
298+
299+
#[test]
300+
fn read_only_arguments_borrow_the_original_value_and_support_aliases() {
301+
fn check(args: &[&DataType], env: &VarEnv, _: &Datastore) -> Result<DataType, QueryError> {
302+
assert!(std::ptr::eq(args[0], env.get("events").unwrap()));
303+
Ok(DataType::Bool(true))
304+
}
305+
let ds = Datastore::new_in_memory(false);
306+
let mut env = VarEnv::new();
307+
env.insert(
308+
"check".into(),
309+
DataType::ReadOnlyFunction("check".into(), check),
310+
);
311+
env.insert(
312+
"events".into(),
313+
DataType::List(vec![DataType::Event(aw_models::Event::default())]),
314+
);
315+
assert_eq!(
316+
run("alias = check; return alias(events);", &mut env, &ds),
317+
DataType::Bool(true)
318+
);
319+
assert!(env.contains_key("events"));
320+
}
321+
322+
#[test]
323+
fn side_effecting_arguments_keep_snapshots_and_resolve_the_function_after_evaluation() {
324+
fn check(args: &[&DataType], env: &VarEnv, _: &Datastore) -> Result<DataType, QueryError> {
325+
assert_eq!(args[0], &DataType::List(vec![DataType::Number(1.0)]));
326+
assert_eq!(env["items"], DataType::List(vec![DataType::Number(2.0)]));
327+
Ok(DataType::Bool(true))
328+
}
329+
let ds = Datastore::new_in_memory(false);
330+
let mut env = VarEnv::new();
331+
env.insert(
332+
"check".into(),
333+
DataType::ReadOnlyFunction("check".into(), check),
334+
);
335+
// Assignment nodes are supported by the interpreter, although the
336+
// current parser only emits them at statement level.
337+
let expr = |node| Expr {
338+
span: crate::lexer::Span {
339+
lo: 0,
340+
hi: 0,
341+
line: 1,
342+
},
343+
node,
344+
};
345+
for rebind in [false, true] {
346+
env.insert("items".into(), DataType::List(vec![DataType::Number(1.0)]));
347+
env.insert(
348+
"f".into(),
349+
if rebind {
350+
DataType::Number(0.0)
351+
} else {
352+
env["check"].clone()
353+
},
354+
);
355+
let mut args = vec![
356+
expr(Expr_::Var("items".into())),
357+
expr(Expr_::Assign(
358+
"items".into(),
359+
Box::new(expr(Expr_::List(vec![expr(Expr_::Number(2.0))]))),
360+
)),
361+
];
362+
if rebind {
363+
args.push(expr(Expr_::Assign(
364+
"f".into(),
365+
Box::new(expr(Expr_::Var("check".into()))),
366+
)));
367+
}
368+
let call = expr(Expr_::Function(
369+
"f".into(),
370+
Box::new(expr(Expr_::List(args))),
371+
));
372+
assert_eq!(
373+
interpret_expr(&mut env, &ds, call).unwrap(),
374+
DataType::Bool(true)
375+
);
376+
}
377+
}
378+
379+
#[test]
380+
fn read_only_builtins_keep_values_available_and_accept_computed_arguments() {
381+
let ds = Datastore::new_in_memory(false);
382+
let mut env = VarEnv::new();
383+
functions::fill_env(&mut env);
384+
let mut event = aw_models::Event::default();
385+
event.duration = chrono::Duration::seconds(3);
386+
env.insert(
387+
"events".into(),
388+
DataType::List(vec![DataType::Event(event)]),
389+
);
390+
assert_eq!(run("total = sum_durations; a = total(events); b = total(events + events); return a + b;", &mut env, &ds), DataType::Number(9.0));
391+
assert_eq!(
392+
run(
393+
"values = [1, 2]; f = contains; a = f(values, 2); return [a, values];",
394+
&mut env,
395+
&ds
396+
),
397+
DataType::List(vec![
398+
DataType::Bool(true),
399+
DataType::List(vec![DataType::Number(1.0), DataType::Number(2.0)])
400+
])
401+
);
402+
}
403+
}

0 commit comments

Comments
 (0)