compiler, runtime: extend the selector expression to maps in templates

This change allows to use the selector expression, in templates, to
access map keys.

Close #937

Co-authored-by: Gianluca Mondini <mondini@open2b.com>

Author: gazerro

internal/compiler/builder_instructions.go

@@ -438,6 +438,9 @@ func (fb *functionBuilder) emitIndex(ki bool, expr, i, dst int8, t reflect.Type,
 	case reflect.Map:
 		op = runtime.OpMapIndex
 		fb.addOperandKinds(0, t.Key().Kind(), t.Elem().Kind())
+	case reflect.Interface:
+		op = runtime.OpMapIndexAny
+		fb.addPosAndPath(pos)
 	case reflect.String:
 		op = runtime.OpIndexString
 	default:

internal/compiler/checker_assignment.go

@@ -53,7 +53,7 @@ func (tc *typechecker) checkAssignment(node *ast.Assignment) {
 		} else {
 			lh = tc.checkExpr(lhExpr)
 		}
-		tc.checkAssignTo(lh, lhExpr)
+		node.Lhs[i] = tc.checkAssignTo(lh, lhExpr)
 		lhs[i] = lh
 	}
 
@@ -107,7 +107,7 @@ func (tc *typechecker) checkAssignmentOperation(node *ast.Assignment) {
 	lh := tc.checkExpr(node.Lhs[0])
 	rh := tc.checkExpr(node.Rhs[0])
 
-	tc.checkAssignTo(lh, node.Lhs[0])
+	node.Lhs[0] = tc.checkAssignTo(lh, node.Lhs[0])
 
 	op := operatorFromAssignmentType(node.Type)
 	_, err := tc.binaryOp(node.Lhs[0], op, node.Rhs[0])
@@ -254,7 +254,7 @@ func (tc *typechecker) checkIncDecStatement(node *ast.Assignment) {
 	}
 
 	lh := tc.checkExpr(node.Lhs[0])
-	tc.checkAssignTo(lh, node.Lhs[0])
+	node.Lhs[0] = tc.checkAssignTo(lh, node.Lhs[0])
 
 	if !isNumeric(lh.Type.Kind()) {
 		panic(tc.errorf(node, "invalid operation: %s (non-numeric type %s)", node, lh))
@@ -466,14 +466,29 @@ func (tc *typechecker) declareVariable(lh *ast.Identifier, typ reflect.Type) {
 }
 
 // checkAssignTo checks that it is possible to assign to the expression expr.
-func (tc *typechecker) checkAssignTo(ti *typeInfo, expr ast.Expression) {
-	if ti.Addressable() && !ti.IsMacroDeclaration() || tc.isMapIndexing(expr) {
-		return
+// The caller must replace expr in the AST with the returned expression.
+func (tc *typechecker) checkAssignTo(ti *typeInfo, expr ast.Expression) ast.Expression {
+	if ti.IsKeySelector() {
+		if ti.replacement != nil {
+			expr = ti.replacement.(ast.Expression)
+			return expr
+		}
+	} else if ti.Addressable() && !ti.IsMacroDeclaration() || tc.isMapIndexing(expr) {
+		return expr
 	}
 	format := "cannot assign to %s"
 	switch e := expr.(type) {
 	case *ast.Selector:
-		if tc.isMapIndexing(e.Expr) {
+		if ti.IsKeySelector() {
+			for {
+				if t := tc.compilation.typeInfos[e.Expr]; t.replacement != nil {
+					break
+				}
+				e = e.Expr.(*ast.Selector)
+			}
+			expr = e.Expr
+			format = "cannot index %s (map index expression of type interface{})"
+		} else if tc.isMapIndexing(e.Expr) {
 			format = "cannot assign to struct field %s in map"
 		}
 	case *ast.Index:
@@ -592,6 +607,14 @@ func (tc *typechecker) rebalancedRightSide(node ast.Node) []ast.Expression {
 			tc.compilation.typeInfos[v1] = &typeInfo{Type: ti.Type}
 			tc.compilation.typeInfos[v2] = untypedBoolTypeInfo
 			return []ast.Expression{v1, v2}
+		case *ast.Selector:
+			if ti := tc.checkExpr(rhExpr); ti.IsKeySelector() {
+				v1 := ast.NewSelector(v.Pos(), v.Expr, v.Ident)
+				v2 := ast.NewSelector(v.Pos(), v.Expr, v.Ident)
+				tc.compilation.typeInfos[v1] = &typeInfo{Type: ti.Type, Properties: ti.Properties, replacement: ti.replacement}
+				tc.compilation.typeInfos[v2] = untypedBoolTypeInfo
+				return []ast.Expression{v1, v2}
+			}
 		case *ast.Index:
 			v1 := ast.NewIndex(v.Pos(), v.Expr, v.Index)
 			v2 := ast.NewIndex(v.Pos(), v.Expr, v.Index)

internal/compiler/checker_expressions.go

@@ -749,6 +749,10 @@ func (tc *typechecker) typeof(expr ast.Expression, typeExpected bool) *typeInfo
 		if mv, ok := tc.checkMethodValue(t, expr); ok {
 			return mv
 		}
+		// Key selector.
+		if ti, ok := tc.checkKeySelector(t, expr); ok {
+			return ti
+		}
 		// Field selector.
 		return tc.checkFieldSelector(t, expr)
 
@@ -2550,6 +2554,30 @@ func (tc *typechecker) checkMethodValue(t *typeInfo, expr *ast.Selector) (*typeI
 	}, true
 }
 
+// checkKeySelector checks a key selector.
+func (tc *typechecker) checkKeySelector(t *typeInfo, expr *ast.Selector) (*typeInfo, bool) {
+
+	if tc.opts.mod != templateMod {
+		return nil, false
+	}
+
+	switch {
+	case t.Type.Kind() == reflect.Map:
+		// Type must be 'map[K]E' where K is a string type or the 'interface{}' type and E is any type.
+		if k := t.Type.Key(); t.Type.Name() != "" || k.Kind() != reflect.String && k != emptyInterfaceType {
+			panic(tc.errorf(expr, "invalid operation: cannot select %s (type %s does not support key selection)", expr, t.Type))
+		}
+		// Remember to replace 'm.x' with 'm["x"]'.
+		replacement := ast.NewIndex(expr.Pos(), expr.Expr, ast.NewBasicLiteral(expr.Pos(), ast.StringLiteral, "`"+expr.Ident+"`"))
+		tc.checkExpr(replacement)
+		return &typeInfo{Type: t.Type.Elem(), Properties: propertyKeySelector, replacement: replacement}, true
+	case t.IsKeySelector() && t.Type == emptyInterfaceType:
+		return &typeInfo{Type: t.Type, Properties: t.Properties}, true
+	}
+
+	return nil, false
+}
+
 // checkFieldSelector checks a field selector.
 func (tc *typechecker) checkFieldSelector(t *typeInfo, expr *ast.Selector) *typeInfo {
 

internal/compiler/checker_template_test.go

@@ -44,11 +44,17 @@ var stringArrayTypeInfo = &typeInfo{Type: reflect.ArrayOf(2, stringType), Proper
 var boolSliceTypeInfo = &typeInfo{Type: reflect.SliceOf(boolType), Properties: propertyAddressable}
 var boolArrayTypeInfo = &typeInfo{Type: reflect.ArrayOf(2, boolType), Properties: propertyAddressable}
 var interfaceSliceTypeInfo = &typeInfo{Type: reflect.SliceOf(emptyInterfaceType), Properties: propertyAddressable}
+var stringToAnyMapTypeInfo = &typeInfo{Type: reflect.MapOf(stringType, emptyInterfaceType), Properties: propertyAddressable}
 
 var stringToIntMapTypeInfo = &typeInfo{Type: reflect.MapOf(stringType, intType), Properties: propertyAddressable}
 var intToStringMapTypeInfo = &typeInfo{Type: reflect.MapOf(intType, stringType), Properties: propertyAddressable}
+var intToAnyMapTypeInfo = &typeInfo{Type: reflect.MapOf(intType, emptyInterfaceType), Properties: propertyAddressable}
+var anyToIntMapTypeInfo = &typeInfo{Type: reflect.MapOf(emptyInterfaceType, intType), Properties: propertyAddressable}
+var definedStringToStringMapTypeInfo = &typeInfo{Type: reflect.MapOf(definedStringTypeInfo.Type, stringType), Properties: propertyAddressable}
 var definedIntToStringMapTypeInfo = &typeInfo{Type: reflect.MapOf(definedIntTypeInfo.Type, stringType), Properties: propertyAddressable}
 
+var stringToStringToIntMapTypeInfo = &typeInfo{Type: reflect.MapOf(stringType, reflect.MapOf(stringType, intType)), Properties: propertyAddressable}
+
 var definedIntTypeInfo = &typeInfo{Type: reflect.TypeOf(definedInt(0)), Properties: propertyAddressable}
 var definedIntSliceTypeInfo = &typeInfo{Type: reflect.SliceOf(definedIntTypeInfo.Type), Properties: propertyAddressable}
 
@@ -70,6 +76,12 @@ func tiMarkdown() *typeInfo {
 	return &typeInfo{Type: formatTypes[ast.FormatMarkdown]}
 }
 
+type TF struct {
+	F int
+}
+
+var stringToTFMapTypeInfo = &typeInfo{Type: reflect.MapOf(stringType, reflect.TypeOf(TF{})), Properties: propertyAddressable}
+
 var checkerTemplateExprs = []struct {
 	src   string
 	ti    *typeInfo
@@ -156,6 +168,16 @@ var checkerTemplateExprs = []struct {
 	// conversion from markdown to html
 	{`html(a)`, tiHTMLConst("<h1>title</h1>"), map[string]*typeInfo{"a": tiMarkdownConst("# title")}},
 	{`html(a)`, tiHTML(), map[string]*typeInfo{"a": tiMarkdown()}},
+
+	// key selector
+	{`m.x`, tiInt(), map[string]*typeInfo{"m": stringToIntMapTypeInfo}},
+	{`m.x`, tiString(), map[string]*typeInfo{"m": definedStringToStringMapTypeInfo}},
+	{`m.x`, tiInt(), map[string]*typeInfo{"m": anyToIntMapTypeInfo}},
+	{`m.x.y`, tiInt(), map[string]*typeInfo{"m": stringToStringToIntMapTypeInfo}},
+	{`m.x.y`, tiInterface(), map[string]*typeInfo{"m": stringToAnyMapTypeInfo}},
+	{`m.x.y.z`, tiInterface(), map[string]*typeInfo{"m": stringToAnyMapTypeInfo}},
+	{`m.x.nil`, tiInterface(), map[string]*typeInfo{"m": stringToAnyMapTypeInfo}},
+	{`m.x.F`, tiInt(), map[string]*typeInfo{"m": stringToTFMapTypeInfo}},
 }
 
 func TestCheckerTemplateExpressions(t *testing.T) {
@@ -236,6 +258,14 @@ var checkerTemplateExprErrors = []struct {
 	// slicing of a format type
 	{`a[1:2]`, tierr(1, 5, `invalid operation a[1:2] (slice of compiler.html)`), map[string]*typeInfo{"a": tiHTMLConst("<b>a</b>")}},
 	{`a[1:2]`, tierr(1, 5, `invalid operation a[1:2] (slice of compiler.html)`), map[string]*typeInfo{"a": tiHTML()}},
+
+	// key selector
+	{`m.x`, tierr(1, 5, `invalid operation: cannot select m.x (type map[int]string does not support key selection)`), map[string]*typeInfo{"m": intToStringMapTypeInfo}},
+	{`m.x`, tierr(1, 5, `invalid operation: cannot select m.x (type map[int]interface {} does not support key selection)`), map[string]*typeInfo{"m": intToAnyMapTypeInfo}},
+	{`m.x.y`, tierr(1, 7, `m.x.y undefined (type int has no field or method y)`), map[string]*typeInfo{"m": stringToIntMapTypeInfo}},
+	{`m.x.y.z`, tierr(1, 7, `m.x.y undefined (type int has no field or method y)`), map[string]*typeInfo{"m": stringToIntMapTypeInfo}},
+	{`nil.x.y.z`, tierr(1, 4, `use of untyped nil`), nil},
+	{`m._`, tierr(1, 5, `cannot refer to blank field or method`), map[string]*typeInfo{"m": stringToAnyMapTypeInfo}},
 }
 
 func TestCheckerTemplateExpressionErrors(t *testing.T) {
@@ -613,6 +643,23 @@ var checkerTemplateStmts = []struct {
 	{src: `{% L: select %}{% default %}{% break L %}{% end %}`, expected: ok},
 	{src: `{% _ = func() { L: goto L } %}`, expected: ok},
 	{src: `{% L: for %}{% _ = func() { goto L } %}{% end %}`, expected: `label L not defined`},
+
+	// Key selector.
+	{src: `{% m := map[string]int{} %}{% m.x = 5 %}{% m.x += 1 %}{% m.x++ %}{{ m.x + 2 }}`, expected: ok},
+	{src: `{% m := map[DS]int{} %}{% m.x = 5 %}{% m.x += 1 %}{% m.x++ %}{{ m.x + 2 }}`, expected: ok},
+	{src: `{% m := map[string]bool{} %}{% m.nil = true %}{{ m.nil && false }}`, expected: ok},
+	{src: `{% m := map[interface{}]string{} %}{% m.x = "a" %}{{ m.a + "b" }}`, expected: ok},
+	{src: `{% m := map[string]interface{}{} %}{% m.x = 6.89 %}{{ m.x.(float64) - 1.4 }}`, expected: ok},
+	{src: `{% m := map[string]interface{}{} %}{% m.x = map[string]interface{}{} %}{% m.x.y = true %}`, expected: `cannot index m.x (map index expression of type interface{})`},
+	{src: `{% m := map[interface{}]interface{}{} %}{% m.x = map[string]interface{}{} %}{% m.x.y = 'v' %}`, expected: `cannot index m.x (map index expression of type interface{})`},
+	{src: `{% m := map[string]map[string]int{} %}{% m.x = map[string]int{} %}{% m.x.y = 3 %}{% m.x.y += 1 %}{% m.x.y++ %}{{ m.x.y * 2 }}`, expected: ok},
+	{src: `{% m := map[string]map[string]interface{}{} %}{% m.x = map[string]interface{}{} %}{% m.x.y = "a" %}{{ m.x.y.(string) + "b" }}`, expected: ok},
+	{src: `{% m := map[string]int{} %}{% m.x = "a" %}`, expected: `cannot use "a" (type untyped string) as type int in assignment`},
+	{src: `{% m := map[interface{}]int{} %}{% m.x = "a" %}`, expected: `cannot use "a" (type untyped string) as type int in assignment`},
+	{src: `{% m := map[string]int{} %}{% m.x := "a" %}`, expected: `non-name m.x on left side of :=`},
+	{src: `{% m := map[interface{}]int{} %}{% m.x := "a" %}`, expected: `non-name m.x on left side of :=`},
+	{src: `{% m := map[string]int{} %}{{ m.x + "a" }}`, expected: `invalid operation: m.x + "a" (cannot convert "a" (type untyped string) to type int)`},
+	{src: `{% m := map[interface{}]int{} %}{{ m.x + "a" }}`, expected: `invalid operation: m.x + "a" (cannot convert "a" (type untyped string) to type int)`},
 }
 
 func TestCheckerTemplatesStatements(t *testing.T) {
@@ -633,6 +680,7 @@ func TestCheckerTemplatesStatements(t *testing.T) {
 			"Ui": native.UntypedNumericConst("5"),
 			"Uf": native.UntypedNumericConst("5.0"),
 			"R":  'r',
+			"DS": reflect.TypeOf(definedString("")),
 		},
 	}
 	for _, cas := range checkerTemplateStmts {

internal/compiler/disassembler.go

@@ -595,6 +595,10 @@ func disassembleInstruction(fn *runtime.Function, globals []Global, addr runtime
 		s += " " + disassembleOperand(fn, a, reflect.Interface, false)
 		s += " " + disassembleOperand(fn, b, getKind('b', fn, addr), k)
 		s += " " + disassembleOperand(fn, c, getKind('c', fn, addr), false)
+	case runtime.OpMapIndexAny:
+		s += " " + disassembleOperand(fn, a, reflect.Interface, false)
+		s += " " + disassembleOperand(fn, b, reflect.String, k)
+		s += " " + disassembleOperand(fn, c, reflect.Interface, false)
 	case runtime.OpMethodValue:
 		s += " " + disassembleOperand(fn, a, reflect.Interface, false)
 		s += " " + disassembleOperand(fn, b, reflect.String, true)
@@ -1067,7 +1071,8 @@ var operationName = [...]string{
 
 	runtime.OpMakeStruct: "MakeStruct",
 
-	runtime.OpMapIndex: "MapIndex",
+	runtime.OpMapIndex:    "MapIndex",
+	runtime.OpMapIndexAny: "MapIndex",
 
 	runtime.OpMethodValue: "MethodValue",
 

internal/compiler/emitter_assignment.go

@@ -447,6 +447,21 @@ func (em *emitter) assignValuesToAddresses(addresses []address, values []ast.Exp
 		addresses[0].assign(false, value, valueType)
 		addresses[1].assign(false, okReg, okType)
 
+	case *ast.Selector: // key selector.
+		exprType := em.typ(valueExpr.Expr)
+		expr := em.emitExpr(valueExpr.Expr, exprType)
+		key := em.fb.makeStringValue(valueExpr.Ident)
+		value := em.fb.newRegister(reflect.Interface)
+		okType := addresses[1].addressedType
+		okReg := em.fb.newRegister(reflect.Bool)
+		pos := valueExpr.Pos()
+		em.fb.emitIndex(true, expr, key, value, exprType, pos, false)
+		em.fb.emitMove(true, 1, okReg, reflect.Bool)
+		em.fb.emitIf(false, 0, runtime.ConditionOK, 0, reflect.Interface, pos)
+		em.fb.emitMove(true, 0, okReg, reflect.Bool)
+		addresses[0].assign(false, value, emptyInterfaceType)
+		addresses[1].assign(false, okReg, okType)
+
 	case *ast.TypeAssertion:
 		typ := em.typ(valueExpr.Type)
 		expr := em.emitExpr(valueExpr.Expr, emptyInterfaceType)

internal/compiler/emitter_expressions.go

@@ -263,31 +263,7 @@ func (em *emitter) _emitExpr(expr ast.Expression, dstType reflect.Type, reg int8
 
 	case *ast.Index:
 
-		exprType := em.typ(expr.Expr)
-		exprReg := em.emitExpr(expr.Expr, exprType)
-		var indexType reflect.Type
-		if exprType.Kind() == reflect.Map {
-			indexType = exprType.Key()
-		} else {
-			indexType = intType
-		}
-		index, kindex := em.emitExprK(expr.Index, indexType)
-		var elemType reflect.Type
-		if exprType.Kind() == reflect.String {
-			elemType = uint8Type
-		} else {
-			elemType = exprType.Elem()
-		}
-		pos := expr.Pos()
-		if canEmitDirectly(elemType.Kind(), dstType.Kind()) {
-			em.fb.emitIndex(kindex, exprReg, index, reg, exprType, pos, true)
-			return reg, false
-		}
-		em.fb.enterStack()
-		tmp := em.fb.newRegister(elemType.Kind())
-		em.fb.emitIndex(kindex, exprReg, index, tmp, exprType, pos, true)
-		em.changeRegister(false, tmp, reg, elemType, dstType)
-		em.fb.exitStack()
+		em.emitIndex(expr, reg, dstType)
 
 	case *ast.Render:
 
@@ -704,11 +680,59 @@ func (em *emitter) emitCompositeLiteral(expr *ast.CompositeLiteral, reg int8, ds
 	return reg, false
 }
 
+// emitIndex emits an index in register reg.
+func (em *emitter) emitIndex(v *ast.Index, reg int8, dstType reflect.Type) {
+	exprType := em.typ(v.Expr)
+	exprReg := em.emitExpr(v.Expr, exprType)
+	var indexType reflect.Type
+	if exprType.Kind() == reflect.Map {
+		indexType = exprType.Key()
+	} else {
+		indexType = intType
+	}
+	index, kindex := em.emitExprK(v.Index, indexType)
+	var elemType reflect.Type
+	if exprType.Kind() == reflect.String {
+		elemType = uint8Type
+	} else {
+		elemType = exprType.Elem()
+	}
+	pos := v.Pos()
+	if canEmitDirectly(elemType.Kind(), dstType.Kind()) {
+		em.fb.emitIndex(kindex, exprReg, index, reg, exprType, pos, true)
+		return
+	}
+	em.fb.enterStack()
+	tmp := em.fb.newRegister(elemType.Kind())
+	em.fb.emitIndex(kindex, exprReg, index, tmp, exprType, pos, true)
+	em.changeRegister(false, tmp, reg, elemType, dstType)
+	em.fb.exitStack()
+}
+
 // emitSelector emits selector in register reg.
 func (em *emitter) emitSelector(v *ast.Selector, reg int8, dstType reflect.Type) {
 
 	ti := em.ti(v)
 
+	// Key selector.
+	if ti.IsKeySelector() {
+		// Key selector on the empty interface type.
+		if typ := em.typ(v.Expr); typ == emptyInterfaceType {
+			exprReg := em.emitExpr(v.Expr, typ)
+			keyReg := em.fb.makeStringValue(v.Ident)
+			pos := v.Pos()
+			em.fb.enterStack()
+			dst := em.fb.newRegister(reflect.Interface)
+			em.fb.emitIndex(true, exprReg, keyReg, dst, typ, pos, false)
+			em.changeRegister(false, dst, reg, typ, dstType)
+			em.fb.exitStack()
+			return
+		}
+		// Key selector on a map type.
+		em.emitIndex(ti.replacement.(*ast.Index), reg, dstType)
+		return
+	}
+
 	// Method value on concrete and interface values.
 	if ti.MethodType == methodValueConcrete || ti.MethodType == methodValueInterface {
 		expr := v.Expr

internal/compiler/typeinfo.go

@@ -7,6 +7,7 @@ package compiler
 import (
 	"reflect"
 
+	"github.com/open2b/scriggo/ast"
 	"github.com/open2b/scriggo/internal/runtime"
 )
 
@@ -24,6 +25,7 @@ const (
 	propertyHasValue                                              // has a value
 	propertyIsMacroDeclaration                                    // is macro declaration
 	propertyMacroDeclaredInFileWithExtends                        // is macro declared in file with extends
+	propertyKeySelector                                           // is a key selector
 )
 
 // A typeInfo holds the type checking information. For example, every expression
@@ -39,6 +41,7 @@ type typeInfo struct {
 	MethodType        methodType   // Method type.
 	value             interface{}  // value; for packages has type *Package.
 	valueType         reflect.Type // When value is a native type holds the original type of value.
+	replacement       ast.Node     // Replacement node.
 }
 
 // methodType represents the type of a method, intended as a combination of a
@@ -131,6 +134,11 @@ func (ti *typeInfo) MacroDeclaredInExtendingFile() bool {
 	return ti.Properties&propertyMacroDeclaredInFileWithExtends != 0
 }
 
+// IsKeySelector reports whether it is a key selector.
+func (ti *typeInfo) IsKeySelector() bool {
+	return ti.Properties&propertyKeySelector != 0
+}
+
 // TypeName returns the name of the type. If it is an alias, it returns the
 // name of the alias. Panics if it is not a type.
 func (ti *typeInfo) TypeName() string {