From b2e6e2d0d05121e5ef845a0cea93f19f9c762a21 Mon Sep 17 00:00:00 2001
From: Martin Atkins <mart@degeneration.co.uk>
Date: Fri, 16 Jun 2017 07:28:29 -0700
Subject: [PATCH] zclsyntax: look in parent EvalContexts for functions

If we can't find a function in the given EvalContext, we must traverse
the context chain until either we run out of contexts or we find a
matching function.

At present there is no reason for a non-root context to have any
functions, so this will always traverse to the root. This may change in
future if we introduce constructs that define local functions.
---
 zcl/eval_context.go              |  6 ++++++
 zcl/zclsyntax/expression.go      | 36 ++++++++++++++++++++++----------
 zcl/zclsyntax/expression_test.go | 12 +++++++++++
 3 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/zcl/eval_context.go b/zcl/eval_context.go
index a956533..d2e4961 100644
--- a/zcl/eval_context.go
+++ b/zcl/eval_context.go
@@ -17,3 +17,9 @@ type EvalContext struct {
 func (ctx *EvalContext) NewChild() *EvalContext {
 	return &EvalContext{parent: ctx}
 }
+
+// Parent returns the parent of the receiver, or nil if the receiver has
+// no parent.
+func (ctx *EvalContext) Parent() *EvalContext {
+	return ctx.parent
+}
diff --git a/zcl/zclsyntax/expression.go b/zcl/zclsyntax/expression.go
index 953b3b6..d3d7bdf 100644
--- a/zcl/zclsyntax/expression.go
+++ b/zcl/zclsyntax/expression.go
@@ -121,21 +121,35 @@ func (e *FunctionCallExpr) walkChildNodes(w internalWalkFunc) {
 func (e *FunctionCallExpr) Value(ctx *zcl.EvalContext) (cty.Value, zcl.Diagnostics) {
 	var diags zcl.Diagnostics
 
-	if ctx == nil || ctx.Functions == nil {
-		return cty.DynamicVal, zcl.Diagnostics{
-			{
-				Severity: zcl.DiagError,
-				Summary:  "Function calls not allowed",
-				Detail:   "Functions may not be called here.",
-				Subject:  &e.NameRange,
-				Context:  e.Range().Ptr(),
-			},
+	var f function.Function
+	exists := false
+	hasNonNilMap := false
+	thisCtx := ctx
+	for thisCtx != nil {
+		if thisCtx.Functions == nil {
+			thisCtx = thisCtx.Parent()
+			continue
 		}
+		hasNonNilMap = true
+		f, exists = thisCtx.Functions[e.Name]
+		if exists {
+			break
+		}
+		thisCtx = thisCtx.Parent()
 	}
 
-	// FIXME: also need to look in ctx.parent, etc
-	f, exists := ctx.Functions[e.Name]
 	if !exists {
+		if !hasNonNilMap {
+			return cty.DynamicVal, zcl.Diagnostics{
+				{
+					Severity: zcl.DiagError,
+					Summary:  "Function calls not allowed",
+					Detail:   "Functions may not be called here.",
+					Subject:  e.Range().Ptr(),
+				},
+			}
+		}
+
 		avail := make([]string, 0, len(ctx.Functions))
 		for name := range ctx.Functions {
 			avail = append(avail, name)
diff --git a/zcl/zclsyntax/expression_test.go b/zcl/zclsyntax/expression_test.go
index 3b45752..3dda1bc 100644
--- a/zcl/zclsyntax/expression_test.go
+++ b/zcl/zclsyntax/expression_test.go
@@ -417,6 +417,18 @@ upper(
 			}),
 			0,
 		},
+		{
+			`{for k, v in {hello: "world"}: upper(k) => upper(v) if k == "hello"}`,
+			&zcl.EvalContext{
+				Functions: map[string]function.Function{
+					"upper": stdlib.UpperFunc,
+				},
+			},
+			cty.ObjectVal(map[string]cty.Value{
+				"HELLO": cty.StringVal("WORLD"),
+			}),
+			0,
+		},
 		{
 			`{for k, v in ["world"]: k => v if k == 0}`,
 			nil,