From 95d83984870c078c014dd6bd6ad7809ed5872635 Mon Sep 17 00:00:00 2001
From: R-unic <roleyspersonal@gmail.com>
Date: Thu, 5 Oct 2023 16:48:56 -0400
Subject: [PATCH] feat: implement a maximum recursion depth

---
 src/code-analysis/parser/ast/statements/if.ts |  2 +-
 .../parser/ast/statements/while.ts            |  2 +-
 src/runtime/interpreter.ts                    | 37 ++++++++++++++++---
 3 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/src/code-analysis/parser/ast/statements/if.ts b/src/code-analysis/parser/ast/statements/if.ts
index 41557132..0fb91573 100644
--- a/src/code-analysis/parser/ast/statements/if.ts
+++ b/src/code-analysis/parser/ast/statements/if.ts
@@ -3,7 +3,7 @@ import AST from "..";
 
 export class IfStatement extends AST.Statement {
   public constructor(
-    public readonly token: Token,
+    public readonly token: Token<undefined>,
     public readonly condition: AST.Expression,
     public readonly body: AST.Statement,
     public readonly elseBranch?: AST.Statement
diff --git a/src/code-analysis/parser/ast/statements/while.ts b/src/code-analysis/parser/ast/statements/while.ts
index 566cfa1d..a5dbe85f 100644
--- a/src/code-analysis/parser/ast/statements/while.ts
+++ b/src/code-analysis/parser/ast/statements/while.ts
@@ -3,7 +3,7 @@ import AST from "..";
 
 export class WhileStatement extends AST.Statement {
   public constructor(
-    public readonly token: Token,
+    public readonly token: Token<undefined>,
     public readonly condition: AST.Expression,
     public readonly body: AST.Statement,
     public readonly elseBranch?: AST.Statement
diff --git a/src/runtime/interpreter.ts b/src/runtime/interpreter.ts
index f295e111..9fb6bf38 100644
--- a/src/runtime/interpreter.ts
+++ b/src/runtime/interpreter.ts
@@ -6,6 +6,8 @@ import type { Callable } from "./types/callable";
 import type Binder from "../code-analysis/type-checker/binder";
 import type Resolver from "../code-analysis/resolver";
 import Scope from "./scope";
+import PFunction from "./types/function";
+import HookedException from "./hooked-exceptions";
 import Intrinsics from "./intrinsics";
 import Syntax from "../code-analysis/syntax/syntax-type";
 import AST from "../code-analysis/parser/ast";
@@ -28,17 +30,20 @@ import type { VariableDeclarationStatement } from "../code-analysis/parser/ast/s
 import type { BlockStatement } from "../code-analysis/parser/ast/statements/block";
 import type { IfStatement } from "../code-analysis/parser/ast/statements/if";
 import type { WhileStatement } from "../code-analysis/parser/ast/statements/while";
-import { IndexExpression } from "../code-analysis/parser/ast/expressions";
+import type { IndexExpression } from "../code-analysis/parser/ast/expressions";
 import { PropertyAssignmentExpression } from "../code-analysis/parser/ast/expressions/property-assignment";
-import { FunctionDeclarationStatement } from "../code-analysis/parser/ast/statements/function-declaration";
-import PFunction from "./types/function";
-import { ReturnStatement } from "../code-analysis/parser/ast/statements/return";
-import HookedException from "./hooked-exceptions";
+import type { FunctionDeclarationStatement } from "../code-analysis/parser/ast/statements/function-declaration";
+import type { ReturnStatement } from "../code-analysis/parser/ast/statements/return";
+
+const MAX_RECURSION_DEPTH = 1200;
 
 export default class Interpreter implements AST.Visitor.Expression<ValueType>, AST.Visitor.Statement<void> {
   public readonly globals = new Scope;
   public scope = this.globals;
 
+  private loopLevel = 0;
+  private recursionDepth = 1;
+
   public constructor(
     public readonly runner: P,
     public readonly resolver: Resolver,
@@ -62,14 +67,24 @@ export default class Interpreter implements AST.Visitor.Expression<ValueType>, A
   }
 
   public visitWhileStatement(stmt: WhileStatement): void {
+    this.loopLevel++;
     const inverted = stmt.token.syntax === Syntax.Until;
-    while (inverted ? !this.evaluate(stmt.condition) : this.evaluate(stmt.condition))
+    let depth = 0;
+
+    while (inverted ? !this.evaluate(stmt.condition) : this.evaluate(stmt.condition)) {
+      this.startRecursion(stmt.token);
       this.execute(stmt.body);
+      depth++;
+    }
+
+    this.endRecursion(depth);
+    this.loopLevel--;
   }
 
   public visitIfStatement(stmt: IfStatement): void {
     const condition = this.evaluate(stmt.condition);
     const inverted = stmt.token.syntax === Syntax.Unless;
+
     if (inverted ? !condition : condition)
       this.execute(stmt.body)
     else if (stmt.elseBranch)
@@ -275,4 +290,14 @@ export default class Interpreter implements AST.Visitor.Expression<ValueType>, A
     } else
       return statements.accept(this);
   }
+
+  private startRecursion(token: Token<undefined>): void {
+    this.recursionDepth++;
+    if (this.recursionDepth < MAX_RECURSION_DEPTH) return;
+    throw new RuntimeError(`Stack overflow: Recursion depth of ${MAX_RECURSION_DEPTH} exceeded`, token);
+  }
+
+  private endRecursion(level = 1): void {
+    this.recursionDepth -= level;
+  }
 }
\ No newline at end of file