Compile ForCond statements

2 files changed, 77 insertions, 2 deletions

diff --git a/pkg/lang/compiler/compiler.go b/pkg/lang/compiler/compiler.go
index cad4328..0356411 100644
--- a/pkg/lang/compiler/compiler.go
+++ b/pkg/lang/compiler/compiler.go
@@ -43,8 +43,11 @@ func (comp *Compiler) compileStmt(t *code.Builder, stmt ast.Stmt) error {
 		decl := stmt.Value.(ast.StmtVarDecl)
 		err = comp.compileVarDeclStmt(t, decl)
 	case ast.StmtKindIf:
-		ifstmt := stmt.Value.(ast.StmtIf)
-		err = comp.compileIfStmt(t, ifstmt)
+		ifStmt := stmt.Value.(ast.StmtIf)
+		err = comp.compileIfStmt(t, ifStmt)
+	case ast.StmtKindForCond:
+		forCondStmt := stmt.Value.(ast.StmtForCond)
+		err = comp.compileForCondStmt(t, forCondStmt)
 	case ast.StmtKindExpr:
 		expr := stmt.Value.(ast.StmtExpr).Value
 		err = comp.compileExpr(t, expr)
@@ -137,6 +140,47 @@ func (comp *Compiler) compileIfStmt(t *code.Builder, ifStmt ast.StmtIf) error {
 	return nil
 }
 
+func (comp *Compiler) compileForCondStmt(t *code.Builder, forCondStmt ast.StmtForCond) error {
+	// Parts:
+	// 1. Condition check: Decides whether the loop should run
+	// 2. Condition jump: Jumps to the end of the for if condition was false
+	// 3. Do block: Does something
+	// 4. Repeat jump: Jumps back to start
+
+	// Do block
+	doTarget := code.NewBuilder()
+	if err := comp.compileBlockNode(&doTarget, forCondStmt.Do); err != nil {
+		return err
+	}
+
+	jmpLength := 9
+
+	conditionTarget := code.NewBuilder()
+	if !forCondStmt.Cond.IsEmpty() {
+		// Condition check
+		if err := comp.compileExpr(&conditionTarget, forCondStmt.Cond); err != nil {
+			return err
+		}
+
+		endOfFor := conditionTarget.Len() + doTarget.Len() + jmpLength*2
+
+		// Condition jump
+		conditionTarget.AppendOp(code.OpJf)
+		conditionTarget.AppendReferenceToPc(int64(endOfFor))
+	}
+
+	subUnit := conditionTarget
+	subUnit.AppendBuilder(doTarget)
+
+	// Repeat jump
+	subUnit.AppendOp(code.OpJmp)
+	subUnit.AppendReferenceToPc(int64(0)) // Start of the for
+
+	t.AppendBuilder(subUnit)
+
+	return nil
+}
+
 func (comp *Compiler) compileExpr(t *code.Builder, expr ast.Expr) error {
 	switch expr.Kind {
 	case ast.ExprKindBinary:
diff --git a/pkg/lang/compiler/compiler_test.go b/pkg/lang/compiler/compiler_test.go
index 88413c7..1a759f8 100644
--- a/pkg/lang/compiler/compiler_test.go
+++ b/pkg/lang/compiler/compiler_test.go
@@ -234,6 +234,37 @@ func TestNestedIfs(t *testing.T) {
 	mustCompileTo(t, src, expected)
 }
 
+func TestForCond(t *testing.T) {
+	src := `
+	for true {
+		for false {
+			1
+		}
+	}
+	`
+
+	expected := `
+	@start:
+	push_true
+	jf @end
+
+	@inner_start:
+	push_false
+	jf @inner_end
+
+	push_int 1
+	jmp @inner_start
+
+	@inner_end:
+	jmp @start
+
+	@end
+	halt
+	`
+
+	mustCompileTo(t, src, expected)
+}
+
 func mustCompileTo(t *testing.T, src, expected string) {
 	scanner := scanner.New(strings.NewReader(src))
 	tokens, err := scanner.Scan()