Parser: support 'alloc' vars (#168)

* Create skeleton to plan for 'alloc' var logic

* Support 'alloc' variables in the parser

* Add some comments to plan and panic on special_decl_init

* add some more todos

src/emitter/module_emitter.rs

@@ -137,6 +137,7 @@ impl<'a, 'b, 'c, 'd, 'e, 'f> ModuleEmitter<'a, 'b, 'c, 'd, 'e, 'f> {
         let func = FunctionBuilder::new(params, results);
         self.emitting_func = Some(func);
 
+        // TODO(alloc) -- load alloc vars
         if self.map_lib_adapter.is_used {
             let fid = self.map_lib_adapter.get_map_init_fid(self.app_wasm, err);
             if let Some(func) = &mut self.emitting_func {
@@ -147,6 +148,7 @@ impl<'a, 'b, 'c, 'd, 'e, 'f> ModuleEmitter<'a, 'b, 'c, 'd, 'e, 'f> {
         // emit the function body
         self.emit_body(&[], block, err);
 
+        // TODO(alloc) -- save alloc vars
         // emit the function
         if let Some(func) = self.emitting_func.take() {
             let fid = func.finish_module(self.app_wasm);

src/emitter/utils.rs

@@ -116,16 +116,22 @@ pub fn emit_stmt<'a, T: Opcode<'a> + MacroOpcode<'a> + AddLocal>(
                 )
             }
         }
-        Statement::ReportDecl { .. } => emit_report_decl_stmt(
-            stmt,
-            injector,
-            table,
-            mem_tracker,
-            map_lib_adapter,
-            report_var_metadata,
-            err_msg,
-            err,
-        ),
+        Statement::AllocDecl { is_report, .. } => {
+            if *is_report {
+                emit_report_decl_stmt(
+                    stmt,
+                    injector,
+                    table,
+                    mem_tracker,
+                    map_lib_adapter,
+                    report_var_metadata,
+                    err_msg,
+                    err,
+                )
+            } else {
+                panic!()
+            }
+        }
         Statement::SetMap { .. } => emit_set_map_stmt(
             stmt,
             injector,
@@ -241,7 +247,11 @@ fn emit_report_decl_stmt<'a, T: Opcode<'a> + MacroOpcode<'a> + AddLocal>(
     err_msg: &str,
     err: &mut ErrorGen,
 ) -> bool {
-    if let Statement::ReportDecl { decl, .. } = stmt {
+    // TODO(alloc)
+    //   call lang_features.alloc_vars.rewriting IF doing rewriting...
+    //   ...will need to thread injection method through
+    //   (ignore this statement on wizard target since it's already handled)
+    if let Statement::AllocDecl { decl, .. } = stmt {
         return match &**decl {
             Statement::Decl { ty, var_id, .. } => {
                 // look up in symbol table

src/generator/mod.rs

@@ -351,7 +351,7 @@ impl<T: GeneratingVisitor> WhammVisitorMut<bool> for T {
 
                 is_success
             }
-            Statement::ReportDecl { decl, .. } => self.visit_stmt(decl),
+            Statement::AllocDecl { decl, .. } => self.visit_stmt(decl),
             Statement::SetMap { map, key, val, .. } => {
                 let mut is_success = true;
                 is_success &= self.visit_expr(map);

src/generator/rewriting/simple_ast.rs

@@ -302,7 +302,11 @@ impl WhammVisitor<()> for SimpleASTBuilder<'_, '_> {
         trace!("Entering: BehaviorTreeBuilder::visit_stmt");
         // for checking for report_decls
         match stmt {
-            Statement::ReportDecl { .. } => self.curr_num_reports += 1,
+            Statement::AllocDecl { is_report, .. } => {
+                if *is_report {
+                    self.curr_num_reports += 1
+                }
+            }
             Statement::If { conseq, alt, .. } => {
                 self.visit_block(conseq);
                 self.visit_block(alt);

src/generator/wizard/metadata_collector.rs

@@ -234,8 +234,12 @@ impl WhammVisitor<()> for WizardProbeMetadataCollector<'_, '_, '_> {
             Statement::Decl { .. } => {
                 // ignore
             }
-            Statement::ReportDecl { .. } => {
-                self.curr_probe.incr_reports();
+            Statement::AllocDecl { is_report, .. } => {
+                if *is_report {
+                    self.curr_probe.incr_reports();
+                }
+                // change this to save off data to allocate
+                todo!()
             }
             Statement::Assign { var_id, expr, .. } => {
                 if let Expr::VarId { name, .. } = var_id {

src/generator/wizard/mod.rs

@@ -121,6 +121,7 @@ impl WizardGenerator<'_, '_, '_, '_, '_, '_, '_, '_, '_> {
         let mut params = vec![];
         let mut param_str = "".to_string();
 
+        // TODO(alloc) -- handle $alloc param (if there are alloc vars)
         for (local_id, (param_name, param_ty)) in param_reqs.iter().enumerate() {
             // handle param list
             params.push(whamm_type_to_wasm_type(param_ty));
@@ -168,7 +169,6 @@ impl WizardGenerator<'_, '_, '_, '_, '_, '_, '_, '_, '_> {
     }
 
     fn visit_wiz_probe(&mut self, probe: &mut WizardProbe) {
-        // TODO -- handle provided functions (provider, package, event, probe)
         self.set_curr_loc(self.create_curr_loc(probe));
 
         let (pred_fid, pred_param_str) = if let Some(pred) = &mut probe.predicate {

src/lang_features/alloc_vars/mod.rs

@@ -0,0 +1,2 @@
+pub mod rewriting;
+pub mod wizard;

src/lang_features/alloc_vars/rewriting.rs

@@ -0,0 +1,19 @@
+use crate::parser::types::DataType;
+use crate::verifier::types::VarAddr;
+use std::collections::HashMap;
+
+pub fn allocate_vars(_to_alloc: Vec<(String, DataType)>) -> HashMap<String, (VarAddr, DataType)> {
+    // Called once per probe match point with `alloc` OR `report` vars.
+
+    // result: The new global variables: name -> (addr, ty)
+    //     as a hashmap to enable caller to place in SymbolTable and handle report variables
+    //   Add these VarAddrs to the symbol table.
+    //   Can now emit the rest of the probe body logic as normal.
+
+    // NOTE: `decl_init` statements should be run ONCE
+
+    // See utils.rs/`emit_report_decl_stmt`
+    //    basically want to do just refactor to call out to THIS function instead (more modular)
+    //    this function will also handle if the var is a report variable
+    todo!()
+}

src/lang_features/alloc_vars/wizard.rs

@@ -0,0 +1,49 @@
+use crate::parser::types::DataType;
+use crate::verifier::types::VarAddr;
+use orca_wasm::ir::id::FunctionID;
+use std::collections::HashMap;
+
+pub fn create_alloc_func(_to_alloc: Vec<DataType>) -> FunctionID {
+    // Called once per probe definition with `alloc` OR `report` vars.
+
+    // $alloc description:
+    // Will generate a function that allocates the memory required.
+    // The order of the data will go in the order of the `to_alloc` param.
+    // The function will return the memory offset to use by the probe logic.
+    // Will also need to update some global value that keeps track of the already (use memory allocator?)
+    // allocated memory space. Make sure to check that memory size is big enough!
+
+    // NOTE: `decl_init` statements should be run ONCE (can be in $alloc func)
+    todo!()
+}
+
+pub fn load_alloc_vars(
+    _alloc_mem_offset: VarAddr,
+    _alloc_vars: Vec<(String, DataType)>,
+) -> HashMap<String, (VarAddr, DataType)> {
+    // increment by bytes used by each var's DataType as we load them
+    // will be used to load the next variable AND to save in the VarAddr!
+    let _used_bytes = 0;
+
+    // alloc_mem_offset: parameter that specifies the result of calling $alloc (should be wasm local var)
+    // alloc_vars: Vec<(var_name, var_ty)>, in the order they should appear in memory starting at
+    //     the offset `alloc_mem_offset`
+    // result: The new local variables: name -> (addr, ty)
+    //     as a hashmap to enable caller to place in SymbolTable and handle report variables
+    //     the new VarAddr will be a pointer into memory with an offset (real addr should be alloc_mem_offset + len_of_prev_vars)
+
+    // At start of probe logic, pull the current values of the 'alloc' variables from memory.
+    //   Add these VarAddrs to the symbol table.
+    //   Can now emit the rest of the probe body logic as normal.
+
+    todo!()
+}
+
+pub fn save_alloc_vars(
+    _alloc_mem_offset: VarAddr,
+    _allocated_vars: HashMap<String, (VarAddr, DataType)>,
+) {
+    // At end of probe logic, save the values of 'alloc' variables back into memory.
+
+    todo!()
+}

src/lang_features/libraries/core/io/mod.rs

@@ -131,7 +131,11 @@ impl WhammVisitor<bool> for IOPackage {
     }
 
     fn visit_stmt(&mut self, stmt: &Statement) -> bool {
-        matches!(stmt, Statement::ReportDecl { .. })
+        if let Statement::AllocDecl { is_report, .. } = stmt {
+            *is_report
+        } else {
+            false
+        }
     }
 
     fn visit_expr(&mut self, _expr: &Expr) -> bool {

src/lang_features/mod.rs

@@ -1 +1,2 @@
+pub mod alloc_vars;
 pub mod libraries;

src/parser/print_visitor.rs

@@ -421,8 +421,11 @@ impl WhammVisitor<String> for AsStrVisitor {
                     self.visit_expr(val)
                 )
             }
-            Statement::ReportDecl { decl, .. } => {
-                format!("report {}", self.visit_stmt(decl))
+            Statement::AllocDecl {
+                decl, is_report, ..
+            } => {
+                let report = if *is_report { "report" } else { "" };
+                format!("{report} alloc {}", self.visit_stmt(decl))
             }
         }
     }

src/parser/tests.rs

@@ -179,6 +179,48 @@ wasm:opcode:br:before {
         i --;
     }
     "#,
+    // report variables
+    r#"
+        i32 a;
+        report i32 c;
+        wasm::br:before {
+            a = 1;
+            report bool b;
+        }
+    "#,
+    // TODO -- uncomment when we've supported special_decl_init
+    // r#"
+    //     i32 a;
+    //     report i32 c;
+    //     wasm::br:before {
+    //         a = 1;
+    //         report bool b = true;
+    //     }
+    // "#,
+    // alloc variables
+    r#"
+        i32 a;
+        alloc i32 c;
+        wasm::br:before {
+            a = 1;
+            alloc bool b;
+        }
+    "#,
+    // TODO -- uncomment when we've supported special_decl_init
+    // r#"
+    //     i32 a;
+    //     alloc i32 c;
+    //     wasm::br:before {
+    //         a = 1;
+    //         alloc bool b = true;
+    //     }
+    // "#,
+    // special variables
+    r#"
+        alloc report i32 c;
+        report alloc i32 c;
+        wasm::br:before {}
+    "#,
     // Comments
     r#"
 /* comment */
@@ -429,6 +471,40 @@ map<i32, i32> arg0;
             count[] = my_fn();
         }
     "#,
+    // use report multiple times
+    r#"
+        i32 a;
+        report alloc report i32 c;
+        wasm::br:before {
+            a = 1;
+            report bool b;
+        }
+    "#,
+    r#"
+        i32 a;
+        report i32 c;
+        wasm::br:before {
+            a = 1;
+            report alloc report bool b;
+        }
+    "#,
+    // use alloc multiple times
+    r#"
+        i32 a;
+        alloc report alloc i32 c;
+        wasm::br:before {
+            a = 1;
+            alloc bool b;
+        }
+    "#,
+    r#"
+        i32 a;
+        alloc i32 c;
+        wasm::br:before {
+            a = 1;
+            alloc report alloc bool b;
+        }
+    "#,
 ];
 const SPECIAL: &[&str] = &["BEGIN { }", "END { }", "wasm:::alt { }"];
 
@@ -490,7 +566,7 @@ pub fn get_ast(script: &str, err: &mut ErrorGen) -> Whamm {
 }
 
 fn is_valid_script(script: &str, err: &mut ErrorGen) -> bool {
-    parse_script(&script.to_string(), err).is_some()
+    parse_script(&script.to_string(), err).is_some() && !err.has_errors
 }
 
 pub fn run_test_on_valid_list(scripts: Vec<String>, err: &mut ErrorGen) {

src/parser/types.rs

@@ -289,7 +289,10 @@ pub enum Statement {
         alt: Block,
         loc: Option<Location>,
     },
-    ReportDecl {
+    // all report variables must be allocated,
+    // but not all alloc variables must be reported
+    AllocDecl {
+        is_report: bool,
         decl: Box<Statement>,
         loc: Option<Location>,
     },
@@ -302,7 +305,7 @@ impl Statement {
             | Statement::Return { loc, .. }
             | Statement::Assign { loc, .. }
             | Statement::SetMap { loc, .. }
-            | Statement::ReportDecl { loc, .. }
+            | Statement::AllocDecl { loc, .. }
             | Statement::Expr { loc, .. } => loc,
         }
     }

src/parser/whamm.pest

@@ -69,26 +69,32 @@ TYPE = _{ TY_U32 | TY_I32 | TY_F32 | TY_U64 | TY_I64 | TY_F64 | TY_BOOL | TY_STR
 // ---- Statements ----
 // ====================
 
-//has to be "_stmt" because of rust rules
+statement = { ( decl_init | fn_call | special_decl | decl | assignment | incrementor | decrementor | ret ) ~ ";" }
+
+// control flow
 if_stmt = { "if" ~ "(" ~ expr ~ ")" ~ block ~ (else_stmt | elif) ? }
 elif = { "elif" ~ "(" ~ expr ~ ")" ~ block ~ (else_stmt | elif) ? }
 else_stmt = { "else" ~  block }
+block = { "{" ~ (if_stmt | statement)* ~ "}" }
+ret = { "return" ~ expr ? }
+
+// calls
 arg = { tuple | expr | val | ternary }
 fn_call = { ID ~ "(" ~ ( arg )? ~ ( "," ~ arg )* ~ ")" }
-block = { "{" ~ (if_stmt | statement)* ~ "}" }
 
+// unop
+incrementor = { (get_map | ID) ~ "++"}
+decrementor = { (get_map | ID) ~ "--"}
 
 // var ops
-ret = { "return" ~ expr ? }
-
-report_declaration = { "report" ~ declaration }
+REPORT = @{ "report" }
+ALLOC = @{ "alloc" }
+var_decorators = { ( REPORT | ALLOC )+ }
+special_decl = { var_decorators ~ decl }
 assignment = { (get_map | ID) ~ "=" ~ (ternary | expr) }
-declaration = { TYPE ~ !RESERVED_KEYWORDS ~ ID }
+decl = { TYPE ~ !RESERVED_KEYWORDS ~ ID }
+decl_init = { (special_decl | decl) ~ "=" ~ (ternary | expr) }
 
-initialize = { TYPE ~ ID ~ "=" ~ (ternary | expr) }
-statement = { ( ( initialize | fn_call | report_declaration | declaration | assignment | incrementor | decrementor | ret ) ~ ";" )+ }
-incrementor = { (get_map | ID) ~ "++"}
-decrementor = { (get_map | ID) ~ "--"}
 // =====================
 // ---- Expressions ----
 // =====================