From fd71a46edeb061d725a63b2ecc568ae6f33cd0ce Mon Sep 17 00:00:00 2001
From: Olivier Abrivard <olivier@abrivard.fr>
Date: Mon, 30 Sep 2024 11:41:52 +0200
Subject: [PATCH] Add methods on classes

---
 src/main/fr/celticinfo/lox/Interpreter.kt     |  9 ++++-
 src/main/fr/celticinfo/lox/LoxClass.kt        |  6 ++-
 src/main/fr/celticinfo/lox/LoxInstance.kt     |  5 +++
 src/main/fr/celticinfo/lox/Resolver.kt        |  8 +++-
 src/test/fr/celticinfo/lox/InterpreterTest.kt | 38 ++++++++++++++++++-
 src/test/fr/celticinfo/lox/ParserTest.kt      | 20 ++++++++++
 6 files changed, 81 insertions(+), 5 deletions(-)

diff --git a/src/main/fr/celticinfo/lox/Interpreter.kt b/src/main/fr/celticinfo/lox/Interpreter.kt
index ed33de9..3930de3 100644
--- a/src/main/fr/celticinfo/lox/Interpreter.kt
+++ b/src/main/fr/celticinfo/lox/Interpreter.kt
@@ -48,7 +48,14 @@ class Interpreter: ExprVisitor<Any?>, StmtVisitor<Unit>{
 
     override fun visitClassStmt(stmt: ClassStmt) {
         environment.define(stmt.name.lexeme, null)
-        val klass = LoxClass(stmt.name.lexeme)
+
+        val methods = stmt.methods.associate { method ->
+            val function = LoxFunction(method, environment)
+            method.name.lexeme to function
+        }
+
+        val klass = LoxClass(stmt.name.lexeme, methods)
+
         environment.assign(stmt.name, klass)
     }
 
diff --git a/src/main/fr/celticinfo/lox/LoxClass.kt b/src/main/fr/celticinfo/lox/LoxClass.kt
index ea44ad2..0325805 100644
--- a/src/main/fr/celticinfo/lox/LoxClass.kt
+++ b/src/main/fr/celticinfo/lox/LoxClass.kt
@@ -1,11 +1,15 @@
 package fr.celticinfo.lox
 
-class LoxClass(private val name: String) : LoxCallable {
+class LoxClass(private val name: String, private val methods: Map<String, LoxFunction>) : LoxCallable {
     override fun call(interpreter: Interpreter, arguments: List<Any?>): Any? {
         val instance = LoxInstance(this)
         return instance
     }
 
+    fun findMethod(name: String): LoxFunction? {
+        return methods[name]
+    }
+
     override fun arity() = 0
 
     override fun toString() = name
diff --git a/src/main/fr/celticinfo/lox/LoxInstance.kt b/src/main/fr/celticinfo/lox/LoxInstance.kt
index f104c4a..120c548 100644
--- a/src/main/fr/celticinfo/lox/LoxInstance.kt
+++ b/src/main/fr/celticinfo/lox/LoxInstance.kt
@@ -8,6 +8,11 @@ class LoxInstance(private val klass: LoxClass) {
             return fields[name.lexeme]
         }
 
+        val method = klass.findMethod(name.lexeme)
+        if (method != null) {
+            return method
+        }
+
         throw RuntimeError(name, "Undefined property '${name.lexeme}'.")
     }
 
diff --git a/src/main/fr/celticinfo/lox/Resolver.kt b/src/main/fr/celticinfo/lox/Resolver.kt
index 9fb92ea..d2d69e7 100644
--- a/src/main/fr/celticinfo/lox/Resolver.kt
+++ b/src/main/fr/celticinfo/lox/Resolver.kt
@@ -20,6 +20,11 @@ class Resolver(private val interpreter: Interpreter) : ExprVisitor<Unit>, StmtVi
     override fun visitClassStmt(stmt: ClassStmt) {
         declare(stmt.name)
         define(stmt.name)
+
+        for (method in stmt.methods) {
+            val declaration = FunctionType.METHOD
+            resolveFunction(method, declaration)
+        }
     }
 
     override fun visitExpression(stmt: Expression) {
@@ -172,5 +177,6 @@ class Resolver(private val interpreter: Interpreter) : ExprVisitor<Unit>, StmtVi
 
 enum class FunctionType {
     NONE,
-    FUNCTION
+    FUNCTION,
+    METHOD
 }
diff --git a/src/test/fr/celticinfo/lox/InterpreterTest.kt b/src/test/fr/celticinfo/lox/InterpreterTest.kt
index b629cd8..eedd132 100644
--- a/src/test/fr/celticinfo/lox/InterpreterTest.kt
+++ b/src/test/fr/celticinfo/lox/InterpreterTest.kt
@@ -838,7 +838,6 @@ var a = "global";
         }
     }
 
-
     @Test
     fun `valid code with properties initialization and access`() {
         val standardOut = System.out
@@ -873,4 +872,39 @@ var a = "global";
             System.setOut(standardOut)
         }
     }
-}
+
+    @Test
+    fun `valid code with method`() {
+        val standardOut = System.out
+        val outputStreamCaptor = ByteArrayOutputStream()
+
+        System.setOut(PrintStream(outputStreamCaptor))
+
+        try {
+            val code = """
+                class Bacon {
+                  eat() {
+                    print "Crunch crunch crunch!";
+                  }
+                }
+                
+                Bacon().eat();
+            """.trimIndent()
+            val scanner = Scanner(code)
+            val tokens = scanner.scanTokens()
+            val parser = Parser(tokens)
+            val statements = parser.parse()
+            assertEquals(2, statements.size)
+
+            val interpreter = Interpreter()
+
+            val resolver = Resolver(interpreter)
+            resolver.resolve(statements)
+
+            interpreter.interpret(statements)
+            val output = outputStreamCaptor.toString().trim()
+            assertEquals("Crunch crunch crunch!", output)
+        } finally {
+            System.setOut(standardOut)
+        }
+    }}
diff --git a/src/test/fr/celticinfo/lox/ParserTest.kt b/src/test/fr/celticinfo/lox/ParserTest.kt
index fdc91ff..14bbe56 100644
--- a/src/test/fr/celticinfo/lox/ParserTest.kt
+++ b/src/test/fr/celticinfo/lox/ParserTest.kt
@@ -291,4 +291,24 @@ class ParserTest {
         assertTrue(statements[2] is Expression)
         assertTrue(statements[3] is Print)
     }
+
+    @Test
+    fun `valid code with methods`() {
+        val code = """
+            class Bacon {
+              eat() {
+                print "Crunch crunch crunch!";
+              }
+            }
+            
+            Bacon().eat();
+        """.trimIndent()
+        val scanner = Scanner(code)
+        val tokens = scanner.scanTokens()
+        val parser = Parser(tokens)
+        val statements = parser.parse()
+        assertEquals(2,statements.size)
+        assertTrue(statements[0] is ClassStmt)
+        assertTrue(statements[1] is Expression)
+    }
 }
\ No newline at end of file