From 65f61fcd900998b7bfd312da8e2705728b985b04 Mon Sep 17 00:00:00 2001
From: Alex Zenla <alex@fork.run>
Date: Tue, 5 Sep 2023 01:35:27 -0700
Subject: [PATCH] evaluator: implement an elaborate system that should mean
 that recursive functions don't require reaching far up the stack

---
 .../pizza/pork/evaluator/EvaluationVisitor.kt  |  6 ++++--
 .../pizza/pork/evaluator/FastVariableCache.kt  | 18 ++++++++++++++++++
 .../kotlin/gay/pizza/pork/evaluator/Scope.kt   | 17 ++++++++++++++---
 3 files changed, 36 insertions(+), 5 deletions(-)
 create mode 100644 evaluator/src/main/kotlin/gay/pizza/pork/evaluator/FastVariableCache.kt

diff --git a/evaluator/src/main/kotlin/gay/pizza/pork/evaluator/EvaluationVisitor.kt b/evaluator/src/main/kotlin/gay/pizza/pork/evaluator/EvaluationVisitor.kt
index 5e05b07..1d5d096 100644
--- a/evaluator/src/main/kotlin/gay/pizza/pork/evaluator/EvaluationVisitor.kt
+++ b/evaluator/src/main/kotlin/gay/pizza/pork/evaluator/EvaluationVisitor.kt
@@ -4,6 +4,7 @@ import gay.pizza.pork.ast.*
 
 class EvaluationVisitor(root: Scope) : NodeVisitor<Any> {
   private var currentScope: Scope = root
+  private var insideFastCachePreservation = false
 
   override fun visitIntLiteral(node: IntLiteral): Any = node.value
   override fun visitStringLiteral(node: StringLiteral): Any = node.text
@@ -28,7 +29,7 @@ class EvaluationVisitor(root: Scope) : NodeVisitor<Any> {
 
   override fun visitLambda(node: Lambda): CallableFunction {
     return CallableFunction { arguments ->
-      currentScope = currentScope.fork()
+      currentScope = currentScope.fork(inheritFastCache = insideFastCachePreservation)
       for ((index, argumentSymbol) in node.arguments.withIndex()) {
         currentScope.define(argumentSymbol.id, arguments.values[index])
       }
@@ -104,7 +105,8 @@ class EvaluationVisitor(root: Scope) : NodeVisitor<Any> {
   override fun visitFunctionDefinition(node: FunctionDefinition): Any {
     val blockFunction = visitBlock(node.block) as BlockFunction
     val function = CallableFunction { arguments ->
-      currentScope = currentScope.fork()
+      currentScope = currentScope.fork(inheritFastCache = insideFastCachePreservation)
+      currentScope.fastVariableCache.put(node.symbol.id, currentScope.value(node.symbol.id))
       for ((index, argumentSymbol) in node.arguments.withIndex()) {
         currentScope.define(argumentSymbol.id, arguments.values[index])
       }
diff --git a/evaluator/src/main/kotlin/gay/pizza/pork/evaluator/FastVariableCache.kt b/evaluator/src/main/kotlin/gay/pizza/pork/evaluator/FastVariableCache.kt
new file mode 100644
index 0000000..96c7bbf
--- /dev/null
+++ b/evaluator/src/main/kotlin/gay/pizza/pork/evaluator/FastVariableCache.kt
@@ -0,0 +1,18 @@
+package gay.pizza.pork.evaluator
+
+class FastVariableCache {
+  private val cache = mutableMapOf<String, Any>()
+
+  fun put(key: String, value: Any) {
+    cache[key] = value
+  }
+
+  fun lookup(key: String): Any = cache[key] ?: NotFound
+  fun has(key: String): Boolean = cache.containsKey(key)
+
+  fun invalidate(key: String) {
+    cache.remove(key)
+  }
+
+  object NotFound
+}
diff --git a/evaluator/src/main/kotlin/gay/pizza/pork/evaluator/Scope.kt b/evaluator/src/main/kotlin/gay/pizza/pork/evaluator/Scope.kt
index b18e104..e736522 100644
--- a/evaluator/src/main/kotlin/gay/pizza/pork/evaluator/Scope.kt
+++ b/evaluator/src/main/kotlin/gay/pizza/pork/evaluator/Scope.kt
@@ -1,10 +1,15 @@
 package gay.pizza.pork.evaluator
 
-class Scope(val parent: Scope? = null, inherits: List<Scope> = emptyList()) {
+class Scope(
+  val parent: Scope? = null,
+  inherits: List<Scope> = emptyList(),
+  val fastVariableCache: FastVariableCache = FastVariableCache()
+) {
   private val inherited = inherits.toMutableList()
   private val variables = mutableMapOf<String, Any>()
 
   fun define(name: String, value: Any) {
+    fastVariableCache.invalidate(name)
     if (variables.containsKey(name)) {
       throw RuntimeException("Variable '${name}' is already defined")
     }
@@ -12,6 +17,11 @@ class Scope(val parent: Scope? = null, inherits: List<Scope> = emptyList()) {
   }
 
   fun value(name: String): Any {
+    val fastCached = fastVariableCache.lookup(name)
+    if (fastCached != FastVariableCache.NotFound) {
+      return fastCached
+    }
+
     val value = valueOrNotFound(name)
     if (value == NotFound) {
       throw RuntimeException("Variable '${name}' not defined.")
@@ -48,8 +58,9 @@ class Scope(val parent: Scope? = null, inherits: List<Scope> = emptyList()) {
     return value.call(arguments)
   }
 
-  fun fork(): Scope {
-    return Scope(this)
+  fun fork(inheritFastCache: Boolean = false): Scope {
+    val fastCache = if (inheritFastCache) fastVariableCache else FastVariableCache()
+    return Scope(this, fastVariableCache = fastCache)
   }
 
   fun leave(): Scope {