Add variables

ast/node.go

@@ -151,6 +151,13 @@ type ConditionalNode struct {
 	Exp2 Node
 }
 
+type VariableDeclaratorNode struct {
+	base
+	Name  string
+	Value Node
+	Expr  Node
+}
+
 type ArrayNode struct {
 	base
 	Nodes []Node

ast/print.go

@@ -130,6 +130,10 @@ func (n *PointerNode) String() string {
 	return "#"
 }
 
+func (n *VariableDeclaratorNode) String() string {
+	return fmt.Sprintf("let %s = %s; %s", n.Name, n.Value.String(), n.Expr.String())
+}
+
 func (n *ConditionalNode) String() string {
 	var cond, exp1, exp2 string
 	if _, ok := n.Cond.(*ConditionalNode); ok {

ast/visitor.go

@@ -45,6 +45,9 @@ func Walk(node *Node, v Visitor) {
 	case *ClosureNode:
 		Walk(&n.Node, v)
 	case *PointerNode:
+	case *VariableDeclaratorNode:
+		Walk(&n.Value, v)
+		Walk(&n.Expr, v)
 	case *ConditionalNode:
 		Walk(&n.Cond, v)
 		Walk(&n.Exp1, v)

checker/checker.go

@@ -58,10 +58,17 @@ func Check(tree *parser.Tree, config *conf.Config) (t reflect.Type, err error) {
 type visitor struct {
 	config      *conf.Config
 	collections []reflect.Type
+	scopes      []scope
 	parents     []ast.Node
 	err         *file.Error
 }
 
+type scope struct {
+	name  string
+	vtype reflect.Type
+	info  info
+}
+
 type info struct {
 	method bool
 	fn     *builtin.Function
@@ -104,6 +111,8 @@ func (v *visitor) visit(node ast.Node) (reflect.Type, info) {
 		t, i = v.ClosureNode(n)
 	case *ast.PointerNode:
 		t, i = v.PointerNode(n)
+	case *ast.VariableDeclaratorNode:
+		t, i = v.VariableDeclaratorNode(n)
 	case *ast.ConditionalNode:
 		t, i = v.ConditionalNode(n)
 	case *ast.ArrayNode:
@@ -135,6 +144,9 @@ func (v *visitor) NilNode(*ast.NilNode) (reflect.Type, info) {
 }
 
 func (v *visitor) IdentifierNode(node *ast.IdentifierNode) (reflect.Type, info) {
+	if s, ok := v.lookupVariable(node.Value); ok {
+		return s.vtype, s.info
+	}
 	if node.Value == "$env" {
 		return mapType, info{}
 	}
@@ -862,6 +874,23 @@ func (v *visitor) PointerNode(node *ast.PointerNode) (reflect.Type, info) {
 	return v.error(node, "cannot use %v as array", collection)
 }
 
+func (v *visitor) VariableDeclaratorNode(node *ast.VariableDeclaratorNode) (reflect.Type, info) {
+	vtype, vinfo := v.visit(node.Value)
+	v.scopes = append(v.scopes, scope{node.Name, vtype, vinfo})
+	t, i := v.visit(node.Expr)
+	v.scopes = v.scopes[:len(v.scopes)-1]
+	return t, i
+}
+
+func (v *visitor) lookupVariable(name string) (scope, bool) {
+	for i := len(v.scopes) - 1; i >= 0; i-- {
+		if v.scopes[i].name == name {
+			return v.scopes[i], true
+		}
+	}
+	return scope{}, false
+}
+
 func (v *visitor) ConditionalNode(node *ast.ConditionalNode) (reflect.Type, info) {
 	c, _ := v.visit(node.Cond)
 	if !isBool(c) && !isAny(c) {

checker/checker_test.go

@@ -125,6 +125,7 @@ var successTests = []string{
 	"Any.A?.B == nil",
 	"(Any.Bool ?? Bool) > 0",
 	"Bool ?? Bool",
+	"let foo = 1; foo == 1",
 }
 
 func TestCheck(t *testing.T) {

compiler/compiler.go

@@ -28,6 +28,7 @@ func Compile(tree *parser.Tree, config *conf.Config) (program *Program, err erro
 		locations:      make([]file.Location, 0),
 		constantsIndex: make(map[interface{}]int),
 		functionsIndex: make(map[string]int),
+		debugInfo:      make(map[string]string),
 	}
 
 	if config != nil {
@@ -50,30 +51,38 @@ func Compile(tree *parser.Tree, config *conf.Config) (program *Program, err erro
 		Node:      tree.Node,
 		Source:    tree.Source,
 		Locations: c.locations,
+		Variables: c.variables,
 		Constants: c.constants,
 		Bytecode:  c.bytecode,
 		Arguments: c.arguments,
 		Functions: c.functions,
-		FuncNames: c.functionNames,
+		DebugInfo: c.debugInfo,
 	}
 	return
 }
 
 type compiler struct {
 	locations      []file.Location
 	bytecode       []Opcode
+	variables      []interface{}
+	scopes         []scope
 	constants      []interface{}
 	constantsIndex map[interface{}]int
 	functions      []Function
-	functionNames  []string
 	functionsIndex map[string]int
+	debugInfo      map[string]string
 	mapEnv         bool
 	cast           reflect.Kind
 	nodes          []ast.Node
 	chains         [][]int
 	arguments      []int
 }
 
+type scope struct {
+	variableName string
+	index        int
+}
+
 func (c *compiler) emitLocation(loc file.Location, op Opcode, arg int) int {
 	c.bytecode = append(c.bytecode, op)
 	current := len(c.bytecode)
@@ -129,6 +138,13 @@ func (c *compiler) addConstant(constant interface{}) int {
 	return p
 }
 
+func (c *compiler) addVariable(name string) int {
+	c.variables = append(c.variables, nil)
+	p := len(c.variables) - 1
+	c.debugInfo[fmt.Sprintf("var_%d", p)] = name
+	return p
+}
+
 // emitFunction adds builtin.Function.Func to the program.Functions and emits call opcode.
 func (c *compiler) emitFunction(fn *builtin.Function, argsLen int) {
 	switch argsLen {
@@ -156,8 +172,8 @@ func (c *compiler) addFunction(fn *builtin.Function) int {
 	}
 	p := len(c.functions)
 	c.functions = append(c.functions, fn.Func)
-	c.functionNames = append(c.functionNames, fn.Name)
 	c.functionsIndex[fn.Name] = p
+	c.debugInfo[fmt.Sprintf("func_%d", p)] = fn.Name
 	return p
 }
 
@@ -209,6 +225,8 @@ func (c *compiler) compile(node ast.Node) {
 		c.ClosureNode(n)
 	case *ast.PointerNode:
 		c.PointerNode(n)
+	case *ast.VariableDeclaratorNode:
+		c.VariableDeclaratorNode(n)
 	case *ast.ConditionalNode:
 		c.ConditionalNode(n)
 	case *ast.ArrayNode:
@@ -227,6 +245,10 @@ func (c *compiler) NilNode(_ *ast.NilNode) {
 }
 
 func (c *compiler) IdentifierNode(node *ast.IdentifierNode) {
+	if index, ok := c.lookupVariable(node.Value); ok {
+		c.emit(OpLoadVar, index)
+		return
+	}
 	if node.Value == "$env" {
 		c.emit(OpLoadEnv)
 		return
@@ -742,10 +764,36 @@ func (c *compiler) ClosureNode(node *ast.ClosureNode) {
 	c.compile(node.Node)
 }
 
-func (c *compiler) PointerNode(node *ast.PointerNode) {
+func (c *compiler) PointerNode(_ *ast.PointerNode) {
 	c.emit(OpPointer)
 }
 
+func (c *compiler) VariableDeclaratorNode(node *ast.VariableDeclaratorNode) {
+	c.compile(node.Value)
+	index := c.addVariable(node.Name)
+	c.emit(OpStore, index)
+	c.beginScope(node.Name, index)
+	c.compile(node.Expr)
+	c.endScope()
+}
+
+func (c *compiler) beginScope(name string, index int) {
+	c.scopes = append(c.scopes, scope{name, index})
+}
+
+func (c *compiler) endScope() {
+	c.scopes = c.scopes[:len(c.scopes)-1]
+}
+
+func (c *compiler) lookupVariable(name string) (int, bool) {
+	for i := len(c.scopes) - 1; i >= 0; i-- {
+		if c.scopes[i].variableName == name {
+			return c.scopes[i].index, true
+		}
+	}
+	return 0, false
+}
+
 func (c *compiler) ConditionalNode(node *ast.ConditionalNode) {
 	c.compile(node.Cond)
 	otherwise := c.emit(OpJumpIfFalse, placeholder)

expr_test.go

@@ -948,6 +948,18 @@ func TestExpr(t *testing.T) {
 			`1 /* one */ + 2 // two`,
 			3,
 		},
+		{
+			`let x = 1; x + 2`,
+			3,
+		},
+		{
+			`map(1..3, let x = #; let y = x * x; y * y)`,
+			[]interface{}{1, 16, 81},
+		},
+		{
+			`map(1..2, let x = #; map(2..3, let y = #; x + y))`,
+			[]interface{}{[]interface{}{3, 4}, []interface{}{4, 5}},
+		},
 	}
 
 	for _, tt := range tests {

parser/lexer/lexer_test.go

@@ -188,6 +188,17 @@ func TestLex(t *testing.T) {
 				{Kind: EOF},
 			},
 		},
+		{
+			`let foo = bar;`,
+			[]Token{
+				{Kind: Operator, Value: "let"},
+				{Kind: Identifier, Value: "foo"},
+				{Kind: Operator, Value: "="},
+				{Kind: Identifier, Value: "bar"},
+				{Kind: Operator, Value: ";"},
+				{Kind: EOF},
+			},
+		},
 	}
 
 	for _, test := range tests {

parser/lexer/state.go

@@ -37,7 +37,7 @@ func root(l *lexer) stateFn {
 		l.emit(Bracket)
 	case strings.ContainsRune(")]}", r):
 		l.emit(Bracket)
-	case strings.ContainsRune("#,:%+-^", r): // single rune operator
+	case strings.ContainsRune("#,:;%+-^", r): // single rune operator
 		l.emit(Operator)
 	case strings.ContainsRune("&!=*<>", r): // possible double rune operator
 		l.accept("&=*")
@@ -124,6 +124,8 @@ loop:
 				return not
 			case "in", "or", "and", "matches", "contains", "startsWith", "endsWith":
 				l.emit(Operator)
+			case "let":
+				l.emit(Operator)
 			default:
 				l.emit(Identifier)
 			}

parser/parser.go

@@ -110,6 +110,12 @@ func (p *parser) expect(kind Kind, values ...string) {
 // parse functions
 
 func (p *parser) parseExpression(precedence int) Node {
+	if precedence == 0 {
+		if p.current.Is(Operator, "let") {
+			return p.parseVariableDeclaration()
+		}
+	}
+
 	nodeLeft := p.parsePrimary()
 
 	prevOperator := ""
@@ -179,6 +185,21 @@ func (p *parser) parseExpression(precedence int) Node {
 	return nodeLeft
 }
 
+func (p *parser) parseVariableDeclaration() Node {
+	p.expect(Operator, "let")
+	variableName := p.current
+	p.expect(Identifier)
+	p.expect(Operator, "=")
+	value := p.parseExpression(0)
+	p.expect(Operator, ";")
+	node := p.parseExpression(0)
+	return &VariableDeclaratorNode{
+		Name:  variableName.Value,
+		Value: value,
+		Expr:  node,
+	}
+}
+
 func (p *parser) parseConditional(node Node) Node {
 	var expr1, expr2 Node
 	for p.current.Is(Operator, "?") && p.err == nil {

parser/parser_test.go

@@ -452,6 +452,17 @@ func TestParse(t *testing.T) {
 				Callee: &IdentifierNode{Value: "ok"},
 				Arguments: []Node{
 					&BoolNode{Value: true}}}},
+		{
+			`let foo = a + b; foo + c`,
+			&VariableDeclaratorNode{
+				Name: "foo",
+				Value: &BinaryNode{Operator: "+",
+					Left:  &IdentifierNode{Value: "a"},
+					Right: &IdentifierNode{Value: "b"}},
+				Expr: &BinaryNode{Operator: "+",
+					Left:  &IdentifierNode{Value: "foo"},
+					Right: &IdentifierNode{Value: "c"}}},
+		},
 	}
 	for _, test := range parseTests {
 		actual, err := parser.ParseWithConfig(test.input, &conf.Config{Pipes: true})

vm/opcodes.go

@@ -7,6 +7,8 @@ const (
 	OpPush
 	OpPushInt
 	OpPop
+	OpStore
+	OpLoadVar
 	OpLoadConst
 	OpLoadField
 	OpLoadFast