For range

compiler/compiler/for.go

@@ -1,8 +1,20 @@
 package compiler
 
-import "github.com/zegl/tre/compiler/parser"
+import (
+	"fmt"
+	"github.com/zegl/tre/compiler/parser"
+)
 
 func (c *Compiler) compileForNode(v *parser.ForNode) {
+	if v.IsThreeTypeFor {
+		c.compileForThreeType(v)
+		return
+	}
+
+	c.compileForRange(v)
+}
+
+func (c *Compiler) compileForThreeType(v *parser.ForNode) {
 	// TODO: create a new context-block for code running inside the for loop
 	c.compile([]parser.Node{
 		v.BeforeLoop,
@@ -48,6 +60,97 @@ func (c *Compiler) compileForNode(v *parser.ForNode) {
 	c.contextLoopContinue = c.contextLoopContinue[0 : len(c.contextLoopContinue)-1]
 }
 
+func (c *Compiler) compileForRange(v *parser.ForNode) {
+	// A for range that iterates over a slice is just syntactic sugar
+	// for k, v := range a
+	// for k := 0; k < len(a); k++ { v := a[k] }
+	var modifiedBlock []parser.Node
+
+	// The node/value that we're iterating over
+	var rangeItem parser.Node
+
+	if forAlloc, ok := v.BeforeLoop.(*parser.AllocNode); ok {
+		forAllocRange := forAlloc.Val.(*parser.RangeNode)
+		rangeItem = forAllocRange.Item
+	} else if forRange, ok := v.BeforeLoop.(*parser.RangeNode); ok {
+		rangeItem = forRange.Item
+	} else {
+		panic("unexpected for/range beforeLoop type: " + fmt.Sprintf("%T %+v", v.BeforeLoop, v.BeforeLoop))
+	}
+
+	// Alloc the value of rangeItem and save it in a variable
+	rangeItemName := getVarName("range-item")
+	c.compileAllocNode(&parser.AllocNode{
+		Name: rangeItemName,
+		Val:  rangeItem,
+	})
+
+	// Call and alloc len() and save it in a variable
+	forItemLenName := getVarName("range-item-len")
+	c.compileAllocNode(&parser.AllocNode{
+		Name: forItemLenName,
+		Val: &parser.CallNode{
+			Function:  &parser.NameNode{Name: "len"},
+			Arguments: []parser.Node{&parser.NameNode{Name: rangeItemName}},
+		},
+	})
+
+	forKeyName := getVarName("for-key")
+
+	// Ranges that use the key and value
+	if forAlloc, ok := v.BeforeLoop.(*parser.AllocNode); ok {
+		var keyName string
+		if forAlloc.MultiNames == nil {
+			keyName = forAlloc.Name
+		} else {
+			keyName = forAlloc.MultiNames.Names[0].Name
+		}
+
+		// Assignment of key
+		modifiedBlock = append(modifiedBlock, &parser.AllocNode{
+			Name: keyName,
+			Val:  &parser.NameNode{Name: forKeyName},
+		})
+
+		// Assignment of value
+		if forAlloc.MultiNames != nil && len(forAlloc.MultiNames.Names) >= 2 {
+			modifiedBlock = append(modifiedBlock, &parser.AllocNode{
+				Name: forAlloc.MultiNames.Names[1].Name,
+				Val:  &parser.LoadArrayElement{Array: &parser.NameNode{Name: rangeItemName}, Pos: &parser.NameNode{Name: forKeyName}},
+			})
+		}
+	}
+	modifiedBlock = append(modifiedBlock, v.Block...)
+
+	typeCastedKey := &parser.TypeCastNode{
+		Type: &parser.SingleTypeNode{SourceName: "int32", TypeName: "int32"},
+		Val:  &parser.NameNode{Name: forKeyName},
+	}
+
+	c.compileForThreeType(&parser.ForNode{
+		BeforeLoop: &parser.AllocNode{Name: forKeyName, Val: &parser.ConstantNode{Type: parser.NUMBER, Value: 0}},
+
+		Condition: &parser.OperatorNode{
+			Left:     typeCastedKey,
+			Operator: parser.OP_LT,
+			Right: &parser.NameNode{
+				Name: forItemLenName,
+			},
+		},
+
+		AfterIteration: &parser.AssignNode{
+			Target: &parser.NameNode{Name: forKeyName},
+			Val: &parser.OperatorNode{
+				Left:     &parser.NameNode{Name: forKeyName},
+				Operator: parser.OP_ADD,
+				Right:    &parser.ConstantNode{Type: parser.NUMBER, Value: 1},
+			},
+		},
+
+		Block: modifiedBlock,
+	})
+}
+
 func (c *Compiler) compileBreakNode(v *parser.BreakNode) {
 	c.contextBlock.NewBr(c.contextLoopBreak[len(c.contextLoopBreak)-1])
 }

compiler/lexer/keywords.go

@@ -16,4 +16,5 @@ var keywords = map[string]struct{}{
 	"true":      {},
 	"false":     {},
 	"interface": {},
+	"range":     {},
 }

compiler/parser/for.go

@@ -2,38 +2,67 @@ package parser
 
 import (
 	"fmt"
-
 	"github.com/zegl/tre/compiler/lexer"
 )
 
+// ForNode creates a new for-loop
+type ForNode struct {
+	baseNode
+
+	BeforeLoop     Node
+	Condition      *OperatorNode
+	AfterIteration Node
+	Block          []Node
+
+	IsThreeTypeFor bool
+}
+
+func (f ForNode) String() string {
+	return fmt.Sprintf("For(%s; %s; %s) {\n\t%s\n}", f.BeforeLoop, f.Condition, f.AfterIteration, f.Block)
+}
+
 func (p *parser) parseFor() *ForNode {
 	res := &ForNode{}
 
 	p.i++
-	beforeLoop := p.parseUntil(lexer.Item{Type: lexer.SEPARATOR, Val: ";"})
+	beforeLoop, reachedItem := p.parseUntilEither([]lexer.Item{
+		{Type: lexer.SEPARATOR, Val: ";"}, // three type for
+		{Type: lexer.SEPARATOR, Val: "{"}, // range type for
+	})
+
 	if len(beforeLoop) != 1 {
 		panic("Expected only one beforeLoop in for loop")
 	}
-	res.BeforeLoop = beforeLoop[0]
 
-	p.i++
-	loopCondition := p.parseUntil(lexer.Item{Type: lexer.SEPARATOR, Val: ";"})
-	if len(loopCondition) != 1 {
-		panic("Expected only one condition in for loop")
+	isThreeTypeFor := false
+	if reachedItem.Val == ";" {
+		isThreeTypeFor = true
+		res.IsThreeTypeFor = true
 	}
 
-	if conditionNode, ok := loopCondition[0].(*OperatorNode); ok {
-		res.Condition = conditionNode
-	} else {
-		panic(fmt.Sprintf("Expected OperatorNode in for loop. Got: %T: %+v", loopCondition[0], loopCondition[0]))
-	}
 
-	p.i++
-	afterIteration := p.parseUntil(lexer.Item{Type: lexer.SEPARATOR, Val: "{"})
-	if len(afterIteration) != 1 {
-		panic("Expected only one afterIteration in for loop")
+	res.BeforeLoop = beforeLoop[0]
+
+	if isThreeTypeFor {
+		p.i++
+		loopCondition := p.parseUntil(lexer.Item{Type: lexer.SEPARATOR, Val: ";"})
+		if len(loopCondition) != 1 {
+			panic("Expected only one condition in for loop")
+		}
+
+		if conditionNode, ok := loopCondition[0].(*OperatorNode); ok {
+			res.Condition = conditionNode
+		} else {
+			panic(fmt.Sprintf("Expected OperatorNode in for loop. Got: %T: %+v", loopCondition[0], loopCondition[0]))
+		}
+
+		p.i++
+		afterIteration := p.parseUntil(lexer.Item{Type: lexer.SEPARATOR, Val: "{"})
+		if len(afterIteration) != 1 {
+			panic("Expected only one afterIteration in for loop")
+		}
+		res.AfterIteration = afterIteration[0]
 	}
-	res.AfterIteration = afterIteration[0]
 
 	p.i++
 	res.Block = p.parseUntil(lexer.Item{Type: lexer.SEPARATOR, Val: "}"})

compiler/parser/node.go

@@ -312,20 +312,6 @@ func (d DeclarePackageNode) String() string {
 	return "DeclarePackageNode(" + d.PackageName + ")"
 }
 
-// ForNode creates a new for-loop
-type ForNode struct {
-	baseNode
-
-	BeforeLoop     Node
-	Condition      *OperatorNode
-	AfterIteration Node
-	Block          []Node
-}
-
-func (f ForNode) String() string {
-	return fmt.Sprintf("For(%s; %s; %s) {\n\t%s\n}", f.BeforeLoop, f.Condition, f.AfterIteration, f.Block)
-}
-
 // BreakNode breaks out of the current for loop
 type BreakNode struct {
 	baseNode

compiler/parser/parser.go

@@ -485,6 +485,10 @@ func (p *parser) parseOne(withAheadParse bool) (res Node) {
 				Value: v,
 			}
 		}
+
+		if current.Val == "range" {
+			return p.parseRange()
+		}
 	}
 
 	p.printInput()
@@ -772,12 +776,24 @@ func (p *parser) lookAhead(steps int) lexer.Item {
 }
 
 func (p *parser) parseUntil(until lexer.Item) []Node {
-	var res []Node
+	n, _ := p.parseUntilEither([]lexer.Item{until})
+	return n
+}
 
+// parseUntilEither reads lexer items until it finds one that equals to a item in "untils"
+// The list of parsed nodes is returned in res. The lexer item that stopped the iteration
+// is returned in "reached"
+func (p *parser) parseUntilEither(untils []lexer.Item) (res []Node, reached lexer.Item) {
 	for {
 		current := p.input[p.i]
-		if current.Type == until.Type && current.Val == until.Val {
-			return res
+
+		log.Printf("current: %+v", current)
+
+		// Check if we have reached the end
+		for _, until := range untils {
+			if current.Type == until.Type && current.Val == until.Val {
+				return res, until
+			}
 		}
 
 		// Ignore comma

compiler/parser/range.go

@@ -0,0 +1,19 @@
+package parser
+
+import (
+	"fmt"
+)
+
+type RangeNode struct {
+	baseNode
+	Item Node
+}
+
+func (r RangeNode) String() string {
+	return fmt.Sprintf("range %s", r.Item)
+}
+
+func (p *parser) parseRange() *RangeNode {
+	p.i++
+	return &RangeNode{Item: p.parseOne(true)}
+}

compiler/testdata/slice-range.go

@@ -0,0 +1,46 @@
+package main
+
+import "external"
+
+
+func ff() []int {
+	return []int{40, 50}
+}
+
+func main() {
+	s := []int{10, 20, 30}
+	// 0 10
+	// 1 20
+	// 2 30
+	for k, v := range s {
+		external.Printf("%d %d\n", k, v)
+	}
+
+	// 0
+	// 1
+	// 2
+	for k := range s {
+	 	external.Printf("%d\n", k)
+	}
+
+	// AAA
+	// AAA
+	// AAA
+	for range s {
+		external.Printf("%s\n", "AAA")
+	}
+
+	// 0 40
+	// 1 50
+	f := ff()
+	for k, v := range f {
+		external.Printf("%d %d\n", k, v)
+	}
+
+	// 0 40
+	// 1 50
+	for k, v := range ff() {
+		external.Printf("%d %d\n", k, v)
+	}
+}
+