Implemented switch statements

examples/fizzbuzz/fizzbuzz.q

@@ -9,18 +9,18 @@ fizzbuzz(n) {
 	x := 1
 
 	loop {
-		// TODO: implement switch statement
-		if x % 15 == 0 {
+		switch {
+			x % 15 == 0 {
 				print("FizzBuzz", 8)
-		} else {
-			if x % 5 == 0 {
-				print("Buzz", 4)
-			} else {
-				if x % 3 == 0 {
-					print("Fizz", 4)
-				} else {
-					log.number(x)
 			}
+			x % 5 == 0 {
+				print("Buzz", 4)
+			}
+			x % 3 == 0 {
+				print("Fizz", 4)
+			}
+			_ {
+				log.number(x)
 			}
 		}
 
@@ -28,9 +28,9 @@ fizzbuzz(n) {
 
 		if x > n {
 			return
-		} else {
-			print(" ", 1)
 		}
+
+		print(" ", 1)
 	}
 }
 

src/build/ast/Count.go

@@ -20,11 +20,6 @@ func Count(body AST, buffer []byte, kind token.Kind, name string) uint8 {
 		case *Define:
 			count += node.Expression.Count(buffer, kind, name)
 
-		case *Return:
-			if node.Value != nil {
-				count += node.Value.Count(buffer, kind, name)
-			}
-
 		case *If:
 			count += node.Condition.Count(buffer, kind, name)
 			count += Count(node.Body, buffer, kind, name)
@@ -33,6 +28,20 @@ func Count(body AST, buffer []byte, kind token.Kind, name string) uint8 {
 		case *Loop:
 			count += Count(node.Body, buffer, kind, name)
 
+		case *Return:
+			if node.Value != nil {
+				count += node.Value.Count(buffer, kind, name)
+			}
+
+		case *Switch:
+			for _, c := range node.Cases {
+				if c.Condition != nil {
+					count += c.Condition.Count(buffer, kind, name)
+				}
+
+				count += Count(c.Body, buffer, kind, name)
+			}
+
 		default:
 			panic("unknown AST type")
 		}

src/build/ast/Switch.go

@@ -0,0 +1,16 @@
+package ast
+
+import (
+	"git.akyoto.dev/cli/q/src/build/expression"
+)
+
+// Switch represents a switch statement.
+type Switch struct {
+	Cases []Case
+}
+
+// Case represents a case inside a switch.
+type Case struct {
+	Condition *expression.Expression
+	Body      AST
+}

src/build/ast/parseKeyword.go

@@ -51,6 +51,21 @@ func parseKeyword(tokens token.List, source []byte, nodes AST) (Node, error) {
 		value := expression.Parse(tokens[1:])
 		return &Return{Value: value}, nil
 
+	case token.Switch:
+		blockStart := tokens.IndexKind(token.BlockStart)
+		blockEnd := tokens.LastIndexKind(token.BlockEnd)
+
+		if blockStart == -1 {
+			return nil, errors.New(errors.MissingBlockStart, nil, tokens[0].End())
+		}
+
+		if blockEnd == -1 {
+			return nil, errors.New(errors.MissingBlockEnd, nil, tokens[len(tokens)-1].End())
+		}
+
+		cases, err := parseSwitch(tokens[blockStart+1:blockEnd], source)
+		return &Switch{Cases: cases}, err
+
 	default:
 		return nil, errors.New(&errors.KeywordNotImplemented{Keyword: tokens[0].Text(source)}, nil, tokens[0].Position)
 	}

src/build/ast/parseSwitch.go

@@ -0,0 +1,37 @@
+package ast
+
+import (
+	"git.akyoto.dev/cli/q/src/build/expression"
+	"git.akyoto.dev/cli/q/src/build/token"
+)
+
+// parseSwitch generates the cases inside a switch statement.
+func parseSwitch(tokens token.List, source []byte) ([]Case, error) {
+	var cases []Case
+
+	err := EachInstruction(tokens, func(caseTokens token.List) error {
+		blockStart, _, body, err := block(caseTokens, source)
+
+		if err != nil {
+			return err
+		}
+
+		conditionTokens := caseTokens[:blockStart]
+		var condition *expression.Expression
+
+		if len(conditionTokens) == 1 && conditionTokens[0].Kind == token.Identifier && conditionTokens[0].Text(source) == "_" {
+			condition = nil
+		} else {
+			condition = expression.Parse(conditionTokens)
+		}
+
+		cases = append(cases, Case{
+			Condition: condition,
+			Body:      body,
+		})
+
+		return nil
+	})
+
+	return cases, err
+}

src/build/core/CompileASTNode.go

@@ -23,10 +23,6 @@ func (f *Function) CompileASTNode(node ast.Node) error {
 		f.Fold(node.Expression)
 		return f.CompileDefinition(node)
 
-	case *ast.Return:
-		f.Fold(node.Value)
-		return f.CompileReturn(node)
-
 	case *ast.If:
 		f.Fold(node.Condition)
 		return f.CompileIf(node)
@@ -34,6 +30,17 @@ func (f *Function) CompileASTNode(node ast.Node) error {
 	case *ast.Loop:
 		return f.CompileLoop(node)
 
+	case *ast.Return:
+		f.Fold(node.Value)
+		return f.CompileReturn(node)
+
+	case *ast.Switch:
+		for _, c := range node.Cases {
+			f.Fold(c.Condition)
+		}
+
+		return f.CompileSwitch(node)
+
 	default:
 		panic("unknown AST type")
 	}

src/build/core/CompileSwitch.go

@@ -0,0 +1,55 @@
+package core
+
+import (
+	"fmt"
+
+	"git.akyoto.dev/cli/q/src/build/asm"
+	"git.akyoto.dev/cli/q/src/build/ast"
+)
+
+// CompileSwitch compiles a multi-branch instruction.
+func (f *Function) CompileSwitch(s *ast.Switch) error {
+	f.count.multiBranch++
+	end := fmt.Sprintf("%s_switch_%d_end", f.UniqueName, f.count.multiBranch)
+
+	for _, branch := range s.Cases {
+		if branch.Condition == nil {
+			f.PushScope(branch.Body, f.File.Bytes)
+			err := f.CompileAST(branch.Body)
+
+			if err != nil {
+				return err
+			}
+
+			f.PopScope()
+			break
+		}
+
+		f.count.branch++
+
+		var (
+			success = fmt.Sprintf("%s_case_%d_true", f.UniqueName, f.count.branch)
+			fail    = fmt.Sprintf("%s_case_%d_false", f.UniqueName, f.count.branch)
+			err     = f.CompileCondition(branch.Condition, success, fail)
+		)
+
+		if err != nil {
+			return err
+		}
+
+		f.AddLabel(success)
+		f.PushScope(branch.Body, f.File.Bytes)
+		err = f.CompileAST(branch.Body)
+
+		if err != nil {
+			return err
+		}
+
+		f.Jump(asm.JUMP, end)
+		f.PopScope()
+		f.AddLabel(fail)
+	}
+
+	f.AddLabel(end)
+	return nil
+}

src/build/core/Function.go

@@ -24,6 +24,7 @@ type Function struct {
 type counter struct {
 	assert      int
 	branch      int
+	multiBranch int
 	data        int
 	loop        int
 	subBranch   int

src/build/token/Kind.go

@@ -69,5 +69,6 @@ const (
 	Import                      // import
 	Loop                        // loop
 	Return                      // return
+	Switch                      // switch
 	_keywordsEnd                // </keywords>
 )

src/build/token/Tokenize.go

@@ -151,6 +151,8 @@ func Tokenize(buffer []byte) List {
 					kind = Loop
 				case "return":
 					kind = Return
+				case "switch":
+					kind = Switch
 				}
 
 				tokens = append(tokens, Token{Kind: kind, Position: position, Length: Length(len(identifier))})