aw-query: Add preprocessing and refcounting

johan-bjareholt · johan-bjareholt · commit 072753bb32d1 · 2020-11-01T23:58:23.000+01:00
diff --git a/aw-query/src/functions.rs b/aw-query/src/functions.rs
@@ -7,97 +7,97 @@ pub type QueryFn =
     fn(args: Vec<DataType>, env: &VarEnv, ds: &Datastore) -> Result<DataType, QueryError>;
 
 pub fn fill_env(env: &mut VarEnv) {
-    env.insert(
+    env.declare_static(
         "print".to_string(),
         DataType::Function("print".to_string(), qfunctions::print),
     );
-    env.insert(
+    env.declare_static(
         "query_bucket".to_string(),
         DataType::Function("query_bucket".to_string(), qfunctions::query_bucket),
     );
-    env.insert(
+    env.declare_static(
         "query_bucket_names".to_string(),
         DataType::Function(
             "query_bucket_names".to_string(),
             qfunctions::query_bucket_names,
         ),
     );
-    env.insert(
+    env.declare_static(
         "sort_by_duration".to_string(),
         DataType::Function("sort_by_duration".to_string(), qfunctions::sort_by_duration),
     );
-    env.insert(
+    env.declare_static(
         "sort_by_timestamp".to_string(),
         DataType::Function(
             "sort_by_timestamp".to_string(),
             qfunctions::sort_by_timestamp,
         ),
     );
-    env.insert(
+    env.declare_static(
         "sum_durations".to_string(),
         DataType::Function("sum_durations".to_string(), qfunctions::sum_durations),
     );
-    env.insert(
+    env.declare_static(
         "limit_events".to_string(),
         DataType::Function("limit_events".to_string(), qfunctions::limit_events),
     );
-    env.insert(
+    env.declare_static(
         "contains".to_string(),
         DataType::Function("contains".to_string(), qfunctions::contains),
     );
-    env.insert(
+    env.declare_static(
         "flood".to_string(),
         DataType::Function("flood".to_string(), qfunctions::flood),
     );
-    env.insert(
+    env.declare_static(
         "find_bucket".to_string(),
         DataType::Function("find_bucket".to_string(), qfunctions::find_bucket),
     );
-    env.insert(
+    env.declare_static(
         "merge_events_by_keys".to_string(),
         DataType::Function(
             "merge_events_by_keys".to_string(),
             qfunctions::merge_events_by_keys,
         ),
     );
-    env.insert(
+    env.declare_static(
         "chunk_events_by_key".to_string(),
         DataType::Function(
             "chunk_events_by_key".to_string(),
             qfunctions::chunk_events_by_key,
         ),
     );
-    env.insert(
+    env.declare_static(
         "filter_keyvals".to_string(),
         DataType::Function("filter_keyvals".to_string(), qfunctions::filter_keyvals),
     );
-    env.insert(
+    env.declare_static(
         "filter_keyvals_regex".to_string(),
         DataType::Function(
             "filter_keyvals_regex".to_string(),
             qfunctions::filter_keyvals_regex,
         ),
     );
-    env.insert(
+    env.declare_static(
         "filter_period_intersect".to_string(),
         DataType::Function(
             "filter_period_intersect".to_string(),
             qfunctions::filter_period_intersect,
         ),
     );
-    env.insert(
+    env.declare_static(
         "split_url_events".to_string(),
         DataType::Function("split_url_events".to_string(), qfunctions::split_url_events),
     );
-    env.insert(
+    env.declare_static(
         "concat".to_string(),
         DataType::Function("concat".to_string(), qfunctions::concat),
     );
-    env.insert(
+    env.declare_static(
         "categorize".to_string(),
         DataType::Function("categorize".into(), qfunctions::categorize),
     );
-    env.insert(
+    env.declare_static(
         "tag".to_string(),
         DataType::Function("tag".into(), qfunctions::tag),
     );
@@ -521,7 +521,7 @@ mod validate {
     }
 
     pub fn get_timeinterval(env: &VarEnv) -> Result<TimeInterval, QueryError> {
-        let interval_str = match env.get("TIMEINTERVAL") {
+        let interval_str = match env.deprecated_get("TIMEINTERVAL") {
             Some(data_ti) => match data_ti {
                 DataType::String(ti_str) => ti_str,
                 _ => {
diff --git a/aw-query/src/interpret.rs b/aw-query/src/interpret.rs
@@ -1,40 +1,28 @@
 use std::collections::HashMap;
 
-use crate::functions;
-
 use aw_datastore::Datastore;
-use aw_models::TimeInterval;
 
 use crate::ast::*;
 use crate::DataType;
 use crate::QueryError;
-
-pub type VarEnv = HashMap<String, DataType>;
-
-fn init_env(ti: &TimeInterval) -> VarEnv {
-    let mut env = HashMap::new();
-    env.insert("TIMEINTERVAL".to_string(), DataType::String(ti.to_string()));
-    functions::fill_env(&mut env);
-    env
-}
+use crate::VarEnv;
 
 pub fn interpret_prog(
     p: Program,
-    ti: &TimeInterval,
+    env: &mut VarEnv,
     ds: &Datastore,
 ) -> Result<DataType, QueryError> {
-    let mut env = init_env(ti);
     for expr in p.stmts {
-        interpret_expr(&mut env, ds, expr)?;
+        interpret_expr(env, ds, expr)?;
     }
-    match env.remove("RETURN") {
+    match env.take("RETURN") {
         Some(ret) => Ok(ret),
         None => Err(QueryError::EmptyQuery()),
     }
 }
 
 fn interpret_expr(
-    env: &mut HashMap<String, DataType>,
+    env: &mut VarEnv,
     ds: &Datastore,
     expr: Expr,
 ) -> Result<DataType, QueryError> {
@@ -184,9 +172,8 @@ fn interpret_expr(
             env.insert(var, val);
             Ok(DataType::None())
         }
-        // FIXME: avoid clone, it's slow
-        Var(var) => match env.get(&var) {
-            Some(v) => Ok(v.clone()),
+        Var(var) => match env.take(&var) {
+            Some(v) => Ok(v),
             None => Err(QueryError::VariableNotDefined(var.to_string())),
         },
         Bool(lit) => Ok(DataType::Bool(lit)),
@@ -215,7 +202,7 @@ fn interpret_expr(
                 DataType::List(l) => l,
                 _ => unreachable!(),
             };
-            let var = match env.get(&fname[..]) {
+            let var = match env.take(&fname[..]) {
                 Some(v) => v,
                 None => return Err(QueryError::VariableNotDefined(fname.clone())),
             };
diff --git a/aw-query/src/lib.rs b/aw-query/src/lib.rs
@@ -21,9 +21,11 @@ mod lexer;
     unused_braces
 )]
 mod parser;
+mod varenv;
+mod preprocess;
 
 pub use crate::datatype::DataType;
-pub use crate::interpret::VarEnv;
+pub use crate::varenv::VarEnv;
 
 // TODO: add line numbers to errors
 // (works during lexing, but not during parsing I believe)
@@ -50,6 +52,14 @@ impl fmt::Display for QueryError {
     }
 }
 
+fn init_env(ti: &TimeInterval) -> VarEnv {
+    let mut env = VarEnv::new();
+    env.declare_static("TIMEINTERVAL".to_string(), DataType::String(ti.to_string()));
+    env.declare_static("RETURN".to_string(), DataType::None());
+    functions::fill_env(&mut env);
+    env
+}
+
 pub fn query(code: &str, ti: &TimeInterval, ds: &Datastore) -> Result<DataType, QueryError> {
     let lexer = lexer::Lexer::new(code);
     let program = match parser::parse(lexer) {
@@ -60,5 +70,7 @@ pub fn query(code: &str, ti: &TimeInterval, ds: &Datastore) -> Result<DataType,
             return Err(QueryError::ParsingError(format!("{:?}", e)));
         }
     };
-    interpret::interpret_prog(program, ti, ds)
+    let mut env = init_env(ti);
+    preprocess::preprocess_prog(&program, &mut env, ds)?;
+    interpret::interpret_prog(program, &mut env, ds)
 }
diff --git a/aw-query/src/preprocess.rs b/aw-query/src/preprocess.rs
@@ -0,0 +1,87 @@
+use aw_datastore::Datastore;
+
+use crate::ast::*;
+use crate::DataType;
+use crate::QueryError;
+use crate::VarEnv;
+
+pub fn preprocess_prog(
+    p: &Program,
+    env: &mut VarEnv,
+    ds: &Datastore,
+) -> Result<(), QueryError> {
+    for expr in &p.stmts {
+        preprocess_expr(env, ds, expr)?;
+    }
+    Ok(())
+}
+
+fn preprocess_expr(
+    env: &mut VarEnv,
+    ds: &Datastore,
+    expr: &Expr,
+) -> Result<(), QueryError> {
+    use crate::ast::Expr_::*;
+    match &expr.node {
+        Var(ref var) => env.add_ref(var)?,
+        Add(ref a, ref b) => {
+            preprocess_expr(env, ds, a)?;
+            preprocess_expr(env, ds, b)?;
+        }
+        Sub(ref a, ref b) => {
+            preprocess_expr(env, ds, a)?;
+            preprocess_expr(env, ds, b)?;
+        }
+        Mul(ref a, ref b) => {
+            preprocess_expr(env, ds, a)?;
+            preprocess_expr(env, ds, b)?;
+        }
+        Div(ref a, ref b) => {
+            preprocess_expr(env, ds, a)?;
+            preprocess_expr(env, ds, b)?;
+        }
+        Mod(ref a, ref b) => {
+            preprocess_expr(env, ds, a)?;
+            preprocess_expr(env, ds, b)?;
+        }
+        Equal(ref a, ref b) => {
+            preprocess_expr(env, ds, a)?;
+            preprocess_expr(env, ds, b)?;
+        }
+        Assign(ref var, ref b) => {
+            preprocess_expr(env, ds, b)?;
+            env.declare(var.to_string());
+        },
+        Function(ref fname, ref args) => {
+            env.add_ref(fname)?;
+            preprocess_expr(env, ds, args)?;
+        }
+        If(ref ifs) => {
+            for (cond, block) in ifs {
+                // TODO: could be optimized?
+                preprocess_expr(env, ds, cond)?;
+                for expr in block {
+                    preprocess_expr(env, ds, expr)?;
+                }
+            }
+        }
+        List(list) => {
+            for entry in list {
+                preprocess_expr(env, ds, entry)?;
+            }
+        }
+        Dict(d) => {
+            for (key, val_uninterpreted) in d {
+                preprocess_expr(env, ds, val_uninterpreted)?;
+            }
+        }
+        Return(e) => {
+            println!("return!");
+            preprocess_expr(env, ds, e)?;
+        },
+        Bool(_lit) => (),
+        Number(_lit) => (),
+        String(_lit) => (),
+    };
+    Ok(())
+}
diff --git a/aw-query/src/varenv.rs b/aw-query/src/varenv.rs
