bake: allow user functions in variables and vice-versa

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

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
+}