bake: allow user functions in variables and vice-versa

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>

util/userfunc/decode.go

@@ -2,7 +2,6 @@ package userfunc
 
 import (
 	"github.com/hashicorp/hcl/v2"
-	"github.com/zclconf/go-cty/cty"
 	"github.com/zclconf/go-cty/cty/function"
 )
 
@@ -53,7 +52,7 @@ func decodeUserFunctions(body hcl.Body, blockType string, contextFunc ContextFun
 	}
 
 	funcs = make(map[string]function.Function)
-Blocks:
+
 	for _, block := range content.Blocks {
 		name := block.Labels[0]
 		funcContent, funcDiags := block.Body.Content(funcBodySchema)
@@ -68,88 +67,12 @@ Blocks:
 		if funcContent.Attributes["variadic_param"] != nil {
 			varParamExpr = funcContent.Attributes["variadic_param"].Expr
 		}
-
-		var params []string
-		var varParam string
-
-		paramExprs, paramsDiags := hcl.ExprList(paramsExpr)
-		diags = append(diags, paramsDiags...)
-		if paramsDiags.HasErrors() {
+		f, funcDiags := NewFunction(paramsExpr, varParamExpr, resultExpr, getBaseCtx)
+		if funcDiags.HasErrors() {
+			diags = append(diags, funcDiags...)
 			continue
 		}
-		for _, paramExpr := range paramExprs {
-			param := hcl.ExprAsKeyword(paramExpr)
-			if param == "" {
-				diags = append(diags, &hcl.Diagnostic{
-					Severity: hcl.DiagError,
-					Summary:  "Invalid param element",
-					Detail:   "Each parameter name must be an identifier.",
-					Subject:  paramExpr.Range().Ptr(),
-				})
-				continue Blocks
-			}
-			params = append(params, param)
-		}
-
-		if varParamExpr != nil {
-			varParam = hcl.ExprAsKeyword(varParamExpr)
-			if varParam == "" {
-				diags = append(diags, &hcl.Diagnostic{
-					Severity: hcl.DiagError,
-					Summary:  "Invalid variadic_param",
-					Detail:   "The variadic parameter name must be an identifier.",
-					Subject:  varParamExpr.Range().Ptr(),
-				})
-				continue
-			}
-		}
-
-		spec := &function.Spec{}
-		for _, paramName := range params {
-			spec.Params = append(spec.Params, function.Parameter{
-				Name: paramName,
-				Type: cty.DynamicPseudoType,
-			})
-		}
-		if varParamExpr != nil {
-			spec.VarParam = &function.Parameter{
-				Name: varParam,
-				Type: cty.DynamicPseudoType,
-			}
-		}
-		impl := func(args []cty.Value) (cty.Value, error) {
-			ctx := getBaseCtx()
-			ctx = ctx.NewChild()
-			ctx.Variables = make(map[string]cty.Value)
-
-			// The cty function machinery guarantees that we have at least
-			// enough args to fill all of our params.
-			for i, paramName := range params {
-				ctx.Variables[paramName] = args[i]
-			}
-			if spec.VarParam != nil {
-				varArgs := args[len(params):]
-				ctx.Variables[varParam] = cty.TupleVal(varArgs)
-			}
-
-			result, diags := resultExpr.Value(ctx)
-			if diags.HasErrors() {
-				// Smuggle the diagnostics out via the error channel, since
-				// a diagnostics sequence implements error. Caller can
-				// type-assert this to recover the individual diagnostics
-				// if desired.
-				return cty.DynamicVal, diags
-			}
-			return result, nil
-		}
-		spec.Type = func(args []cty.Value) (cty.Type, error) {
-			val, err := impl(args)
-			return val.Type(), err
-		}
-		spec.Impl = func(args []cty.Value, retType cty.Type) (cty.Value, error) {
-			return impl(args)
-		}
-		funcs[name] = function.New(spec)
+		funcs[name] = f
 	}
 
 	return funcs, remain, diags

util/userfunc/public.go

@@ -2,6 +2,7 @@ package userfunc
 
 import (
 	"github.com/hashicorp/hcl/v2"
+	"github.com/zclconf/go-cty/cty"
 	"github.com/zclconf/go-cty/cty/function"
 )
 
@@ -40,3 +41,85 @@ type ContextFunc func() *hcl.EvalContext
 func DecodeUserFunctions(body hcl.Body, blockType string, context ContextFunc) (funcs map[string]function.Function, remain hcl.Body, diags hcl.Diagnostics) {
 	return decodeUserFunctions(body, blockType, context)
 }
+
+// NewFunction creates a new function instance from preparsed HCL expressions.
+func NewFunction(paramsExpr, varParamExpr, resultExpr hcl.Expression, getBaseCtx func() *hcl.EvalContext) (function.Function, hcl.Diagnostics) {
+	var params []string
+	var varParam string
+
+	paramExprs, paramsDiags := hcl.ExprList(paramsExpr)
+	if paramsDiags.HasErrors() {
+		return function.Function{}, paramsDiags
+	}
+	for _, paramExpr := range paramExprs {
+		param := hcl.ExprAsKeyword(paramExpr)
+		if param == "" {
+			return function.Function{}, hcl.Diagnostics{{
+				Severity: hcl.DiagError,
+				Summary:  "Invalid param element",
+				Detail:   "Each parameter name must be an identifier.",
+				Subject:  paramExpr.Range().Ptr(),
+			}}
+		}
+		params = append(params, param)
+	}
+
+	if varParamExpr != nil {
+		varParam = hcl.ExprAsKeyword(varParamExpr)
+		if varParam == "" {
+			return function.Function{}, hcl.Diagnostics{{
+				Severity: hcl.DiagError,
+				Summary:  "Invalid variadic_param",
+				Detail:   "The variadic parameter name must be an identifier.",
+				Subject:  varParamExpr.Range().Ptr(),
+			}}
+		}
+	}
+
+	spec := &function.Spec{}
+	for _, paramName := range params {
+		spec.Params = append(spec.Params, function.Parameter{
+			Name: paramName,
+			Type: cty.DynamicPseudoType,
+		})
+	}
+	if varParamExpr != nil {
+		spec.VarParam = &function.Parameter{
+			Name: varParam,
+			Type: cty.DynamicPseudoType,
+		}
+	}
+	impl := func(args []cty.Value) (cty.Value, error) {
+		ctx := getBaseCtx()
+		ctx = ctx.NewChild()
+		ctx.Variables = make(map[string]cty.Value)
+
+		// The cty function machinery guarantees that we have at least
+		// enough args to fill all of our params.
+		for i, paramName := range params {
+			ctx.Variables[paramName] = args[i]
+		}
+		if spec.VarParam != nil {
+			varArgs := args[len(params):]
+			ctx.Variables[varParam] = cty.TupleVal(varArgs)
+		}
+
+		result, diags := resultExpr.Value(ctx)
+		if diags.HasErrors() {
+			// Smuggle the diagnostics out via the error channel, since
+			// a diagnostics sequence implements error. Caller can
+			// type-assert this to recover the individual diagnostics
+			// if desired.
+			return cty.DynamicVal, diags
+		}
+		return result, nil
+	}
+	spec.Type = func(args []cty.Value) (cty.Type, error) {
+		val, err := impl(args)
+		return val.Type(), err
+	}
+	spec.Impl = func(args []cty.Value, retType cty.Type) (cty.Value, error) {
+		return impl(args)
+	}
+	return function.New(spec), nil
+}