Merge pull request #101 from NLnetLabs/testing-in-roto-scripts

Add some basic testing infrastructure for Roto scripts

Author: tertsdiepraam

examples/simple.roto

@@ -1,12 +1,25 @@
+function is_zero(x: IpAddr) -> bool {
+    x == 0.0.0.0
+}
+
 filter-map main(x: IpAddr) {
-    define {
-        y = 0.0.0.0;
+    if is_zero(x) {
+        accept
+    } else {
+        reject
+    }
+}
+
+test is_zero_true {
+    if is_zero(1.1.1.1) {
+        reject
     }
-    apply {
-        if x == y {
-            accept
-        } else {
-            reject
-        }
+    accept
+}
+
+test is_zero_false {
+    if not is_zero(0.0.0.0) {
+        reject
     }
+    accept
 }

examples/simple.rs

@@ -40,5 +40,15 @@ fn main() -> Result<(), roto::RotoReport> {
     let res = func.call(&mut (), "1.1.1.1".parse().unwrap());
     println!("main(1.1.1.1) = {res:?}");
 
+    let is_zero = compiled
+        .get_function::<(), (IpAddr,), bool>("is_zero")
+        .unwrap();
+
+    let res = is_zero.call(&mut (), "0.0.0.0".parse().unwrap());
+    println!("is_zero(0.0.0.0) = {res:?}");
+
+    println!();
+    let _ = compiled.run_tests(());
+
     Ok(())
 }

src/ast.rs

@@ -23,6 +23,7 @@ pub enum Declaration {
     OutputStream(OutputStream),
     Record(RecordTypeDeclaration),
     Function(FunctionDeclaration),
+    Test(Test),
 }
 
 #[derive(Clone, Debug)]
@@ -46,7 +47,7 @@ pub struct FilterMap {
     pub filter_type: FilterType,
     pub ident: Meta<Identifier>,
     pub params: Meta<Params>,
-    pub block: Meta<Block>,
+    pub body: Meta<Block>,
 }
 
 /// A function declaration, including the [`Block`] forming its definition
@@ -58,6 +59,12 @@ pub struct FunctionDeclaration {
     pub body: Meta<Block>,
 }
 
+#[derive(Clone, Debug)]
+pub struct Test {
+    pub ident: Meta<Identifier>,
+    pub body: Meta<Block>,
+}
+
 /// A block of multiple statements
 #[derive(Clone, Debug)]
 pub struct Block {

src/codegen/check.rs

@@ -162,18 +162,6 @@ fn check_roto_type(
     }
 }
 
-pub fn return_type_by_ref(registry: &TypeRegistry, rust_ty: TypeId) -> bool {
-    let Some(rust_ty) = registry.get(rust_ty) else {
-        return false;
-    };
-
-    #[allow(clippy::match_like_matches_macro)]
-    match rust_ty.description {
-        TypeDescription::Verdict(_, _) => true,
-        _ => todo!(),
-    }
-}
-
 /// Parameters of a Roto function
 ///
 /// This trait allows for checking the types against Roto types and converting

src/codegen/mod.rs

@@ -16,17 +16,12 @@ use crate::{
         value::IrType,
         IrFunction,
     },
-    runtime::{
-        context::ContextDescription,
-        ty::{Reflect, GLOBAL_TYPE_REGISTRY},
-        RuntimeConstant,
-    },
+    runtime::{context::ContextDescription, ty::Reflect, RuntimeConstant},
     typechecker::{info::TypeInfo, scope::ScopeRef, types},
-    IrValue,
+    IrValue, Verdict,
 };
 use check::{
-    check_roto_type_reflect, return_type_by_ref, FunctionRetrievalError,
-    RotoParams, TypeMismatch,
+    check_roto_type_reflect, FunctionRetrievalError, RotoParams, TypeMismatch,
 };
 use cranelift::{
     codegen::{
@@ -177,6 +172,7 @@ call_impl!(A, B, C, D, E, F, G);
 pub struct FunctionInfo {
     id: FuncId,
     signature: types::Signature,
+    return_by_ref: bool,
 }
 
 struct ModuleBuilder {
@@ -377,6 +373,7 @@ impl ModuleBuilder {
 
         let mut sig = self.inner.make_signature();
 
+        // This is the parameter for the context
         sig.params
             .push(AbiParam::new(self.cranelift_type(&IrType::Pointer)));
 
@@ -406,6 +403,7 @@ impl ModuleBuilder {
             name.to_string(),
             FunctionInfo {
                 id: func_id,
+                return_by_ref: ir_signature.return_ptr,
                 signature: signature.clone(),
             },
         );
@@ -1072,6 +1070,52 @@ impl<'c> FuncGen<'c> {
 }
 
 impl Module {
+    pub fn run_tests<Ctx: 'static>(
+        &mut self,
+        mut ctx: Ctx,
+    ) -> Result<(), ()> {
+        let tests: Vec<_> = self
+            .functions
+            .keys()
+            .filter(|x| x.starts_with("test#"))
+            .map(Clone::clone)
+            .collect();
+
+        let total = tests.len();
+        let total_width = total.to_string().len();
+        let mut successes = 0;
+        let mut failures = 0;
+
+        for (n, test) in tests.into_iter().enumerate() {
+            let n = n + 1;
+            let test_display = test.strip_prefix("test#").unwrap();
+            print!("Test {n:>total_width$} / {total}: {test_display}... ");
+            let test_fn = self
+                .get_function::<Ctx, (), Verdict<(), ()>>(&test)
+                .unwrap();
+
+            match test_fn.call(&mut ctx) {
+                Verdict::Accept(()) => {
+                    successes += 1;
+                    println!("\x1B[92mok\x1B[m");
+                }
+                Verdict::Reject(()) => {
+                    failures += 1;
+                    println!("\x1B[91mfail\x1B[m");
+                }
+            }
+        }
+        println!(
+            "Ran {total} tests, {successes} succeeded, {failures} failed"
+        );
+
+        if failures == 0 {
+            Result::Ok(())
+        } else {
+            Result::Err(())
+        }
+    }
+
     pub fn get_function<Ctx: 'static, Params: RotoParams, Return: Reflect>(
         &mut self,
         name: &str,
@@ -1109,14 +1153,10 @@ impl Module {
             )
         })?;
 
-        let registry = GLOBAL_TYPE_REGISTRY.lock().unwrap();
-        let return_by_ref =
-            return_type_by_ref(&registry, TypeId::of::<Return>());
-
         let func_ptr = self.inner.0.get_finalized_function(id);
         Ok(TypedFunc {
             func: func_ptr,
-            return_by_ref,
+            return_by_ref: function_info.return_by_ref,
             _module: self.inner.clone(),
             _ty: PhantomData,
         })

src/codegen/tests.rs

@@ -905,6 +905,69 @@ fn use_context() {
     assert_eq!(output, Verdict::Accept(11));
 }
 
+#[test]
+fn use_a_roto_function() {
+    let s = src!(
+        "
+        function double(x: i32) -> i32 {
+            2 * x
+        }"
+    );
+
+    let mut p = compile(s);
+    let f = p.get_function::<(), (i32,), i32>("double").unwrap();
+    let output = f.call(&mut (), 2);
+    assert_eq!(output, 4);
+
+    let output = f.call(&mut (), 16);
+    assert_eq!(output, 32);
+}
+
+#[test]
+fn use_a_test() {
+    let s = src!(
+        "
+        function double(x: i32) -> i32 {
+            x # oops! not correct
+        }
+        
+        test check_double {
+            if double(4) != 8 {
+                reject;
+            }
+            if double(16) != 32 {
+                reject;
+            }
+            accept
+        }
+        "
+    );
+
+    let mut p = compile(s);
+    p.run_tests(()).unwrap_err();
+
+    let s = src!(
+        "
+        function double(x: i32) -> i32 {
+            2 * x
+        }
+        
+        test check_double {
+            if double(4) != 8 {
+                reject;
+            }
+            if double(16) != 32 {
+                reject;
+            }
+            accept
+        }
+        "
+    );
+
+    let mut p = compile(s);
+    p.run_tests(()).unwrap();
+}
+
 #[test]
 fn string() {
     let s = src!(

src/lower/ir.rs

@@ -367,8 +367,8 @@ impl<'a> IrPrinter<'a> {
             } => format!(
                 "{}: {ty} = {}({}, {})",
                 self.var(to),
-                self.operand(ctx),
                 self.ident(func),
+                self.operand(ctx),
                 args.iter()
                     .map(|a| format!(
                         "{} = {}",

src/lower/match_expr.rs

@@ -213,12 +213,13 @@ impl Lowerer<'_> {
             let ty = self.type_info.type_of(&arm.body);
             let val = self.block(&arm.body);
             if let Some(val) = val {
-                let ty = self.lower_type(&ty);
-                self.add(Instruction::Assign {
-                    to: out.clone(),
-                    val,
-                    ty,
-                });
+                if let Some(ty) = self.lower_type(&ty) {
+                    self.add(Instruction::Assign {
+                        to: out.clone(),
+                        val,
+                        ty,
+                    });
+                }
                 any_assigned = true;
             }
             self.add(Instruction::Jump(continue_lbl));
@@ -267,15 +268,17 @@ impl Lowerer<'_> {
                     1 + self.type_info.padding_of(&ty, 1, self.runtime);
                 let val =
                     self.read_field(examinee.clone().into(), offset, &ty);
-                let ty = self.lower_type(&ty);
-                self.add(Instruction::Assign {
-                    to: Var {
-                        scope,
-                        kind: VarKind::Explicit(ident),
-                    },
-                    val,
-                    ty,
-                });
+                if let Some(val) = val {
+                    let ty = self.lower_type(&ty).unwrap();
+                    self.add(Instruction::Assign {
+                        to: Var {
+                            scope,
+                            kind: VarKind::Explicit(ident),
+                        },
+                        val,
+                        ty,
+                    });
+                }
             }
 
             let ident = Identifier::from(format!("guard_{}", i + 1));