From 0ff39a3773144a40a3250f4c10f7927817ba64e3 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Tue, 14 Apr 2015 09:09:18 -0700
Subject: [PATCH] config/lang: support math on variables through implicits

---
 config/lang/check_types.go | 33 +++++++++++++++++++++------------
 config/lang/eval_test.go   | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 12 deletions(-)

diff --git a/config/lang/check_types.go b/config/lang/check_types.go
index 0396eb1f3..4fbcd731a 100644
--- a/config/lang/check_types.go
+++ b/config/lang/check_types.go
@@ -100,20 +100,29 @@ func (tc *typeCheckArithmetic) TypeCheck(v *TypeCheck) (ast.Node, error) {
 		exprs[len(tc.n.Exprs)-1-i] = v.StackPop()
 	}
 
-	// Determine the resulting type we want
+	// Determine the resulting type we want. We do this by going over
+	// every expression until we find one with a type we recognize.
+	// We do this because the first expr might be a string ("var.foo")
+	// and we need to know what to implicit to.
 	mathFunc := "__builtin_IntMath"
 	mathType := ast.TypeInt
-	switch v := exprs[0]; v {
-	case ast.TypeInt:
-		mathFunc = "__builtin_IntMath"
-		mathType = v
-	case ast.TypeFloat:
-		mathFunc = "__builtin_FloatMath"
-		mathType = v
-	default:
-		return nil, fmt.Errorf(
-			"Math operations can only be done with ints and floats, got %s",
-			v)
+	for _, v := range exprs {
+		exit := true
+		switch v {
+		case ast.TypeInt:
+			mathFunc = "__builtin_IntMath"
+			mathType = v
+		case ast.TypeFloat:
+			mathFunc = "__builtin_FloatMath"
+			mathType = v
+		default:
+			exit = false
+		}
+
+		// We found the type, so leave
+		if exit {
+			break
+		}
 	}
 
 	// Verify the args
diff --git a/config/lang/eval_test.go b/config/lang/eval_test.go
index 450a8abce..44f25d6fd 100644
--- a/config/lang/eval_test.go
+++ b/config/lang/eval_test.go
@@ -134,6 +134,40 @@ func TestEval(t *testing.T) {
 			ast.TypeString,
 		},
 
+		{
+			"foo ${bar+1}",
+			&ast.BasicScope{
+				VarMap: map[string]ast.Variable{
+					"bar": ast.Variable{
+						Value: "41",
+						Type:  ast.TypeString,
+					},
+				},
+			},
+			false,
+			"foo 42",
+			ast.TypeString,
+		},
+
+		{
+			"foo ${bar+baz}",
+			&ast.BasicScope{
+				VarMap: map[string]ast.Variable{
+					"bar": ast.Variable{
+						Value: "41",
+						Type:  ast.TypeString,
+					},
+					"baz": ast.Variable{
+						Value: "1",
+						Type:  ast.TypeString,
+					},
+				},
+			},
+			false,
+			"foo 42",
+			ast.TypeString,
+		},
+
 		{
 			"foo ${rand()}",
 			&ast.BasicScope{