From f1875736e37776c59ce8e41d4862692ce08f9df6 Mon Sep 17 00:00:00 2001
From: Mattia Giambirtone <nocturn9x@nocturn9x.space>
Date: Wed, 7 Dec 2022 09:15:29 +0100
Subject: [PATCH] Added switch statements

---
 src/backend/vm.nim             |  2 ++
 src/frontend/compiler.nim      | 22 ++++++++++++++++++++
 src/frontend/meta/ast.nim      | 17 +++++++++++++++-
 src/frontend/meta/bytecode.nim |  4 +++-
 src/frontend/meta/token.nim    |  2 +-
 src/frontend/parser.nim        | 25 +++++++++++++++++++++++
 src/util/symbols.nim           |  5 +++--
 tests/switch.pn                | 37 ++++++++++++++++++++++++++++++++++
 8 files changed, 109 insertions(+), 5 deletions(-)
 create mode 100644 tests/switch.pn

diff --git a/src/backend/vm.nim b/src/backend/vm.nim
index ef9ad76..47fb969 100644
--- a/src/backend/vm.nim
+++ b/src/backend/vm.nim
@@ -851,6 +851,8 @@ proc dispatch*(self: var PeonVM) =
                 # Pops N elements off the call stack
                 for _ in 0..<int(self.readShort()):
                     discard self.popc()
+            of DupTop:
+                self.push(self.peek())
             # Jump opcodes
             of Jump:
                 # Absolute jump
diff --git a/src/frontend/compiler.nim b/src/frontend/compiler.nim
index 1290c69..09a583e 100644
--- a/src/frontend/compiler.nim
+++ b/src/frontend/compiler.nim
@@ -2629,6 +2629,26 @@ proc namedBlock(self: Compiler, node: NamedBlockStmt) =
     self.endScope()
 
 
+proc switchStmt(self: Compiler, node: SwitchStmt) =
+    ## Compiles switch statements
+    self.expression(node.switch)
+    var ifJump: int = -1
+    var thenJumps: seq[int] = @[]
+    for branch in node.branches:
+        self.emitByte(DupTop, branch.body.token.line)
+        self.expression(branch.cond)
+        self.emitByte(Equal, branch.body.token.line)
+        ifJump = self.emitJump(JumpIfFalsePop, branch.body.token.line)
+        self.blockStmt(branch.body)
+        thenJumps.add(self.emitJump(JumpForwards, branch.body.token.line))
+        self.patchJump(ifJump)
+    if not node.default.isNil():
+        self.blockStmt(node.default)
+    for jump in thenJumps:
+        self.patchJump(jump)
+    self.emitByte(OpCode.Pop, node.token.line)
+
+
 proc statement(self: Compiler, node: Statement) =
     ## Compiles all statements
     case node.kind:
@@ -2644,6 +2664,8 @@ proc statement(self: Compiler, node: Statement) =
                 self.printRepl(kind, expression)
             else:
                 self.emitByte(Pop, node.token.line)
+        of NodeKind.switchStmt:
+            self.switchStmt(SwitchStmt(node))
         of NodeKind.namedBlockStmt:
             self.namedBlocks.add(NamedBlock(start: self.chunk.code.len(),
                                             depth: self.depth,
diff --git a/src/frontend/meta/ast.nim b/src/frontend/meta/ast.nim
index 9f4057e..ff63006 100644
--- a/src/frontend/meta/ast.nim
+++ b/src/frontend/meta/ast.nim
@@ -77,7 +77,8 @@ type
         identExpr, # Identifier
         pragmaExpr,
         refExpr,
-        ptrExpr
+        ptrExpr,
+        switchStmt
 
     # Here I would've rather used object variants, and in fact that's what was in
     # place before, but not being able to re-declare a field of the same type in
@@ -291,6 +292,11 @@ type
     
     Ptr* = ref object of Expression
         value*: Expression
+    
+    SwitchStmt* = ref object of Statement
+        switch*: Expression
+        branches*: seq[tuple[cond: Expression, body: BlockStmt]]
+        default*: BlockStmt
 
 
 proc isConst*(self: ASTNode): bool =
@@ -337,6 +343,15 @@ proc newPtrExpr*(expression: Expression, token: Token): Ptr =
     result.token = token
 
 
+proc newSwitchStmt*(switch: Expression, branches: seq[tuple[cond: Expression, body: BlockStmt]], default: BlockStmt, token: Token): SwitchStmt =
+    new(result)
+    result.kind = switchStmt
+    result.switch = switch
+    result.branches = branches
+    result.token = token
+    result.default = default
+
+
 proc newIntExpr*(literal: Token): IntExpr =
     result = IntExpr(kind: intExpr)
     result.literal = literal
diff --git a/src/frontend/meta/bytecode.nim b/src/frontend/meta/bytecode.nim
index ad5f156..f3811ac 100644
--- a/src/frontend/meta/bytecode.nim
+++ b/src/frontend/meta/bytecode.nim
@@ -186,7 +186,8 @@ type
         PopC,        # Pop off the call stack onto the operand stack
         PushC,       # Pop off the operand stack onto the call stack
         SysClock64,  # Pushes the output of a monotonic clock on the stack
-        LoadTOS      # Pushes the top of the call stack onto the operand stack
+        LoadTOS,     # Pushes the top of the call stack onto the operand stack
+        DupTop       # Duplicates the top of the operand stack onto the operand stack
 
 
 # We group instructions by their operation/operand types for easier handling when debugging
@@ -265,6 +266,7 @@ const simpleInstructions* = {Return, LoadNil,
                              Float32LessThan,
                              Float32GreaterOrEqual,
                              Float32LessOrEqual,
+                             DupTop
                             }
 
 # Constant instructions are instructions that operate on the bytecode constant table
diff --git a/src/frontend/meta/token.nim b/src/frontend/meta/token.nim
index 7ead699..46da886 100644
--- a/src/frontend/meta/token.nim
+++ b/src/frontend/meta/token.nim
@@ -38,7 +38,7 @@ type
     Yield, Defer, Try, Except,
     Finally, Type, Operator, Case,
     Enum, From, Ptr, Ref, Object,
-    Export, Block, Template
+    Export, Block, Template, Switch,
 
     # Literal types
     Integer, Float, String, Identifier,
diff --git a/src/frontend/parser.nim b/src/frontend/parser.nim
index 75d4031..20783c1 100644
--- a/src/frontend/parser.nim
+++ b/src/frontend/parser.nim
@@ -1203,12 +1203,37 @@ proc expressionStatement(self: Parser): Statement =
     result.file = self.file
 
 
+proc switchStmt(self: Parser): Statement =
+    ## Parses switch statements
+    let tok = self.peek(-1)
+    let switch = self.expression()
+    self.expect(TokenType.LeftBrace, "expecting '{' after switch condition")
+    var branches: seq[tuple[cond: Expression, body: BlockStmt]] = @[]
+    var match: Expression
+    var body: BlockStmt
+    var default: BlockStmt
+    while not self.check([TokenType.RightBrace, TokenType.Else]) and not self.done():
+        self.expect(TokenType.Case, "expecting at least one 'case' label in switch statement")
+        match = self.expression()
+        self.expect(TokenType.LeftBrace, "expecting '{' after expression match in switch statement")
+        body = BlockStmt(self.blockStmt())
+        branches.add((cond: match, body: body))
+    if self.match(Else):
+        self.expect(TokenType.LeftBrace, "expecting '{' after else clause in switch statement")
+        default = BlockStmt(self.blockStmt())
+    self.expect(TokenType.RightBrace, "missing closing '}' in switch statement")
+    result = newSwitchStmt(switch, branches, default, tok)
+
+
 proc statement(self: Parser): Statement =
     ## Parses statements
     case self.peek().kind:
         of If:
             discard self.step()
             result = self.ifStmt()
+        of Switch:
+            discard self.step()
+            result = self.switchStmt()
         of Assert:
             discard self.step()
             result = self.assertStmt()
diff --git a/src/util/symbols.nim b/src/util/symbols.nim
index 8fcef1b..a8ec1f4 100644
--- a/src/util/symbols.nim
+++ b/src/util/symbols.nim
@@ -46,6 +46,9 @@ proc fillSymbolTable*(tokenizer: Lexer) =
     tokenizer.symbols.addKeyword("return", TokenType.Return)
     tokenizer.symbols.addKeyword("object", Object)
     tokenizer.symbols.addKeyword("export", Export)
+    tokenizer.symbols.addKeyword("block", TokenType.Block)
+    tokenizer.symbols.addKeyword("template", TokenType.Template)
+    tokenizer.symbols.addKeyword("switch", TokenType.Switch)
     # These are more like expressions with a reserved
     # name that produce a value of a builtin type,
     # but we don't need to care about that until
@@ -55,8 +58,6 @@ proc fillSymbolTable*(tokenizer: Lexer) =
     tokenizer.symbols.addKeyword("false", False)
     tokenizer.symbols.addKeyword("ref", TokenType.Ref)
     tokenizer.symbols.addKeyword("ptr", TokenType.Ptr)
-    tokenizer.symbols.addKeyword("block", TokenType.Block)
-    tokenizer.symbols.addKeyword("template", TokenType.Template)
     for sym in [">", "<", "=", "~", "/", "+", "-", "_", "*", "?", "@", ":", "==", "!=",
                 ">=", "<=", "+=", "-=", "/=", "*=", "**=", "!", "%", "&", "|", "^",
                 ">>", "<<"]:
diff --git a/tests/switch.pn b/tests/switch.pn
new file mode 100644
index 0000000..16985cb
--- /dev/null
+++ b/tests/switch.pn
@@ -0,0 +1,37 @@
+import std;
+
+
+switch 2 + 2 {
+    case 4 {
+        print(true);
+    }
+    case 3 {
+        print(false);
+    }
+    else {
+        print(false);
+    }
+}
+
+
+switch 2 + 1 {
+    case 4 {
+        print(false);
+    }
+    case 3 {
+        print(true);
+    }
+}
+
+
+switch 2 + 3 {
+    case 4 {
+        print(false);
+    }
+    case 3 {
+        print(false);
+    }
+    else {
+        print(true);
+    }
+}