simple_args() test

src/jit/engine.rs

@@ -52,7 +52,7 @@ impl Engine {
     pub fn translate(&mut self, fn_impl: &ast::FnImpl) -> Result<(), String> {
         let int = self.module.target_config().pointer_type();
         let builder = FunctionBuilder::new(&mut self.ctx.func, &mut self.builder_context);
-        let trans = super::translate::FunctionTranslator { int, builder };
+        let trans = super::translate::FunctionTranslator::new(int, builder);
         trans.translate(fn_impl);
         Ok(())
     }
@@ -74,8 +74,7 @@ mod tests {
         } else {
             panic!("Failed to parse function from test source");
         }
-
-    } 
+    }
 
     #[test]
     fn simple_math() {
@@ -84,4 +83,12 @@ mod tests {
         let code_fn = unsafe { std::mem::transmute::<_, fn() -> i64>(code_ptr) };
         assert_eq!(code_fn(), 9);
     }
+
+    #[test]
+    fn simple_args() {
+        let source = "simple_args(i64 a, i64 b) i64 { a * (a + b) }";
+        let code_ptr = jit_compile(source);
+        let code_fn = unsafe { std::mem::transmute::<_, fn(i64, i64) -> i64>(code_ptr) };
+        assert_eq!(code_fn(2, 3), 10);
+    }
 }

src/jit/translate.rs

@@ -7,9 +7,18 @@ use cranelift::prelude::*;
 pub struct FunctionTranslator<'a> {
     pub int: types::Type,
     pub builder: FunctionBuilder<'a>,
+    pub locals: Vec<(String, Value)>,
 }
 
 impl<'a> FunctionTranslator<'a> {
+    pub fn new(int: types::Type, builder: FunctionBuilder<'a>) -> Self {
+        Self {
+            int,
+            builder,
+            locals: Vec::new(),
+        }
+    }
+
     pub fn translate(mut self, fn_impl: &ast::FnImpl) {
         let signature = &fn_impl.def.signature;
 
@@ -24,7 +33,11 @@ impl<'a> FunctionTranslator<'a> {
             None
         };
 
-        for _ in signature.args.iter() {
+        for ast::FnArg { type_name, .. } in signature.args.iter() {
+            if type_name != &"i64" {
+                unimplemented!("Non-i64 function arg types are unimplemented");
+            }
+
             self.builder
                 .func
                 .signature
@@ -46,6 +59,13 @@ impl<'a> FunctionTranslator<'a> {
         self.builder.switch_to_block(entry_block);
         self.builder.seal_block(entry_block);
 
+        for (i, arg) in signature.args.iter().enumerate() {
+            if let Some(name) = arg.name {
+                let val = self.builder.block_params(entry_block)[i];
+                self.locals.push((name.to_string(), val));
+            }
+        }
+
         if fn_impl.body.statements.len() > 0 {
             unimplemented!("Function body statements are unimplemented");
         }
@@ -62,6 +82,9 @@ impl<'a> FunctionTranslator<'a> {
     pub fn translate_expr(&mut self, expr: &ast::Expr) -> Value {
         use ast::Expr::*;
         match expr {
+            Local(name) => self
+                .get_local(name)
+                .expect(&format!("Unrecognized local {}", name)),
             Literal(ast::Literal::DecimalInteger(literal)) => {
                 // TODO parse integers while building ast so that codegen doesn't have to
                 let val: i64 = literal.parse().unwrap();
@@ -89,4 +112,8 @@ impl<'a> FunctionTranslator<'a> {
             _ => unimplemented!("Binary operation: {:#?}", op),
         }
     }
+
+    pub fn get_local(&self, name: &str) -> Option<Value> {
+        self.locals.iter().rev().find(|l| &l.0 == name).map(|l| l.1)
+    }
 }