Fix invalid generation when using in-rule and not_in rule in map.keys…

….string (#847) ## Summary I fixed invalid generation for Go and cc when using in-rule and not_in rule in map.keys.string ## Issue Closes #785 --------- Co-authored-by: Chris Roche <[email protected]>
bufbuild · Apr 24, 2023 · c87857e · c87857e
1 parent a3c8c99
commit c87857e
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 1 deletion.
diff --git a/templates/cc/msg.go b/templates/cc/msg.go
@@ -91,7 +91,21 @@ const msgTpl = `
 	{{ end }}{{ end }}
         {{ end }}{{ end }}
 
-        {{ if has .Rules "Keys"}}{{ if .Rules.Keys }}
+	{{ if has .Rules "Keys"}}{{ if .Rules.Keys }}
+	{{ if has .Rules.Keys.GetString_ "In" }} {{ if .Rules.Keys.GetString_.In }}
+	const std::set<string> {{ lookup .Field "InLookup" }} = {
+			{{- range .Rules.Keys.GetString_.In }}
+				{{ inKey $f . }},
+			{{- end }}
+		};
+	{{ end }}{{ end }}
+	{{ if has .Rules.Keys.GetString_ "NotIn" }} {{ if .Rules.Keys.GetString_.NotIn }}
+	const std::set<string> {{ lookup .Field "NotInLookup" }} = {
+			{{- range .Rules.Keys.GetString_.NotIn }}
+				{{ inKey $f . }},
+			{{- end }}
+		};
+	{{ end }}{{ end }}
 	{{ if has .Rules.Keys.GetString_ "Pattern" }} {{ if .Rules.Keys.GetString_.Pattern }}
 		const re2::RE2 {{ lookup .Field "Pattern" }}(re2::StringPiece({{ lit .Rules.Keys.GetString_.GetPattern }},
                                               sizeof({{ lit .Rules.Keys.GetString_.GetPattern }}) - 1));

diff --git a/templates/goshared/msg.go b/templates/goshared/msg.go
@@ -234,6 +234,20 @@ var _ interface{
 	{{ end }}{{ end }}
 
 	{{ if has .Rules "Keys"}}{{ if .Rules.Keys }}
+	{{ if has .Rules.Keys.GetString_ "In" }} {{ if .Rules.Keys.GetString_.In }}
+		var {{ lookup .Field "InLookup" }} = map[{{ inType .Field .Rules.Keys.GetString_.In }}]struct{}{
+			{{- range .Rules.Keys.GetString_.In }}
+				{{ inKey $f . }}: {},
+			{{- end }}
+		}
+	{{ end }}{{ end }}
+	{{ if has .Rules.Keys.GetString_ "NotIn" }} {{ if .Rules.Keys.GetString_.NotIn }}
+		var {{ lookup .Field "NotInLookup" }} = map[{{ inType .Field .Rules.Keys.GetString_.NotIn }}]struct{}{
+			{{- range .Rules.Keys.GetString_.NotIn }}
+				{{ inKey $f . }}: {},
+			{{- end }}
+		}
+	{{ end }}{{ end }}
 	{{ if has .Rules.Keys.GetString_ "Pattern" }} {{ if .Rules.Keys.GetString_.Pattern }}
 		var {{ lookup .Field "Pattern" }} = regexp.MustCompile({{ lit .Rules.Keys.GetString_.GetPattern }})
 	{{ end }}{{ end }}

diff --git a/tests/harness/cases/maps.proto b/tests/harness/cases/maps.proto
@@ -36,3 +36,6 @@ message MultipleMaps {
   map <int32, bool> second = 2 [(validate.rules).map.keys.int32.lt = 0];
   map <int32, bool> third = 3 [(validate.rules).map.keys.int32.gt = 0];
 }
+
+message MapKeysIn {map<string, string> val = 1 [(validate.rules).map.keys.string = {in: ["foo", "bar"]}]; }
+message MapKeysNotIn {map<string, string> val = 1 [(validate.rules).map.keys.string = {not_in: ["foo", "bar"]}]; }
diff --git a/tests/harness/executor/cases.go b/tests/harness/executor/cases.go
@@ -1184,8 +1184,12 @@ var mapCases = []TestCase{
 	{"map - keys - valid", &cases.MapKeys{Val: map[int64]string{-1: "a", -2: "b"}}, 0},
 	{"map - keys - valid (empty)", &cases.MapKeys{Val: map[int64]string{}}, 0},
 	{"map - keys - valid (pattern)", &cases.MapKeysPattern{Val: map[string]string{"A": "a"}}, 0},
+	{"map - keys - valid (in)", &cases.MapKeysIn{Val: map[string]string{"foo": "value"}}, 0},
+	{"map - keys - valid (not_in)", &cases.MapKeysNotIn{Val: map[string]string{"baz": "value"}}, 0},
 	{"map - keys - invalid", &cases.MapKeys{Val: map[int64]string{1: "a"}}, 1},
 	{"map - keys - invalid (pattern)", &cases.MapKeysPattern{Val: map[string]string{"A": "a", "!@#$%^&*()": "b"}}, 1},
+	{"map - keys - invalid (in)", &cases.MapKeysIn{Val: map[string]string{"baz": "value"}}, 1},
+	{"map - keys - invalid (not_in)", &cases.MapKeysNotIn{Val: map[string]string{"foo": "value"}}, 1},
 
 	{"map - values - valid", &cases.MapValues{Val: map[string]string{"a": "Alpha", "b": "Beta"}}, 0},
 	{"map - values - valid (empty)", &cases.MapValues{Val: map[string]string{}}, 0},