Merge pull request RustPython#12 from cthulahoops/lambdas

windelbouwman · web-flow · commit ea0e25fd3bb8 · 2018-07-28T13:12:00.000+02:00
Lambdas and pass arguments to functions.
diff --git a/parser/src/ast.rs b/parser/src/ast.rs
@@ -133,6 +133,10 @@ pub enum Expression {
     Identifier {
         name: String,
     },
+    Lambda {
+        args: Vec<String>,
+        body: Box<Expression>,
+    },
     True,
     False,
     None,
diff --git a/parser/src/parser.rs b/parser/src/parser.rs
@@ -150,4 +150,28 @@ mod tests {
             }
         );
     }
+
+    #[test]
+    fn test_parse_lambda() {
+        let source = String::from("lambda x, y: x * y\n"); // lambda(x, y): x * y");
+        let parse_ast = parse_statement(&source);
+        assert_eq!(
+            parse_ast,
+            Ok(ast::Statement::Expression {
+                expression: ast::Expression::Lambda {
+                    args: vec![String::from("x"), String::from("y")],
+                    body:
+                        Box::new(ast::Expression::Binop {
+                            a: Box::new(ast::Expression::Identifier {
+                                name: String::from("x"),
+                            }),
+                            op: ast::Operator::Mult,
+                            b: Box::new(ast::Expression::Identifier {
+                                name: String::from("y"),
+                            })
+                        })
+                    }
+            })
+        )
+    }
 }
diff --git a/parser/src/python.lalrpop b/parser/src/python.lalrpop
@@ -169,7 +169,11 @@ FuncDef: ast::Statement = {
 };
 
 Parameters: Vec<String> = {
-  "(" <a: Comma<Identifier>> ")" => a,
+  "(" <a: TypedArgsList> ")" => a,
+};
+
+TypedArgsList: Vec<String> = {
+  <a: Comma<Identifier>> => a,
 };
 
 ClassDef: ast::Statement = {
@@ -178,8 +182,17 @@ ClassDef: ast::Statement = {
 
 Test: ast::Expression = {
   <e:OrTest> => e,
+  <e:LambdaDef> => e,
 };
 
+LambdaDef: ast::Expression = {
+  "lambda" <p:TypedArgsList> ":" <b:Expression> =>
+	 ast::Expression::Lambda {
+	  args:p,
+	  body:Box::new(b)
+       }
+}
+
 OrTest: ast::Expression = {
   <e:AndTest> => e,
   <e1:OrTest> "or" <e2:AndTest> => ast::Expression::BoolOp { a: Box::new(e1), op: ast::BooleanOperator::Or, b: Box::new(e2) },
@@ -410,6 +423,7 @@ extern {
         "return" => lexer::Tok::Return,
         "while" => lexer::Tok::While,
         "with" => lexer::Tok::With,
+	  "lambda" => lexer::Tok::Lambda,
         "True" => lexer::Tok::True,
         "False" => lexer::Tok::False,
         "None" => lexer::Tok::None,
diff --git a/vm/src/bytecode.rs b/vm/src/bytecode.rs
@@ -17,13 +17,15 @@ use std::fmt;
 pub struct CodeObject {
     pub instructions: Vec<Instruction>,
     pub label_map: HashMap<Label, usize>,
+    pub arg_names: Vec<String>,
 }
 
 impl CodeObject {
-    pub fn new() -> CodeObject {
+    pub fn new(arg_names : Vec<String>) -> CodeObject {
         CodeObject {
             instructions: Vec::new(),
             label_map: HashMap::new(),
+            arg_names: arg_names,
         }
     }
 }
diff --git a/vm/src/compile.rs b/vm/src/compile.rs
@@ -15,7 +15,7 @@ struct Compiler {
 
 pub fn compile(rt: &mut Executor, source: &String, mode: Mode) -> Result<PyObjectRef, String> {
     let mut compiler = Compiler::new();
-    compiler.push_code_object(CodeObject::new());
+    compiler.push_code_object(CodeObject::new(Vec::new()));
     match mode {
         Mode::Exec => match parser::parse_program(source) {
             Ok(ast) => {
@@ -62,7 +62,7 @@ impl Compiler {
     }
 
     fn push_code_object(&mut self, code_object: CodeObject) {
-        self.code_object_stack.push(CodeObject::new());
+        self.code_object_stack.push(CodeObject::new(Vec::new()));
     }
 
     fn pop_code_object(&mut self) -> CodeObject {
@@ -191,7 +191,7 @@ impl Compiler {
             }
             ast::Statement::FunctionDef { name, args, body } => {
                 // Create bytecode for this function:
-                self.code_object_stack.push(CodeObject::new());
+                self.code_object_stack.push(CodeObject::new(args.to_vec()));
                 self.compile_statements(body);
                 let code = self.code_object_stack.pop().unwrap();
                 self.emit(Instruction::LoadConst {
@@ -444,6 +444,22 @@ impl Compiler {
                     name: name.to_string(),
                 });
             }
+            ast::Expression::Lambda { args, body } => {
+                self.code_object_stack.push(CodeObject::new(args.to_vec()));
+                self.compile_expression(body);
+                self.emit(Instruction::ReturnValue);
+                let code = self.code_object_stack.pop().unwrap();
+                self.emit(Instruction::LoadConst {
+                    value: bytecode::Constant::Code { code: code },
+                });
+                self.emit(Instruction::LoadConst {
+                    value: bytecode::Constant::String {
+                        value: String::from("<lambda>"),
+                    },
+                });
+                // Turn code object into function object:
+                self.emit(Instruction::MakeFunction);
+            }
         }
     }
 
diff --git a/vm/src/vm.rs b/vm/src/vm.rs
@@ -10,7 +10,7 @@ use std::ops::Deref;
 
 use super::builtins;
 use super::bytecode;
-use super::frame::{Block, Frame};
+use super::frame::{Block, Frame, copy_code};
 use super::import::import;
 use super::objlist;
 use super::objstr;
@@ -431,7 +431,12 @@ impl VirtualMachine {
         match f.kind {
             PyObjectKind::RustFunction { ref function } => f.call(self, args),
             PyObjectKind::Function { ref code } => {
-                let frame = Frame::new(code.clone(), HashMap::new(), self.new_dict());
+                let mut locals = self.new_dict();
+                let code_object = copy_code(code.clone());
+                for (name, value) in code_object.arg_names.iter().zip(args) {
+                    locals.set_item(name, value);
+                }
+                let frame = Frame::new(code.clone(), HashMap::new(), locals);
                 self.run_frame(frame)
             }
             PyObjectKind::Class { name: _ } => self.new_instance(func_ref.clone(), args),

Original file line number	Diff line number	Diff line change
`@@ -17,13 +17,15 @@ use std::fmt;`
`17`	`17`	`pub struct CodeObject {`
`18`	`18`	`pub instructions: Vec<Instruction>,`
`19`	`19`	`pub label_map: HashMap<Label, usize>,`
	`20`	`+ pub arg_names: Vec<String>,`
`20`	`21`	`}`
`21`	`22`
`22`	`23`	`impl CodeObject {`
`23`		`- pub fn new() -> CodeObject {`
	`24`	`+ pub fn new(arg_names : Vec<String>) -> CodeObject {`
`24`	`25`	`CodeObject {`
`25`	`26`	`instructions: Vec::new(),`
`26`	`27`	`label_map: HashMap::new(),`
	`28`	`+ arg_names: arg_names,`
`27`	`29`	`}`
`28`	`30`	`}`
`29`	`31`	`}`