feat: support repeat keys in unions

Author: evilmonkeyinc

test/README.md

@@ -268,7 +268,7 @@ This implementation would be closer to the `Scalar consensus` as it does not alw
 |---|---|---|---|---|
 |`$[0,1]`|`["first", "second", "third"]`|`[first second]`|`[first second]`|:white_check_mark:|
 |`$[0,0]`|`["a"]`|`[a a]`|`[a a]`|:white_check_mark:|
-|`$['a','a']`|`{"a":1}`|`[1 1]`|`[1]`|:no_entry:|
+|`$['a','a']`|`{"a":1}`|`[1 1]`|`[1 1]`|:white_check_mark:|
 |`$[?(@.key<3),?(@.key>6)]`|`[{"key": 1}, {"key": 8}, {"key": 3}, {"key": 10}, {"key": 7}, {"key": 2}, {"key": 6}, {"key": 4}]`|none|`[]`|:question:|
 |`$['key','another']`|`{ "key": "value", "another": "entry" }`|`[value entry]`|`[value entry]`|:white_check_mark:|
 |`$['missing','key']`|`{ "key": "value", "another": "entry" }`|`[value]`|`nil`|:no_entry:|

test/union_test.go

@@ -20,9 +20,9 @@ func Test_Union(t *testing.T) {
 			expectedError: "",
 		},
 		{
-			query:         `$['a','a']`, // TODO : need to support repeat keys?
+			query:         `$['a','a']`,
 			data:          `{"a":1}`,
-			expected:      []interface{}{float64(1)},
+			expected:      []interface{}{float64(1), float64(1)},
 			consensus:     []interface{}{float64(1), float64(1)},
 			expectedError: "",
 		},

token/union.go

@@ -141,54 +141,51 @@ func (token *unionToken) getUnionByKey(obj interface{}, keys []string) ([]interf
 		return nil, getInvalidTokenTargetNilError(token.Type(), reflect.Map)
 	}
 
-	keyMap := make(map[string]bool)
-	for _, key := range keys {
-		keyMap[key] = true
-	}
-
 	switch objType.Kind() {
 	case reflect.Map:
 		mapKeys := objVal.MapKeys()
 		sortMapKeys(mapKeys)
 
 		elements := make([]interface{}, 0)
 
+		keysMap := make(map[string]reflect.Value)
 		for _, key := range mapKeys {
-			if keyMap[key.String()] {
-				delete(keyMap, key.String())
+			keysMap[key.String()] = key
+		}
+
+		missingKeys := make([]string, 0)
+
+		for _, requestedKey := range keys {
+			if key, ok := keysMap[requestedKey]; ok {
 				elements = append(elements, objVal.MapIndex(key).Interface())
+			} else {
+				missingKeys = append(missingKeys, requestedKey)
 			}
 		}
 
-		if len(keyMap) > 0 {
-			remaining := make([]string, 0)
-			for key := range keyMap {
-				remaining = append(remaining, key)
-			}
-			sort.Strings(remaining)
-			return nil, getInvalidTokenKeyNotFoundError(token.Type(), strings.Join(remaining, ","))
+		if len(missingKeys) > 0 {
+			sort.Strings(missingKeys)
+			return nil, getInvalidTokenKeyNotFoundError(token.Type(), strings.Join(missingKeys, ","))
 		}
 
 		return elements, nil
 	case reflect.Struct:
 		elements := make([]interface{}, 0)
 
-		mapKeys := getStructFields(objVal, false)
+		keysMap := getStructFields(objVal, false)
+		missingKeys := make([]string, 0)
 
-		for key, field := range mapKeys {
-			if keyMap[key] {
-				delete(keyMap, key)
+		for _, requestedKey := range keys {
+			if field, ok := keysMap[requestedKey]; ok {
 				elements = append(elements, objVal.FieldByName(field.Name).Interface())
+			} else {
+				missingKeys = append(missingKeys, requestedKey)
 			}
 		}
 
-		if len(keyMap) > 0 {
-			remaining := make([]string, 0)
-			for key := range keyMap {
-				remaining = append(remaining, key)
-			}
-			sort.Strings(remaining)
-			return nil, getInvalidTokenKeyNotFoundError(token.Type(), strings.Join(remaining, ","))
+		if len(missingKeys) > 0 {
+			sort.Strings(missingKeys)
+			return nil, getInvalidTokenKeyNotFoundError(token.Type(), strings.Join(missingKeys, ","))
 		}
 
 		return elements, nil

token/union_test.go

@@ -610,6 +610,16 @@ func Test_UnionToken_getUnionByIndex(t *testing.T) {
 				},
 			},
 		},
+		{
+			input: input{
+				token: &unionToken{},
+				obj:   []string{"one", "two", "three"},
+				keys:  []int64{1, 1},
+			},
+			expected: expected{
+				obj: []interface{}{"two", "two"},
+			},
+		},
 	}
 
 	for idx, test := range tests {
@@ -764,6 +774,30 @@ func Test_UnionToken_getUnionByKey(t *testing.T) {
 				err: "union: invalid token key 'gone,missing' not found",
 			},
 		},
+		{
+			input: input{
+				token: &unionToken{},
+				obj: sampleStruct{
+					One: "value",
+				},
+				keys: []string{"one", "one"},
+			},
+			expected: expected{
+				obj: []interface{}{"value", "value"},
+			},
+		},
+		{
+			input: input{
+				token: &unionToken{},
+				obj: map[string]interface{}{
+					"a": "value",
+				},
+				keys: []string{"a", "a", "a"},
+			},
+			expected: expected{
+				obj: []interface{}{"value", "value", "value"},
+			},
+		},
 	}
 
 	for idx, test := range tests {