Handle initialisers

Author: penelopeysm

src/compiler.cpp

@@ -130,7 +130,8 @@ Parser::Parser(std::unique_ptr<Scanner> scanner, ObjFunction* fnptr, GC& gc)
     : scanner(std::move(scanner)), current(SENTINEL_EOF),
       previous(SENTINEL_EOF), errmsg(std::nullopt), gc(gc),
       compiler(
-          std::make_unique<Compiler>(fnptr, nullptr, FunctionType::TOPLEVEL)) {}
+          std::make_unique<Compiler>(fnptr, nullptr, FunctionType::TOPLEVEL)),
+      current_class(nullptr) {}
 
 void Parser::advance() {
   previous = current;
@@ -151,9 +152,12 @@ void Parser::function(bool is_class_method) {
   size_t arity = 0;
   // TODO: This copies the string. Do we need to?
   auto new_fnptr = gc.alloc<ObjFunction>(fn_name, arity);
-  auto new_compiler = std::make_unique<Compiler>(
-      new_fnptr, std::move(compiler),
-      (is_class_method ? FunctionType::CLASSMETHOD : FunctionType::FUNCTION));
+  FunctionType fn_type = is_class_method
+                             ? (fn_name == "init" ? FunctionType::CLASSINIT
+                                                  : FunctionType::CLASSMETHOD)
+                             : FunctionType::FUNCTION;
+  auto new_compiler =
+      std::make_unique<Compiler>(new_fnptr, std::move(compiler), fn_type);
   compiler = std::move(new_compiler);
   compiler->begin_scope();
   // Parse parameters (if there are any).
@@ -257,9 +261,10 @@ ObjFunction* Parser::finalise_function() {
               << "\n";
     return nullptr;
   } else {
-    // If there's no explicit return statement, we tack one on the end that
-    // returns nil.
-    emit_constant(std::monostate());
+    // Just tack a return at the end of the function body. If there was already
+    // one, it won't matter that we do this: it'll be unreachable. But if we
+    // fall off the end of a function body, this will catch us.
+    emit_auto_return_value();
     emit(lox::OpCode::RETURN);
     // Get the function object from the current compiler, and pop it off the
     // compiler stack.
@@ -413,7 +418,7 @@ void Parser::class_declaration() {
   emit(lox::OpCode::CLASS);
   emit(name_constant_index);
 
-  compiler->push_current_class();
+  push_current_class();
 
   // This call will emit code to read from the top of the stack and create
   // either a local or global variable with the class.
@@ -437,7 +442,7 @@ void Parser::class_declaration() {
   // any more.
   emit(lox::OpCode::POP);
 
-  compiler->pop_current_class();
+  pop_current_class();
 }
 
 void Parser::method() {
@@ -596,15 +601,34 @@ void Parser::statement() {
   }
 }
 
+// This function is hit whenever we have a plain `return;` (inside
+// `return_statement`) or if we fall off the end of a function without an
+// explicit return (inside `finalise_function`). It pushes the implicit return
+// value onto the stack (but does not emit the RETURN instruction).
+void Parser::emit_auto_return_value() {
+  if (compiler->get_function_type() == FunctionType::CLASSINIT) {
+    // return the instance, which is happily always at local slot 0.
+    emit(lox::OpCode::GET_LOCAL);
+    emit(static_cast<uint8_t>(0));
+  } else {
+    // No return value; return nil
+    emit_constant(std::monostate());
+  }
+}
+
 void Parser::return_statement() {
   if (compiler->get_function_type() == FunctionType::TOPLEVEL) {
     error("cannot return from top-level code", previous.line);
   }
   // Return value
   if (consume_if(TokenType::SEMICOLON)) {
-    // No return value; return nil
-    emit_constant(std::monostate());
+    // No explicit return value.
+    emit_auto_return_value();
   } else {
+    // Explicit return value.
+    if (compiler->get_function_type() == FunctionType::CLASSINIT) {
+      error("cannot return a value from an initializer", previous.line);
+    }
     expression();
     consume_or_error(TokenType::SEMICOLON, "expected ';' after return value");
   }
@@ -726,13 +750,17 @@ void Parser::number(bool _) {
 }
 
 void Parser::this_(bool _) {
-  // 1. We can't assign to this, hence the false.
-  // 2. We just treat 'this' like a local variable that happens to be at index
-  // zero.
-  if (!compiler->is_in_class()) {
+  // Note that we can't just check if the current function is CLASSMETHOD ||
+  // CLASSINIT, because we might be e.g. inside a nested function inside a
+  // class method. That's why we have a separate bit of info in the Parser that
+  // tracks whether we're currently allowed to use `this` or not.
+  if (!is_in_class()) {
     error("cannot use 'this' outside of a class", previous.line);
     return;
   }
+  // We can't assign to this, hence the false. But apart from that, we just
+  // treat 'this' like a local variable that happens to be at index 0 (which the
+  // constructor of Compiler does for us).
   variable(false);
 }
 

src/compiler.hpp

@@ -31,6 +31,7 @@ enum class FunctionType {
   TOPLEVEL,
   FUNCTION,
   CLASSMETHOD,
+  CLASSINIT,
 };
 
 struct Local {
@@ -45,8 +46,8 @@ class Compiler {
   Compiler(ObjFunction* fnptr, std::unique_ptr<Compiler> parent,
            FunctionType fn_type)
       : scope_depth(0), current_function(fnptr), fn_type(fn_type),
-        parent(std::move(parent)), current_class(nullptr) {
-    if (fn_type == FunctionType::CLASSMETHOD) {
+        parent(std::move(parent)) {
+    if (fn_type == FunctionType::CLASSMETHOD || fn_type == FunctionType::CLASSINIT) {
       // For a class method, we reserve slot 0 for the 'this' variable.
       declare_local("this");
     } else {
@@ -95,28 +96,12 @@ class Compiler {
     current_function->arity = new_arity;
   }
 
-  bool is_in_class() const { return current_class != nullptr; }
-  void push_current_class() {
-    auto new_current_class = std::make_unique<CurrentClass>();
-    new_current_class->enclosing = std::move(current_class);
-    current_class = std::move(new_current_class);
-  }
-  void pop_current_class() {
-    if (current_class == nullptr) {
-      throw std::runtime_error(
-          "Compiler::pop_current_class: no current class to pop");
-    } else {
-      current_class = std::move(current_class->enclosing);
-    }
-  }
-
 private:
   std::vector<Local> locals;
   size_t scope_depth;
   ObjFunction* current_function;
   FunctionType fn_type;
   std::unique_ptr<Compiler> parent;
-  std::unique_ptr<CurrentClass> current_class;
 };
 
 class Parser {
@@ -134,6 +119,22 @@ class Parser {
   GC& gc;
   std::unique_ptr<Compiler> compiler;
 
+  bool is_in_class() const { return current_class != nullptr; }
+  void push_current_class() {
+    auto new_current_class = std::make_unique<CurrentClass>();
+    new_current_class->enclosing = std::move(current_class);
+    current_class = std::move(new_current_class);
+  }
+  void pop_current_class() {
+    if (current_class == nullptr) {
+      throw std::runtime_error(
+          "Parser::pop_current_class: no current class to pop");
+    } else {
+      current_class = std::move(current_class->enclosing);
+    }
+  }
+  std::unique_ptr<CurrentClass> current_class;
+
   // Interact with scanner
   void advance();
   bool consume_or_error(scanner::TokenType type,
@@ -180,6 +181,7 @@ class Parser {
   void define_variable(std::string_view var_name);
   void define_global_variable(std::string_view name);
   void named_variable(std::string_view lexeme, bool can_assign);
+  void emit_auto_return_value();
 
   // Interact with chunk
   void emit(uint8_t byte) { compiler->emit(byte, previous.line); }

src/gc.hpp

@@ -35,7 +35,7 @@ class GC {
     obj->size = obj_size;
 #ifdef LOX_GC_DEBUG
     std::cerr << "        allocated object " << obj->to_repr() << " of size "
-              << obj_size;
+              << obj_size << "\n";
 #endif
     return obj;
   }

src/vm.cpp

@@ -71,6 +71,7 @@ VM::VM(std::unique_ptr<scanner::Scanner> scanner, GC gc)
   stack.reserve(MAX_STACK_SIZE);
   auto top_level_fn = _gc.alloc<ObjFunction>("#toplevel#", size_t(0));
   parser = std::make_unique<Parser>(std::move(scanner), top_level_fn, _gc);
+  initString = _gc.get_string_ptr("init");
 
   // Aggressive GC: run it every time we allocate
   _gc.set_alloc_callback([this]() { this->maybe_gc(); });
@@ -586,13 +587,36 @@ InterpretResult VM::run() {
         }
         case ObjType::CLASS: {
           auto classptr = static_cast<ObjClass*>(objptr);
-          // TODO: we are ignoring arguments right now.
-          if (nargs > 0) {
-            throw std::runtime_error(
-                "class constructors don't take arguments (yet)");
-          }
           ObjInstance* inst = _gc.alloc<ObjInstance>(classptr);
-          stack_replace_top(inst);
+          // Replace the class pointer on the stack with the instance pointer
+          // (so that it doesn't get GC'd). Recall that at this point the class
+          // will have been pushed to the stack, followed by any function
+          // arguments.
+          stack[stack.size() - 1 - nargs] = inst;
+          // Check for an initialiser.
+          auto init_method_itr = classptr->methods.find(initString);
+          if (init_method_itr != classptr->methods.end()) {
+            ObjClosure* init_closure = init_method_itr->second;
+            call(init_closure, nargs);
+            // Once we've finished calling the initialiser, the return value of
+            // init() will be on top of the stack. However, we don't want that.
+            // We want to replace it with the instance that we just created.
+            // Unfortunately, we can't handle that at the VM level -- when we
+            // do call(init_closure, nargs) we modify the call frame to
+            // enter the initialiser, but we don't have any hook into when
+            // its execution finishes. So we have to handle it in the compiler:
+            // when it is compiling an initialiser, any return statements will
+            // have to compile to something that returns the instance.
+          } else {
+            // If there's no init method, we are done in terms of the stack --
+            // we just need to check that the user didn't try to pass any
+            // arguments to the nonexistent initialiser.
+            if (nargs > 0) {
+              throw std::runtime_error(
+                  "class has no initialiser but was initialised with " +
+                  std::to_string(nargs) + " arguments");
+            }
+          }
           break;
         }
         case ObjType::BOUND_METHOD: {
@@ -651,8 +675,9 @@ InterpretResult VM::run() {
 
 void VM::maybe_gc() {
   if (_gc.should_gc()) {
-
     // Mark roots as grey.
+    _gc.mark_as_grey(initString);
+
     for (const auto& v : stack) {
       _gc.mark_as_grey(v);
     }

src/vm.hpp

@@ -69,6 +69,9 @@ class VM {
   // sorted upvalues that haven't been closed yet. They sort in decreasing order
   // of the stack slot that they point to
   std::vector<ObjUpvalue*> open_upvalues;
+  // Interned string for "init", which we use when we need to look up the
+  // initialiser method of a class.
+  ObjString* initString;
 
   CallFrame& current_frame() { return call_frames.back(); }
   Chunk& get_chunk() { return current_frame().closure->function->chunk; }

test.lox

@@ -1 +1,8 @@
-print this;
+class Num {
+  init(x) {
+    this.x = x;
+  }
+}
+
+var n = Num(3);
+print n.x;