Tracing V2 prototype [NOT FOR REVIEW]

base/trace_event/v2/tracing_protoc.py

+#!/usr/bin/env python
+
+# Hack required to give precedence to the Chromium's protobuf library instead
+# of the debian's package.
+import os
+_CHROME_SRC = os.path.abspath(os.path.join(os.path.abspath(__file__), os.pardir,
+                              os.pardir, os.pardir, os.pardir))
+_PROTOBUF_DIR = os.path.join(_CHROME_SRC, 'third_party', 'protobuf', 'python')
+import pkg_resources
+pkg_resources.declare_namespace('google')
+pkg_resources.fixup_namespace_packages(_PROTOBUF_DIR)
+import google
+assert _PROTOBUF_DIR in google.__path__[0]
+import google.protobuf
+assert _PROTOBUF_DIR in google.protobuf.__path__[0]
+import google.protobuf.descriptor
+from google.protobuf.descriptor import FieldDescriptor
+
+import argparse
+import shutil
+import subprocess
+import sys
+
+
+class CppStubGenerator(object):
+  def __init__(self, proto_file_name):
+    self.proto_file_name = proto_file_name
+    self.processed_classes = set()
+    self.includes = set()
+    self.namespace = None
+    self.cpp_namespace = None
+
+
+  def GetCppClassName(self, full_name):
+    assert(full_name.startswith(self.namespace))
+    return full_name[len(self.namespace) + 1:].replace('.', '_')
+
+
+  def GenerateTranslationUnit(self):
+    db = google.protobuf.symbol_database.Default()
+    descriptors = [x.DESCRIPTOR for x in
+                   db.GetMessages([self.proto_file_name]).itervalues()]
+    namespace_parts = descriptors[0].file.package.split('.')
+    self.namespace = '.'.join(namespace_parts)
+    self.cpp_namespace = '::' + '::'.join(namespace_parts)
+    ifdef_guard = '%s_%s_H' % (self.namespace, self.proto_file_name)
+    ifdef_guard = ifdef_guard.upper().replace('.', '_')
+
+    descriptors_code = ''
+    for descriptor in descriptors:
+      descriptors_code += self.GenerateClassRecursive(descriptor)
+
+    code = '// Autogenerated. DO NOT EDIT.\n'
+    code += '#ifndef %s\n' % ifdef_guard
+    code += '#define %s\n\n' % ifdef_guard
+
+    code += '#include "base/trace_event/v2/append_only_proto_message.h"\n'
+
+    for inc in self.includes:
+      code += '#include "%s"\n' % inc
+    code += '\n'
+
+    for ns in namespace_parts:
+      code += 'namespace %s {\n' % ns
+
+    code += '\n'
+
+    for full_name in self.processed_classes:
+      cpp_class_name = self.GetCppClassName(full_name)
+      if not cpp_class_name:
+        continue
+      code += 'class %s;\n' % cpp_class_name
+    code += descriptors_code
+
+    for ns in namespace_parts:
+      code += '}  // namespace %s\n' % ns
+    code += '#endif  // %s\n' % ifdef_guard
+    code += '\n'
+    return code
+
+
+  def GenerateClassRecursive(self, descriptor):
+    assert(isinstance(descriptor, google.protobuf.descriptor.Descriptor))
+    if descriptor.full_name in self.processed_classes:
+      return ''
+    self.processed_classes.add(descriptor.full_name)
+    code = ''
+    for nested_descriptor in descriptor.nested_types:
+      code += self.GenerateClassRecursive(nested_descriptor)
+    code += self.GenerateClass(descriptor)
+    return code
+
+
+  def GenerateClass(self, descriptor):
+    assert(isinstance(descriptor, google.protobuf.descriptor.Descriptor))
+    assert(descriptor.file.name == self.proto_file_name)
+    assert(descriptor.full_name.startswith(self.namespace))
+    _BASE_CLASS = '::base::trace_event::v2::AppendOnlyProtoMessage'
+    _SIMPLE_TYPES = {
+        FieldDescriptor.TYPE_INT32:  ('int32_t',     'AppendVarIntSigned'),
+        FieldDescriptor.TYPE_INT64:  ('int64_t',     'AppendVarIntSigned'),
+        FieldDescriptor.TYPE_UINT64: ('uint64_t',    'AppendVarIntUnsigned'),
+        FieldDescriptor.TYPE_UINT32: ('uint32_t',    'AppendVarIntUnsigned'),
+        FieldDescriptor.TYPE_BOOL:   ('bool',        'AppendVarIntUnsigned'),
+        FieldDescriptor.TYPE_FLOAT:  ('float',       'AppendFloat'),
+        FieldDescriptor.TYPE_DOUBLE: ('double',      'AppendDouble'),
+        FieldDescriptor.TYPE_STRING: ('const char*', 'AppendString'),
+    }
+
+    class_name = self.GetCppClassName(descriptor.full_name)
+    code = '\n'
+
+    for enum_descriptor in descriptor.enum_types:
+      code += self.GenerateEnum(class_name, enum_descriptor)
+
+    # Generate enum types.
+    code += 'class %s : public %s {\n' % (class_name, _BASE_CLASS)
+    code += ' public:\n'
+    for nested_type in descriptor.nested_types:
+      code += '  using %s = %s;\n' % (
+          nested_type.name, self.GetCppClassName(nested_type.full_name))
+    for enum_descriptor in descriptor.enum_types:
+      code += '  using %(enum_name)s = %(class_name)s_%(enum_name)s;\n' % {
+          'enum_name': enum_descriptor.name, 'class_name': class_name}
+
+    # Generate C++ class.
+    code += '\n'
+    code += '  explicit %(class_name)s() : %(base_class)s() {}\n' % {
+        'class_name': class_name,
+        'base_class': _BASE_CLASS }
+    code += '  ~%s() override {}\n' % class_name
+
+    for field in descriptor.fields:
+      assert not field.has_options or not field.GetOptions().packed, 'Packed fields are not supported'
+
+      camel_name = field.camelcase_name[:1].upper() + field.camelcase_name[1:]
+      code += '  static const int k%sFieldNumber = %d;\n' % (camel_name, field.number)
+
+      if field.type in _SIMPLE_TYPES:
+        cpp_type, marshal_fn = _SIMPLE_TYPES[field.type]
+        code += '  void set_%s(%s value) { %s(%d, value); }\n' % (
+            field.name, cpp_type, marshal_fn, field.number)
+
+      elif field.type == FieldDescriptor.TYPE_BYTES:
+        code += '  void set_%s(const void* value, size_t size) { AppendBytes(%d, value, size); }\n' % (
+            field.name, field.number)
+
+      elif field.type == FieldDescriptor.TYPE_ENUM:
+        cpp_type = '%s::%s' % (self.cpp_namespace, self.GetCppClassName(field.enum_type.full_name))
+        code += '  void set_%s(%s value) { AppendVarIntSigned(%d, static_cast<int32_t>(value)); }\n' % (
+            field.name, cpp_type, field.number)
+
+      elif field.type == FieldDescriptor.TYPE_MESSAGE:
+        if field.message_type.file.name != descriptor.file.name:
+          self.includes.add(ProtoToHppFileName(field.message_type.file.name))
+
+        props = {
+            'base_class': _BASE_CLASS,
+            'full_class_name': '%s::%s' % (
+                self.cpp_namespace,
+                self.GetCppClassName(field.message_type.full_name)),
+            'field_name': field.name,
+            'field_number': field.number,
+        }
+        code += '''
+  %(full_class_name)s* add_%(field_name)s() {
+    return BeginNestedMessage<%(full_class_name)s>(%(field_number)d);
+    // %(full_class_name)s* inst = new %(full_class_name)s();
+    // inst->set_buffer_writer(buffer_writer());
+    // BeginNestedMessage(%(field_number)d, ::std::unique_ptr<%(base_class)s>(inst));
+    // return inst;
+  }\n''' % props
+  #       if field.label != FieldDescriptor.LABEL_REPEATED:
+  #         code += '''
+  # void set_allocated_%(field_name)s(%(full_class_name)s* value) {
+  #   BeginNestedMessage(%(field_number)d, ::std::unique_ptr<%(base_class)s>(value));
+  # }\n''' % props
+      else:
+        assert False, 'Field "%s" not supported' % field.name
+
+    code += '};\n\n'
+    return code
+
+
+  def GenerateEnum(self, class_name, enum_descriptor):
+    assert(isinstance(enum_descriptor, google.protobuf.descriptor.EnumDescriptor))
+    code = 'enum %s_%s {\n' % (class_name, enum_descriptor.name)
+    for enum_value in enum_descriptor.values:
+      code += '  %s_%s_%s = %d,\n' % (class_name, enum_descriptor.name,
+                                         enum_value.name, enum_value.number)
+    code += '};\n\n'
+    return code
+
+def ProtoToHppFileName(proto_file_name):
+  return proto_file_name.replace('.proto', '.tracing-pb.h')
+
+def main():
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--protoc', help='Path to the standard protoc compiler.')
+  parser.add_argument('--proto-out-dir', help='Path for generated stubs.')
+  parser.add_argument('--proto-in-dir', help='Path for source protos.')
+  parser.add_argument('input_protos', nargs='+')
+
+  args = parser.parse_args()
+  protoc = os.path.abspath(args.protoc)
+  assert(os.path.exists(protoc))
+  out_dir = os.path.abspath(args.proto_out_dir)
+  if not os.path.isdir(out_dir):
+    os.makedirs(out_dir)
+
+  print 'Stage 1/2: .proto -> python (using official protoc)'
+  for input_proto_path in args.input_protos:
+    print '  ' + input_proto_path
+    cmd = [protoc, '--proto_path=' + args.proto_in_dir,
+           input_proto_path,
+           '--python_out=' + out_dir,
+           '--cpp_out=' + out_dir,
+           '--js_out=' + out_dir]
+    subprocess.check_call(cmd)
+
+  print '\nStage 2/2: python -> Fast C++ stubs'
+  sys.path.append(out_dir)
+  for input_proto_path in args.input_protos:
+    fname = os.path.basename(input_proto_path)
+    module_name = fname[0:-6] + '_pb2'
+    assert os.path.exists(os.path.join(out_dir, module_name + '.py'))
+    out_cpp_file_path = os.path.join(out_dir, ProtoToHppFileName(fname))
+    print '  %-32s -> %s' % (fname, out_cpp_file_path)
+    __import__(module_name)
+    with open(out_cpp_file_path, 'w') as fout:
+      gen = CppStubGenerator(fname)
+      fout.write(gen.GenerateTranslationUnit())
+
+if __name__ == '__main__':
+  main()