From 753c0337bc3f8ed870721621e01ed6dfddf7a759 Mon Sep 17 00:00:00 2001
From: Jeong Yunwon <jeong@youknowone.org>
Date: Mon, 2 May 2022 01:51:30 +0900
Subject: [PATCH] compile_single checks last expr out of loop

---
 src/compile.rs | 29 ++++++++++++++---------------
 1 file changed, 14 insertions(+), 15 deletions(-)

diff --git a/src/compile.rs b/src/compile.rs
index 658e12b..52852af 100644
--- a/src/compile.rs
+++ b/src/compile.rs
@@ -337,31 +337,30 @@ impl Compiler {
     ) -> CompileResult<()> {
         self.symbol_table_stack.push(symbol_table);
 
-        let mut emitted_return = false;
-
-        for (i, statement) in body.iter().enumerate() {
-            let is_last = i == body.len() - 1;
+        let (last, body) = if let Some(splited) = body.split_last() {
+            splited
+        } else {
+            return Ok(());
+        };
 
+        for statement in body {
             if let ast::StmtKind::Expr { value } = &statement.node {
                 self.compile_expression(value)?;
-
-                if is_last {
-                    self.emit(Instruction::Duplicate);
-                    self.emit(Instruction::PrintExpr);
-                    self.emit(Instruction::ReturnValue);
-                    emitted_return = true;
-                } else {
-                    self.emit(Instruction::PrintExpr);
-                }
+                self.emit(Instruction::PrintExpr);
             } else {
                 self.compile_statement(statement)?;
             }
         }
 
-        if !emitted_return {
+        if let ast::StmtKind::Expr { value } = &last.node {
+            self.compile_expression(value)?;
+            self.emit(Instruction::Duplicate);
+            self.emit(Instruction::PrintExpr);
+        } else {
+            self.compile_statement(last)?;
             self.emit_constant(ConstantData::None);
-            self.emit(Instruction::ReturnValue);
         }
+        self.emit(Instruction::ReturnValue);
 
         Ok(())
     }