From 783a2e3dd0dc56aef37e27cdb5871c7decaa3053 Mon Sep 17 00:00:00 2001
From: Emmanuel Benoit <tseeker@nocternity.net>
Date: Wed, 11 Jan 2023 07:46:13 +0100
Subject: [PATCH] Interpreter - Special case for initializer methods

---
 src/interpreter/environment.rs   |  9 +++++++++
 src/interpreter/functions.rs     | 17 ++++++++++++++---
 src/interpreter/interpretable.rs |  4 +++-
 3 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/src/interpreter/environment.rs b/src/interpreter/environment.rs
index b493d46..97e1e46 100644
--- a/src/interpreter/environment.rs
+++ b/src/interpreter/environment.rs
@@ -122,6 +122,15 @@ impl Environment {
         self.values.insert(name.to_owned(), Some(value));
     }
 
+    /// Read a variable from an environment directly. Panics if the symbol
+    /// does not exist.
+    pub fn read(&self, name: &str) -> Value {
+        match self.values.get(name) {
+            Some(Some(v)) => v.clone(),
+            _ => panic!("Symbol {name} does not exist"),
+        }
+    }
+
     /// Read an ancestor from the chain of enclosing environments.
     fn ancestor(&self, distance: usize) -> EnvironmentRef {
         let mut ancestor = self.enclosing.clone().expect("ancestor() called at root");
diff --git a/src/interpreter/functions.rs b/src/interpreter/functions.rs
index 23d98e2..7fd07dc 100644
--- a/src/interpreter/functions.rs
+++ b/src/interpreter/functions.rs
@@ -1,4 +1,4 @@
-use std::{fmt::Display, cell::RefMut};
+use std::{cell::RefMut, fmt::Display};
 
 use itertools::izip;
 
@@ -15,6 +15,7 @@ pub struct Function {
     params: Vec<Token>,
     body: Vec<ast::StmtNode>,
     env: EnvironmentRef,
+    is_initializer: bool,
 }
 
 impl Function {
@@ -23,12 +24,14 @@ impl Function {
         params: &[Token],
         body: &[ast::StmtNode],
         environment: EnvironmentRef,
+        is_initializer: bool,
     ) -> Self {
         Self {
             name: name.cloned(),
             params: params.to_owned(),
             body: body.to_owned(),
             env: environment,
+            is_initializer,
         }
     }
 
@@ -38,6 +41,7 @@ impl Function {
             params: self.params.clone(),
             body: self.body.clone(),
             env: Environment::create_child(&self.env),
+            is_initializer: self.is_initializer,
         }
     }
 
@@ -71,11 +75,18 @@ impl Callable for Function {
             let result = stmt.interpret(&mut child)?;
             match result {
                 InterpreterFlowControl::Result(_) => (),
-                InterpreterFlowControl::Return(v) => return Ok(v),
+                InterpreterFlowControl::Return(v) if !self.is_initializer => return Ok(v),
+                InterpreterFlowControl::Return(_) => {
+                    return Ok(itpr_state.environment.borrow().read("this"))
+                }
                 _ => panic!("unexpected flow control {:?}", result),
             }
         }
-        Ok(Value::Nil)
+        if self.is_initializer {
+            Ok(itpr_state.environment.borrow().read("this"))
+        } else {
+            Ok(Value::Nil)
+        }
     }
 }
 
diff --git a/src/interpreter/interpretable.rs b/src/interpreter/interpretable.rs
index e781ad7..8b39932 100644
--- a/src/interpreter/interpretable.rs
+++ b/src/interpreter/interpretable.rs
@@ -208,6 +208,7 @@ impl StmtNode {
                         &method.params,
                         &method.body,
                         es.environment.clone(),
+                        method.name.lexeme == "init",
                     ),
                 )
             })
@@ -226,6 +227,7 @@ impl StmtNode {
             &decl.params,
             &decl.body,
             es.environment.clone(),
+            false,
         );
         es.environment
             .borrow_mut()
@@ -357,7 +359,7 @@ impl Interpretable for ExprNode {
                 arguments,
             } => self.on_call(es, callee, right_paren, arguments),
             ExprNode::Lambda { params, body } => {
-                let lambda = Function::new(None, params, body, es.environment.clone());
+                let lambda = Function::new(None, params, body, es.environment.clone(), false);
                 Ok(Value::from(lambda).into())
             }
             ExprNode::Get(get_expr) => self.on_get_expression(es, get_expr),