From 36d8147e8e39566f5c1f242a6b5f85e4413e4267 Mon Sep 17 00:00:00 2001
From: Windel Bouwman <windel@windel.nl>
Date: Sun, 11 Aug 2019 20:20:15 +0200
Subject: [PATCH] Improve the situation regarding boolean operations.

---
 bytecode/src/bytecode.rs                    |  18 +-
 compiler/src/compile.rs                     | 185 ++++++++++++--------
 compiler/src/symboltable.rs                 |   5 +-
 parser/src/ast.rs                           |   3 +-
 parser/src/python.lalrpop                   |  30 +++-
 tests/snippets/short_circuit_evaluations.py |  15 ++
 vm/src/frame.rs                             |  24 ++-
 vm/src/stdlib/ast.rs                        |   7 +-
 8 files changed, 191 insertions(+), 96 deletions(-)
 create mode 100644 tests/snippets/short_circuit_evaluations.py

diff --git a/bytecode/src/bytecode.rs b/bytecode/src/bytecode.rs
index 0d5008b8393..2c3d4316364 100644
--- a/bytecode/src/bytecode.rs
+++ b/bytecode/src/bytecode.rs
@@ -138,12 +138,24 @@ pub enum Instruction {
     Jump {
         target: Label,
     },
-    JumpIf {
+    /// Pop the top of the stack, and jump if this value is true.
+    JumpIfTrue {
         target: Label,
     },
+    /// Pop the top of the stack, and jump if this value is false.
     JumpIfFalse {
         target: Label,
     },
+    /// Peek at the top of the stack, and jump if this value is true.
+    /// Otherwise, pop top of stack.
+    JumpIfTrueOrPop {
+        target: Label,
+    },
+    /// Peek at the top of the stack, and jump if this value is false.
+    /// Otherwise, pop top of stack.
+    JumpIfFalseOrPop {
+        target: Label,
+    },
     MakeFunction {
         flags: FunctionOpArg,
     },
@@ -411,8 +423,10 @@ impl Instruction {
             Continue => w!(Continue),
             Break => w!(Break),
             Jump { target } => w!(Jump, label_map[target]),
-            JumpIf { target } => w!(JumpIf, label_map[target]),
+            JumpIfTrue { target } => w!(JumpIfTrue, label_map[target]),
             JumpIfFalse { target } => w!(JumpIfFalse, label_map[target]),
+            JumpIfTrueOrPop { target } => w!(JumpIfTrueOrPop, label_map[target]),
+            JumpIfFalseOrPop { target } => w!(JumpIfFalseOrPop, label_map[target]),
             MakeFunction { flags } => w!(MakeFunction, format!("{:?}", flags)),
             CallFunction { typ } => w!(CallFunction, format!("{:?}", typ)),
             ForIter { target } => w!(ForIter, label_map[target]),
diff --git a/compiler/src/compile.rs b/compiler/src/compile.rs
index a13d3617e83..a0f20c1e08e 100644
--- a/compiler/src/compile.rs
+++ b/compiler/src/compile.rs
@@ -15,6 +15,7 @@ use rustpython_parser::{ast, parser};
 
 type BasicOutputStream = PeepholeOptimizer<CodeObjectStream>;
 
+/// Main structure holding the state of compilation.
 struct Compiler<O: OutputStream = BasicOutputStream> {
     output_stack: Vec<O>,
     scope_stack: Vec<SymbolScope>,
@@ -107,12 +108,6 @@ pub enum Mode {
     Single,
 }
 
-#[derive(Clone, Copy)]
-enum EvalContext {
-    Statement,
-    Expression,
-}
-
 pub(crate) type Label = usize;
 
 impl<O> Default for Compiler<O>
@@ -350,14 +345,14 @@ impl<O: OutputStream> Compiler<O> {
                 match orelse {
                     None => {
                         // Only if:
-                        self.compile_test(test, None, Some(end_label), EvalContext::Statement)?;
+                        self.compile_jump_if(test, false, end_label)?;
                         self.compile_statements(body)?;
                         self.set_label(end_label);
                     }
                     Some(statements) => {
                         // if - else:
                         let else_label = self.new_label();
-                        self.compile_test(test, None, Some(else_label), EvalContext::Statement)?;
+                        self.compile_jump_if(test, false, else_label)?;
                         self.compile_statements(body)?;
                         self.emit(Instruction::Jump { target: end_label });
 
@@ -459,7 +454,7 @@ impl<O: OutputStream> Compiler<O> {
                 // if some flag, ignore all assert statements!
                 if self.optimize == 0 {
                     let end_label = self.new_label();
-                    self.compile_test(test, Some(end_label), None, EvalContext::Statement)?;
+                    self.compile_jump_if(test, true, end_label)?;
                     self.emit(Instruction::LoadName {
                         name: String::from("AssertionError"),
                         scope: bytecode::NameScope::Local,
@@ -1006,7 +1001,7 @@ impl<O: OutputStream> Compiler<O> {
 
         self.set_label(start_label);
 
-        self.compile_test(test, None, Some(else_label), EvalContext::Statement)?;
+        self.compile_jump_if(test, false, else_label)?;
 
         let was_in_loop = self.in_loop;
         self.in_loop = true;
@@ -1118,12 +1113,9 @@ impl<O: OutputStream> Compiler<O> {
             });
 
             // if comparison result is false, we break with this value; if true, try the next one.
-            // (CPython compresses these three opcodes into JUMP_IF_FALSE_OR_POP)
-            self.emit(Instruction::Duplicate);
-            self.emit(Instruction::JumpIfFalse {
+            self.emit(Instruction::JumpIfFalseOrPop {
                 target: break_label,
             });
-            self.emit(Instruction::Pop);
         }
 
         // handle the last comparison
@@ -1256,66 +1248,120 @@ impl<O: OutputStream> Compiler<O> {
         self.emit(Instruction::BinaryOperation { op: i, inplace });
     }
 
-    fn compile_test(
+    /// Implement boolean short circuit evaluation logic.
+    /// https://en.wikipedia.org/wiki/Short-circuit_evaluation
+    ///
+    /// This means, in a boolean statement 'x and y' the variable y will
+    /// not be evaluated when x is false.
+    ///
+    /// The idea is to jump to a label if the expression is either true or false
+    /// (indicated by the condition parameter).
+    fn compile_jump_if(
         &mut self,
         expression: &ast::Expression,
-        true_label: Option<Label>,
-        false_label: Option<Label>,
-        context: EvalContext,
+        condition: bool,
+        target_label: Label,
     ) -> Result<(), CompileError> {
         // Compile expression for test, and jump to label if false
         match &expression.node {
-            ast::ExpressionType::BoolOp { a, op, b } => match op {
-                ast::BooleanOperator::And => {
-                    let f = false_label.unwrap_or_else(|| self.new_label());
-                    self.compile_test(a, None, Some(f), context)?;
-                    self.compile_test(b, true_label, false_label, context)?;
-                    if false_label.is_none() {
-                        self.set_label(f);
+            ast::ExpressionType::BoolOp { op, values } => {
+                match op {
+                    ast::BooleanOperator::And => {
+                        if condition {
+                            // If all values are true.
+                            let end_label = self.new_label();
+                            let (last_value, values) = values.split_last().unwrap();
+
+                            // If any of the values is false, we can short-circuit.
+                            for value in values {
+                                self.compile_jump_if(value, false, end_label)?;
+                            }
+
+                            // It depends upon the last value now: will it be true?
+                            self.compile_jump_if(last_value, true, target_label)?;
+                            self.set_label(end_label);
+                        } else {
+                            // If any value is false, the whole condition is false.
+                            for value in values {
+                                self.compile_jump_if(value, false, target_label)?;
+                            }
+                        }
                     }
-                }
-                ast::BooleanOperator::Or => {
-                    let t = true_label.unwrap_or_else(|| self.new_label());
-                    self.compile_test(a, Some(t), None, context)?;
-                    self.compile_test(b, true_label, false_label, context)?;
-                    if true_label.is_none() {
-                        self.set_label(t);
+                    ast::BooleanOperator::Or => {
+                        if condition {
+                            // If any of the values is true.
+                            for value in values {
+                                self.compile_jump_if(value, true, target_label)?;
+                            }
+                        } else {
+                            // If all of the values are false.
+                            let end_label = self.new_label();
+                            let (last_value, values) = values.split_last().unwrap();
+
+                            // If any value is true, we can short-circuit:
+                            for value in values {
+                                self.compile_jump_if(value, true, end_label)?;
+                            }
+
+                            // It all depends upon the last value now!
+                            self.compile_jump_if(last_value, false, target_label)?;
+                            self.set_label(end_label);
+                        }
                     }
                 }
-            },
+            }
+            ast::ExpressionType::Unop {
+                op: ast::UnaryOperator::Not,
+                a,
+            } => {
+                self.compile_jump_if(a, !condition, target_label)?;
+            }
             _ => {
+                // Fall back case which always will work!
                 self.compile_expression(expression)?;
-                match context {
-                    EvalContext::Statement => {
-                        if let Some(true_label) = true_label {
-                            self.emit(Instruction::JumpIf { target: true_label });
-                        }
-                        if let Some(false_label) = false_label {
-                            self.emit(Instruction::JumpIfFalse {
-                                target: false_label,
-                            });
-                        }
-                    }
-                    EvalContext::Expression => {
-                        if let Some(true_label) = true_label {
-                            self.emit(Instruction::Duplicate);
-                            self.emit(Instruction::JumpIf { target: true_label });
-                            self.emit(Instruction::Pop);
-                        }
-                        if let Some(false_label) = false_label {
-                            self.emit(Instruction::Duplicate);
-                            self.emit(Instruction::JumpIfFalse {
-                                target: false_label,
-                            });
-                            self.emit(Instruction::Pop);
-                        }
-                    }
+                if condition {
+                    self.emit(Instruction::JumpIfTrue {
+                        target: target_label,
+                    });
+                } else {
+                    self.emit(Instruction::JumpIfFalse {
+                        target: target_label,
+                    });
                 }
             }
         }
         Ok(())
     }
 
+    /// Compile a boolean operation as an expression.
+    /// This means, that the last value remains on the stack.
+    fn compile_bool_op(
+        &mut self,
+        op: &ast::BooleanOperator,
+        values: &[ast::Expression],
+    ) -> Result<(), CompileError> {
+        let end_label = self.new_label();
+
+        let (last_value, values) = values.split_last().unwrap();
+        for value in values {
+            self.compile_expression(value)?;
+
+            match op {
+                ast::BooleanOperator::And => {
+                    self.emit(Instruction::JumpIfFalseOrPop { target: end_label });
+                }
+                ast::BooleanOperator::Or => {
+                    self.emit(Instruction::JumpIfTrueOrPop { target: end_label });
+                }
+            }
+        }
+
+        // If all values did not qualify, take the value of the last value:
+        self.compile_expression(last_value)?;
+        self.set_label(end_label);
+        Ok(())
+    }
+
     fn compile_expression(&mut self, expression: &ast::Expression) -> Result<(), CompileError> {
         trace!("Compiling {:?}", expression);
         self.set_source_location(&expression.location);
@@ -1327,12 +1373,7 @@ impl<O: OutputStream> Compiler<O> {
                 args,
                 keywords,
             } => self.compile_call(function, args, keywords)?,
-            BoolOp { .. } => self.compile_test(
-                expression,
-                Option::None,
-                Option::None,
-                EvalContext::Expression,
-            )?,
+            BoolOp { op, values } => self.compile_bool_op(op, values)?,
             Binop { a, op, b } => {
                 self.compile_expression(a)?;
                 self.compile_expression(b)?;
@@ -1527,8 +1568,7 @@ impl<O: OutputStream> Compiler<O> {
             IfExpression { test, body, orelse } => {
                 let no_label = self.new_label();
                 let end_label = self.new_label();
-                self.compile_test(test, Option::None, Option::None, EvalContext::Expression)?;
-                self.emit(Instruction::JumpIfFalse { target: no_label });
+                self.compile_jump_if(test, false, no_label)?;
                 // True case
                 self.compile_expression(body)?;
                 self.emit(Instruction::Jump { target: end_label });
@@ -1745,12 +1785,7 @@ impl<O: OutputStream> Compiler<O> {
 
             // Now evaluate the ifs:
             for if_condition in &generator.ifs {
-                self.compile_test(
-                    if_condition,
-                    None,
-                    Some(start_label),
-                    EvalContext::Statement,
-                )?
+                self.compile_jump_if(if_condition, false, start_label)?
             }
         }
 
@@ -1988,11 +2023,11 @@ mod tests {
                 LoadConst {
                     value: Boolean { value: true }
                 },
-                JumpIf { target: 1 },
+                JumpIfTrue { target: 1 },
                 LoadConst {
                     value: Boolean { value: false }
                 },
-                JumpIf { target: 1 },
+                JumpIfTrue { target: 1 },
                 LoadConst {
                     value: Boolean { value: false }
                 },
@@ -2042,7 +2077,7 @@ mod tests {
                 LoadConst {
                     value: Boolean { value: false }
                 },
-                JumpIf { target: 1 },
+                JumpIfTrue { target: 1 },
                 LoadConst {
                     value: Boolean { value: false }
                 },
diff --git a/compiler/src/symboltable.rs b/compiler/src/symboltable.rs
index c1db24bf325..3792bf4f9d8 100644
--- a/compiler/src/symboltable.rs
+++ b/compiler/src/symboltable.rs
@@ -404,9 +404,8 @@ impl SymbolTableBuilder {
                 self.scan_expression(a)?;
                 self.scan_expression(b)?;
             }
-            BoolOp { a, b, .. } => {
-                self.scan_expression(a)?;
-                self.scan_expression(b)?;
+            BoolOp { values, .. } => {
+                self.scan_expressions(values)?;
             }
             Compare { vals, .. } => {
                 self.scan_expressions(vals)?;
diff --git a/parser/src/ast.rs b/parser/src/ast.rs
index b11e0f9c62c..dc7b2a14d8f 100644
--- a/parser/src/ast.rs
+++ b/parser/src/ast.rs
@@ -148,9 +148,8 @@ pub type Expression = Located<ExpressionType>;
 #[derive(Debug, PartialEq)]
 pub enum ExpressionType {
     BoolOp {
-        a: Box<Expression>,
         op: BooleanOperator,
-        b: Box<Expression>,
+        values: Vec<Expression>,
     },
     Binop {
         a: Box<Expression>,
diff --git a/parser/src/python.lalrpop b/parser/src/python.lalrpop
index 6af9fea61c2..15cea7fb5ab 100644
--- a/parser/src/python.lalrpop
+++ b/parser/src/python.lalrpop
@@ -659,18 +659,32 @@ LambdaDef: ast::Expression = {
 }
 
 OrTest: ast::Expression = {
-    AndTest,
-    <e1:OrTest> <location:@L> "or" <e2:AndTest> => ast::Expression {
-        location,
-        node: ast::ExpressionType::BoolOp { a: Box::new(e1), op: ast::BooleanOperator::Or, b: Box::new(e2) }
+    <e1:AndTest> <location:@L> <e2:("or" AndTest)*> => {
+        if e2.is_empty() {
+            e1
+        } else {
+            let mut values = vec![e1];
+            values.extend(e2.into_iter().map(|e| e.1));
+            ast::Expression {
+                location,
+                node: ast::ExpressionType::BoolOp { op: ast::BooleanOperator::Or, values }
+            }
+        }
     },
 };
 
 AndTest: ast::Expression = {
-    NotTest,
-    <e1:AndTest> <location:@L> "and" <e2:NotTest> => ast::Expression {
-        location,
-        node: ast::ExpressionType::BoolOp { a: Box::new(e1), op: ast::BooleanOperator::And, b: Box::new(e2) }
+    <e1:NotTest> <location:@L> <e2:("and" NotTest)*> => {
+        if e2.is_empty() {
+            e1
+        } else {
+            let mut values = vec![e1];
+            values.extend(e2.into_iter().map(|e| e.1));
+            ast::Expression {
+                location,
+                node: ast::ExpressionType::BoolOp { op: ast::BooleanOperator::And, values }
+            }
+        }
     },
 };
 
diff --git a/tests/snippets/short_circuit_evaluations.py b/tests/snippets/short_circuit_evaluations.py
new file mode 100644
index 00000000000..f83a4137c90
--- /dev/null
+++ b/tests/snippets/short_circuit_evaluations.py
@@ -0,0 +1,15 @@
+
+# Test various cases of short circuit evaluation:
+
+run = 1
+timeTaken = 33
+r = (11, 22, run, run != 1 and "s" or "", timeTaken)
+print(r)
+assert r == (11, 22, 1, '', 33)
+
+
+run = 0
+r = (11, 22, run, run != 1 and "s" or "", timeTaken)
+print(r)
+assert r == (11, 22, 0, 's', 33)
+
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
index 4cbb022a5d8..f338b64135e 100644
--- a/vm/src/frame.rs
+++ b/vm/src/frame.rs
@@ -496,7 +496,7 @@ impl Frame {
                 self.jump(*target);
                 Ok(None)
             }
-            bytecode::Instruction::JumpIf { target } => {
+            bytecode::Instruction::JumpIfTrue { target } => {
                 let obj = self.pop_value();
                 let value = objbool::boolval(vm, obj)?;
                 if value {
@@ -514,6 +514,28 @@ impl Frame {
                 Ok(None)
             }
 
+            bytecode::Instruction::JumpIfTrueOrPop { target } => {
+                let obj = self.last_value();
+                let value = objbool::boolval(vm, obj)?;
+                if value {
+                    self.jump(*target);
+                } else {
+                    self.pop_value();
+                }
+                Ok(None)
+            }
+
+            bytecode::Instruction::JumpIfFalseOrPop { target } => {
+                let obj = self.last_value();
+                let value = objbool::boolval(vm, obj)?;
+                if !value {
+                    self.jump(*target);
+                } else {
+                    self.pop_value();
+                }
+                Ok(None)
+            }
+
             bytecode::Instruction::Raise { argc } => {
                 let cause = match argc {
                     2 => self.get_exception(vm, true)?,
diff --git a/vm/src/stdlib/ast.rs b/vm/src/stdlib/ast.rs
index 9543d5ce13c..19a4e084ae2 100644
--- a/vm/src/stdlib/ast.rs
+++ b/vm/src/stdlib/ast.rs
@@ -332,11 +332,8 @@ fn expression_to_ast(vm: &VirtualMachine, expression: &ast::Expression) -> PyRes
                 operand => expression_to_ast(vm, a)?,
             })
         }
-        BoolOp { a, op, b } => {
-            // Attach values:
-            let py_a = expression_to_ast(vm, a)?.into_object();
-            let py_b = expression_to_ast(vm, b)?.into_object();
-            let py_values = vm.ctx.new_tuple(vec![py_a, py_b]);
+        BoolOp { op, values } => {
+            let py_values = expressions_to_ast(vm, values)?;
 
             let str_op = match op {
                 ast::BooleanOperator::And => "And",