Method DecodeProtoOp

tensorflow/core/kernels/decode_proto_op.cc:597–718

class DecodeProtoOp : public OpKernel {
 public:
  explicit DecodeProtoOp(OpKernelConstruction* context) : OpKernel(context) {
    string descriptor_source;
    OP_REQUIRES_OK(context,
                   context->GetAttr("descriptor_source", &descriptor_source));

    // We always get back a desc_pool, but we may not own it. If we own it,
    // owned_desc_pool_ will be filled in.
    DescriptorPool const* desc_pool;
    OP_REQUIRES_OK(context, GetDescriptorPool(context->env(), descriptor_source,
                                              &desc_pool, &owned_desc_pool_));

    string message_type;
    OP_REQUIRES_OK(context, context->GetAttr("message_type", &message_type));

    const Descriptor* message_desc =
        desc_pool->FindMessageTypeByName(message_type);
    OP_REQUIRES(context, message_desc != nullptr,
                errors::InvalidArgument("No descriptor found for message type ",
                                        message_type));

    std::vector<string> field_names;
    OP_REQUIRES_OK(context, context->GetAttr("field_names", &field_names));
    std::vector<DataType> output_types;
    OP_REQUIRES_OK(context, context->GetAttr("output_types", &output_types));
    OP_REQUIRES(
        context, field_names.size() == output_types.size(),
        errors::InvalidArgument("field_names and output_types attributes must "
                                "have the same length"));

    // Gather the field descriptors and check that requested output types match.
    int field_index = 0;
    std::vector<const FieldDescriptor*> field_descs;
    std::vector<const FieldDescriptor*> exts;
    absl::flat_hash_map<string, const FieldDescriptor*> ext_name_to_field;
    std::vector<const FieldDescriptor*>::iterator ext_it = exts.begin();
    for (const string& name : field_names) {
      auto fd = message_desc->FindFieldByName(name);
      if (fd == nullptr) {
        // If field can't be found in original message, try to find a matching
        // extension (by its full_name). First check a hashmap for a matching
        // extension, and if not found, then iterate through available
        // extensions to find a match (updating the hashmap while iterating.)
        auto lookup_result = ext_name_to_field.find(name);
        if (lookup_result != ext_name_to_field.end()) {
          fd = lookup_result->second;
        } else {
          if (ext_it == exts.begin()) {
            desc_pool->FindAllExtensions(message_desc, &exts);
            ext_it = exts.begin();
          }
          while (ext_it != exts.end()) {
            auto ext_name = (*ext_it)->full_name();
            auto ext_field = *ext_it;
            ++ext_it;

            ext_name_to_field.insert({ext_name, ext_field});
            if (ext_name == name) {
              fd = ext_field;