qapi script: support enum type as discriminator in union

By default, any union will automatically generate a enum type as
"[UnionName]Kind" in C code, and it is duplicated when the discriminator
is specified as a pre-defined enum type in schema. After this patch,
the pre-defined enum type will be really used as the switch case
condition in generated C code, if discriminator is an enum field.

Signed-off-by: Wenchao Xia <wenchaoqemu@gmail.com>
Reviewed-by: Markus Armbruster <armbru@redhat.com>
Signed-off-by: Luiz Capitulino <lcapitulino@redhat.com>

9 files changed, 94 insertions, 17 deletions

diff --git a/docs/qapi-code-gen.txt b/docs/qapi-code-gen.txt
index 0728f36..a2e7921 100644
--- a/docs/qapi-code-gen.txt
+++ b/docs/qapi-code-gen.txt
@@ -123,11 +123,15 @@ And it looks like this on the wire:
 
 Flat union types avoid the nesting on the wire. They are used whenever a
 specific field of the base type is declared as the discriminator ('type' is
-then no longer generated). The discriminator must always be a string field.
+then no longer generated). The discriminator can be a string field or a
+predefined enum field. If it is a string field, a hidden enum type will be
+generated as "[UNION_NAME]Kind". If it is an enum field, a compile time check
+will be done to verify the correctness. It is recommended to use an enum field.
 The above example can then be modified as follows:
 
+ { 'enum': 'BlockdevDriver', 'data': [ 'raw', 'qcow2' ] }
  { 'type': 'BlockdevCommonOptions',
-   'data': { 'driver': 'str', 'readonly': 'bool' } }
+   'data': { 'driver': 'BlockdevDriver', 'readonly': 'bool' } }
  { 'union': 'BlockdevOptions',
    'base': 'BlockdevCommonOptions',
    'discriminator': 'driver',
diff --git a/scripts/qapi-types.py b/scripts/qapi-types.py
index 5885bac..10864ef 100644
--- a/scripts/qapi-types.py
+++ b/scripts/qapi-types.py
@@ -201,14 +201,21 @@ def generate_union(expr):
     base = expr.get('base')
     discriminator = expr.get('discriminator')
 
+    enum_define = discriminator_find_enum_define(expr)
+    if enum_define:
+        discriminator_type_name = enum_define['enum_name']
+    else:
+        discriminator_type_name = '%sKind' % (name)
+
     ret = mcgen('''
 struct %(name)s
 {
-    %(name)sKind kind;
+    %(discriminator_type_name)s kind;
     union {
         void *data;
 ''',
-                name=name)
+                name=name,
+                discriminator_type_name=discriminator_type_name)
 
     for key in typeinfo:
         ret += mcgen('''
@@ -389,8 +396,11 @@ for expr in exprs:
         fdef.write(generate_enum_lookup(expr['enum'], expr['data']))
     elif expr.has_key('union'):
         ret += generate_fwd_struct(expr['union'], expr['data']) + "\n"
-        ret += generate_enum('%sKind' % expr['union'], expr['data'].keys())
-        fdef.write(generate_enum_lookup('%sKind' % expr['union'], expr['data'].keys()))
+        enum_define = discriminator_find_enum_define(expr)
+        if not enum_define:
+            ret += generate_enum('%sKind' % expr['union'], expr['data'].keys())
+            fdef.write(generate_enum_lookup('%sKind' % expr['union'],
+                                            expr['data'].keys()))
         if expr.get('discriminator') == {}:
             fdef.write(generate_anon_union_qtypes(expr))
     else:
diff --git a/scripts/qapi-visit.py b/scripts/qapi-visit.py
index 0baaf60..45ce3a9 100644
--- a/scripts/qapi-visit.py
+++ b/scripts/qapi-visit.py
@@ -259,10 +259,16 @@ def generate_visit_union(expr):
         assert not base
         return generate_visit_anon_union(name, members)
 
-    # There will always be a discriminator in the C switch code, by default it
-    # is an enum type generated silently as "'%sKind' % (name)"
-    ret = generate_visit_enum('%sKind' % name, members.keys())
-    disc_type = '%sKind' % (name)
+    enum_define = discriminator_find_enum_define(expr)
+    if enum_define:
+        # Use the enum type as discriminator
+        ret = ""
+        disc_type = enum_define['enum_name']
+    else:
+        # There will always be a discriminator in the C switch code, by default it
+        # is an enum type generated silently as "'%sKind' % (name)"
+        ret = generate_visit_enum('%sKind' % name, members.keys())
+        disc_type = '%sKind' % (name)
 
     if base:
         base_fields = find_struct(base)['data']
@@ -298,15 +304,16 @@ void visit_type_%(name)s(Visitor *m, %(name)s ** obj, const char *name, Error **
     pop_indent()
 
     if not discriminator:
-        desc_type = "type"
+        disc_key = "type"
     else:
-        desc_type = discriminator
+        disc_key = discriminator
     ret += mcgen('''
-        visit_type_%(name)sKind(m, &(*obj)->kind, "%(type)s", &err);
+        visit_type_%(disc_type)s(m, &(*obj)->kind, "%(disc_key)s", &err);
         if (!err) {
             switch ((*obj)->kind) {
 ''',
-                 name=name, type=desc_type)
+                 disc_type = disc_type,
+                 disc_key = disc_key)
 
     for key in members:
         if not discriminator:
@@ -517,7 +524,11 @@ for expr in exprs:
         ret += generate_visit_list(expr['union'], expr['data'])
         fdef.write(ret)
 
-        ret = generate_decl_enum('%sKind' % expr['union'], expr['data'].keys())
+        enum_define = discriminator_find_enum_define(expr)
+        ret = ""
+        if not enum_define:
+            ret = generate_decl_enum('%sKind' % expr['union'],
+                                     expr['data'].keys())
         ret += generate_declaration(expr['union'], expr['data'])
         fdecl.write(ret)
     elif expr.has_key('enum'):
diff --git a/scripts/qapi.py b/scripts/qapi.py
index eebc8a7..2b43ad2 100644
--- a/scripts/qapi.py
+++ b/scripts/qapi.py
@@ -180,6 +180,25 @@ def find_base_fields(base):
         return None
     return base_struct_define['data']
 
+# Return the discriminator enum define if discriminator is specified as an
+# enum type, otherwise return None.
+def discriminator_find_enum_define(expr):
+    base = expr.get('base')
+    discriminator = expr.get('discriminator')
+
+    if not (discriminator and base):
+        return None
+
+    base_fields = find_base_fields(base)
+    if not base_fields:
+        return None
+
+    discriminator_type = base_fields.get(discriminator)
+    if not discriminator_type:
+        return None
+
+    return find_enum(discriminator_type)
+
 def check_union(expr, expr_info):
     name = expr['union']
     base = expr.get('base')
@@ -254,11 +273,17 @@ def parse_schema(fp):
             add_enum(expr['enum'], expr['data'])
         elif expr.has_key('union'):
             add_union(expr)
-            add_enum('%sKind' % expr['union'])
         elif expr.has_key('type'):
             add_struct(expr)
         exprs.append(expr)
 
+    # Try again for hidden UnionKind enum
+    for expr_elem in schema.exprs:
+        expr = expr_elem['expr']
+        if expr.has_key('union'):
+            if not discriminator_find_enum_define(expr):
+                add_enum('%sKind' % expr['union'])
+
     try:
         check_exprs(schema)
     except QAPIExprError, e:
diff --git a/tests/Makefile b/tests/Makefile
index 6ac9889..88fca31 100644
--- a/tests/Makefile
+++ b/tests/Makefile
@@ -145,7 +145,7 @@ check-qapi-schema-y := $(addprefix tests/qapi-schema/, \
         unclosed-list.json unclosed-object.json unclosed-string.json \
         duplicate-key.json union-invalid-base.json flat-union-no-base.json \
         flat-union-invalid-discriminator.json \
-        flat-union-invalid-branch-key.json)
+        flat-union-invalid-branch-key.json flat-union-reverse-define.json)
 
 GENERATED_HEADERS += tests/test-qapi-types.h tests/test-qapi-visit.h tests/test-qmp-commands.h
 
diff --git a/tests/qapi-schema/flat-union-reverse-define.err b/tests/qapi-schema/flat-union-reverse-define.err
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/tests/qapi-schema/flat-union-reverse-define.err
diff --git a/tests/qapi-schema/flat-union-reverse-define.exit b/tests/qapi-schema/flat-union-reverse-define.exit
new file mode 100644
index 0000000..573541a
--- /dev/null
+++ b/tests/qapi-schema/flat-union-reverse-define.exit
@@ -0,0 +1 @@
+0
diff --git a/tests/qapi-schema/flat-union-reverse-define.json b/tests/qapi-schema/flat-union-reverse-define.json
new file mode 100644
index 0000000..9ea7e72
--- /dev/null
+++ b/tests/qapi-schema/flat-union-reverse-define.json
@@ -0,0 +1,17 @@
+{ 'union': 'TestUnion',
+  'base': 'TestBase',
+  'discriminator': 'enum1',
+  'data': { 'value1': 'TestTypeA',
+            'value2': 'TestTypeB' } }
+
+{ 'type': 'TestBase',
+  'data': { 'enum1': 'TestEnum' } }
+
+{ 'enum': 'TestEnum',
+  'data': [ 'value1', 'value2' ] }
+
+{ 'type': 'TestTypeA',
+  'data': { 'string': 'str' } }
+
+{ 'type': 'TestTypeB',
+  'data': { 'integer': 'int' } }
diff --git a/tests/qapi-schema/flat-union-reverse-define.out b/tests/qapi-schema/flat-union-reverse-define.out
new file mode 100644
index 0000000..03c952e
--- /dev/null
+++ b/tests/qapi-schema/flat-union-reverse-define.out
@@ -0,0 +1,9 @@
+[OrderedDict([('union', 'TestUnion'), ('base', 'TestBase'), ('discriminator', 'enum1'), ('data', OrderedDict([('value1', 'TestTypeA'), ('value2', 'TestTypeB')]))]),
+ OrderedDict([('type', 'TestBase'), ('data', OrderedDict([('enum1', 'TestEnum')]))]),
+ OrderedDict([('enum', 'TestEnum'), ('data', ['value1', 'value2'])]),
+ OrderedDict([('type', 'TestTypeA'), ('data', OrderedDict([('string', 'str')]))]),
+ OrderedDict([('type', 'TestTypeB'), ('data', OrderedDict([('integer', 'int')]))])]
+[{'enum_name': 'TestEnum', 'enum_values': ['value1', 'value2']}]
+[OrderedDict([('type', 'TestBase'), ('data', OrderedDict([('enum1', 'TestEnum')]))]),
+ OrderedDict([('type', 'TestTypeA'), ('data', OrderedDict([('string', 'str')]))]),
+ OrderedDict([('type', 'TestTypeB'), ('data', OrderedDict([('integer', 'int')]))])]