allow initialization of external and TLA code variables from an AST node.

This allows configuration of a VM with an external code variable that has been pre-parsed into an AST node.
This allows a caller to pre-parse a large object into a node and re-use that for multiple VMs. The cost
of converting the object to an ast.Node is incurred only once by the caller as opposed to having this
eagerly incurred (whether or not the object is used) before the eval.

Author: gotwarlost

imports.go

@@ -139,6 +139,15 @@ func (cache *importCache) importString(importedFrom, importedPath string, i *int
 	return makeValueString(data.String()), nil
 }
 
+func nodeToPV(i *interpreter, filename string, node ast.Node) *cachedThunk {
+	env := makeInitialEnv(filename, i.baseStd)
+	return &cachedThunk{
+		env:     &env,
+		body:    node,
+		content: nil,
+	}
+}
+
 func codeToPV(i *interpreter, filename string, code string) *cachedThunk {
 	node, err := program.SnippetToAST(ast.DiagnosticFileName(filename), "", code)
 	if err != nil {
@@ -149,12 +158,7 @@ func codeToPV(i *interpreter, filename string, code string) *cachedThunk {
 		// The same thinking applies to external variables.
 		return &cachedThunk{err: err}
 	}
-	env := makeInitialEnv(filename, i.baseStd)
-	return &cachedThunk{
-		env:     &env,
-		body:    node,
-		content: nil,
-	}
+	return nodeToPV(i, filename, node)
 }
 
 // ImportCode imports code from a path.

interpreter.go

@@ -1194,9 +1194,13 @@ func evaluateStd(i *interpreter) (value, error) {
 func prepareExtVars(i *interpreter, ext vmExtMap, kind string) map[string]*cachedThunk {
 	result := make(map[string]*cachedThunk)
 	for name, content := range ext {
-		if content.isCode {
-			result[name] = codeToPV(i, "<"+kind+":"+name+">", content.value)
-		} else {
+		diagnosticFile := "<" + kind + ":" + name + ">"
+		switch content.kind {
+		case extKindCode:
+			result[name] = codeToPV(i, diagnosticFile, content.value)
+		case extKindNode:
+			result[name] = nodeToPV(i, diagnosticFile, content.node)
+		default:
 			result[name] = readyThunk(makeValueString(content.value))
 		}
 	}

jsonnet_test.go

@@ -2,6 +2,7 @@ package jsonnet
 
 import (
 	"bytes"
+	"encoding/json"
 	"reflect"
 	"strings"
 	"testing"
@@ -265,3 +266,64 @@ func TestTLAReset(t *testing.T) {
 		t.Errorf("unexpected error %v", err)
 	}
 }
+
+func assertVarOutput(t *testing.T, jsonStr string) {
+	var data struct {
+		Var  string `json:"var"`
+		Code string `json:"code"`
+		Node string `json:"node"`
+	}
+	err := json.Unmarshal([]byte(jsonStr), &data)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if data.Var != "var" {
+		t.Errorf("var attribute not correct, want '%s' got '%s'", "var", data.Var)
+	}
+	if data.Code != "code" {
+		t.Errorf("code attribute not correct, want '%s' got '%s'", "code", data.Code)
+	}
+	if data.Node != "node" {
+		t.Errorf("node attribute not correct, want '%s' got '%s'", "node", data.Node)
+	}
+}
+
+func TestExtTypes(t *testing.T) {
+	node, err := SnippetToAST("var.jsonnet", `{ node: 'node' }`)
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	vm := MakeVM()
+	vm.ExtVar("var", "var")
+	vm.ExtCode("code", `{ code: 'code'}`)
+	vm.ExtNode("node", node)
+
+	jsonStr, err := vm.EvaluateAnonymousSnippet(
+		"caller.jsonnet",
+		`{ var: std.extVar('var') } + std.extVar('code') + std.extVar('node')`,
+	)
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	assertVarOutput(t, jsonStr)
+}
+
+func TestTLATypes(t *testing.T) {
+	node, err := SnippetToAST("var.jsonnet", `{ node: 'node' }`)
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	vm := MakeVM()
+	vm.TLAVar("var", "var")
+	vm.TLACode("code", `{ code: 'code'}`)
+	vm.TLANode("node", node)
+
+	jsonStr, err := vm.EvaluateAnonymousSnippet(
+		"caller.jsonnet",
+		`function (var, code, node) { var: var } + code + node`,
+	)
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	assertVarOutput(t, jsonStr)
+}

vm.go

@@ -46,13 +46,23 @@ type VM struct {
 	importCache    *importCache
 }
 
+// extKind indicates the kind of external variable that is being initialized for the VM
+type extKind int
+
+const (
+	extKindVar  extKind = iota // a simple string
+	extKindCode                // a code snippet represented as a string
+	extKindNode                // an ast.Node that is passed in
+)
+
 // External variable or top level argument provided before execution
 type vmExt struct {
-	// jsonnet code to evaluate or string to pass
+	// the kind of external variable that is specified.
+	kind extKind
+	// jsonnet code to evaluate (kind=extKindCode) or string to pass (kind=extKindVar)
 	value string
-	// isCode determines whether it should be evaluated as jsonnet code or
-	// treated as string.
-	isCode bool
+	// the specified node for kind=extKindNode
+	node ast.Node
 }
 
 type vmExtMap map[string]vmExt
@@ -85,13 +95,19 @@ func (vm *VM) flushValueCache() {
 
 // ExtVar binds a Jsonnet external var to the given value.
 func (vm *VM) ExtVar(key string, val string) {
-	vm.ext[key] = vmExt{value: val, isCode: false}
+	vm.ext[key] = vmExt{value: val, kind: extKindVar}
 	vm.flushValueCache()
 }
 
 // ExtCode binds a Jsonnet external code var to the given code.
 func (vm *VM) ExtCode(key string, val string) {
-	vm.ext[key] = vmExt{value: val, isCode: true}
+	vm.ext[key] = vmExt{value: val, kind: extKindCode}
+	vm.flushValueCache()
+}
+
+// ExtNode binds a Jsonnet external code var to the given AST node.
+func (vm *VM) ExtNode(key string, node ast.Node) {
+	vm.ext[key] = vmExt{node: node, kind: extKindNode}
 	vm.flushValueCache()
 }
 
@@ -103,15 +119,21 @@ func (vm *VM) ExtReset() {
 
 // TLAVar binds a Jsonnet top level argument to the given value.
 func (vm *VM) TLAVar(key string, val string) {
-	vm.tla[key] = vmExt{value: val, isCode: false}
+	vm.tla[key] = vmExt{value: val, kind: extKindVar}
 	// Setting a TLA does not require flushing the cache.
 	// Only the results of evaluation of imported files are cached
 	// and the TLAs do not affect these unlike extVars.
 }
 
 // TLACode binds a Jsonnet top level argument to the given code.
 func (vm *VM) TLACode(key string, val string) {
-	vm.tla[key] = vmExt{value: val, isCode: true}
+	vm.tla[key] = vmExt{value: val, kind: extKindCode}
+	// Setting a TLA does not require flushing the cache - see above.
+}
+
+// TLANode binds a Jsonnet top level argument to the given AST node.
+func (vm *VM) TLANode(key string, node ast.Node) {
+	vm.tla[key] = vmExt{node: node, kind: extKindNode}
 	// Setting a TLA does not require flushing the cache - see above.
 }