From 0cbdaefe469e89edcc5b32b19d9c786299410458 Mon Sep 17 00:00:00 2001
From: evomassiny <evomassiny@users.noreply.github.com>
Date: Fri, 1 Nov 2019 05:28:12 +0100
Subject: [PATCH] Fix early return (#197)

Fix early return
---
 src/lib/environment/lexical_environment.rs |  4 ++--
 src/lib/exec.rs                            | 28 +++++++++++++++++++---
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/src/lib/environment/lexical_environment.rs b/src/lib/environment/lexical_environment.rs
index 8d0e6ae920..bfdc95ae9f 100644
--- a/src/lib/environment/lexical_environment.rs
+++ b/src/lib/environment/lexical_environment.rs
@@ -95,8 +95,8 @@ impl LexicalEnvironment {
         self.environment_stack.push_back(env);
     }
 
-    pub fn pop(&mut self) {
-        self.environment_stack.pop_back();
+    pub fn pop(&mut self) -> Option<Environment> {
+        self.environment_stack.pop_back()
     }
 
     pub fn environments(&self) -> impl Iterator<Item = Environment> {
diff --git a/src/lib/exec.rs b/src/lib/exec.rs
index f20aee9970..594fbfb1ff 100644
--- a/src/lib/exec.rs
+++ b/src/lib/exec.rs
@@ -5,7 +5,7 @@ use crate::{
         value::{from_value, to_value, ResultValue, Value, ValueData},
     },
     environment::lexical_environment::{
-        new_declarative_environment, new_function_environment, VariableScope,
+        new_declarative_environment, new_function_environment, EnvironmentType, VariableScope,
     },
     realm::Realm,
     syntax::ast::{
@@ -86,7 +86,6 @@ impl Executor for Interpreter {
                     // early return
                     if self.is_return {
                         obj = val;
-                        self.is_return = false;
                         break;
                     }
                     if e == es.last().expect("unable to get last value") {
@@ -94,7 +93,16 @@ impl Executor for Interpreter {
                     }
                 }
 
-                self.realm.environment.pop();
+                // pop the block env
+                let block_env = self.realm.environment.pop();
+
+                // clear the early return flag `self.is_return`
+                // only when we leave the associated function
+                if let Some(env) = block_env {
+                    if let EnvironmentType::Function = env.deref().borrow().get_environment_type() {
+                        self.is_return = false;
+                    }
+                }
                 Ok(obj)
             }
             ExprDef::Local(ref name) => {
@@ -798,4 +806,18 @@ mod tests {
         "#;
         assert_eq!(exec(boolean_false), String::from("-1"));
     }
+
+    #[test]
+    fn test_early_return() {
+        let early_return = r#"
+        function early_return() {
+            if (true) {
+                return true;
+            }
+            return false;
+        }
+        early_return()
+        "#;
+        assert_eq!(exec(early_return), String::from("true"));
+    }
 }