1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
use super::code_writer::CodeWriter;
use super::rust_types_values::*;
use protobuf::descriptor::*;
use protobuf::descriptorx::*;
use Customize;

struct ExtGen<'a> {
    file: &'a FileDescriptorProto,
    root_scope: &'a RootScope<'a>,
    field: &'a FieldDescriptorProto,
    customize: Customize,
}

impl<'a> ExtGen<'a> {
    fn extendee_rust_name(&self) -> String {
        type_name_to_rust_relative(self.field.get_extendee(), self.file, true, self.root_scope)
    }

    fn repeated(&self) -> bool {
        match self.field.get_label() {
            FieldDescriptorProto_Label::LABEL_REPEATED => true,
            FieldDescriptorProto_Label::LABEL_OPTIONAL => false,
            FieldDescriptorProto_Label::LABEL_REQUIRED => {
                panic!("required ext field: {}", self.field.get_name())
            }
        }
    }

    fn return_type_gen(&self) -> ProtobufTypeGen {
        if self.field.has_type_name() {
            let rust_name_relative = type_name_to_rust_relative(
                self.field.get_type_name(),
                self.file,
                true,
                self.root_scope,
            );
            match self.field.get_field_type() {
                FieldDescriptorProto_Type::TYPE_MESSAGE => {
                    ProtobufTypeGen::Message(rust_name_relative)
                }
                FieldDescriptorProto_Type::TYPE_ENUM => ProtobufTypeGen::Enum(rust_name_relative),
                t => panic!("unknown type: {:?}", t),
            }
        } else {
            ProtobufTypeGen::Primitive(self.field.get_field_type(), PrimitiveTypeVariant::Default)
        }
    }

    fn write(&self, w: &mut CodeWriter) {
        let suffix = if self.repeated() {
            "Repeated"
        } else {
            "Optional"
        };
        let field_type = format!("::protobuf::ext::ExtField{}", suffix);
        w.pub_const(
            &self.field.rust_name(),
            &format!(
                "{}<{}, {}>",
                field_type,
                self.extendee_rust_name(),
                self.return_type_gen().rust_type(&self.customize),
            ),
            &format!(
                "{} {{ field_number: {}, phantom: ::std::marker::PhantomData }}",
                field_type,
                self.field.get_number()
            ),
        );
    }
}

pub(crate) fn write_extensions(
    file: &FileDescriptorProto,
    root_scope: &RootScope,
    w: &mut CodeWriter,
    customize: &Customize,
) {
    if file.get_extension().is_empty() {
        return;
    }

    w.write_line("");
    w.pub_mod("exts", |w| {
        w.write_line("use protobuf::Message as Message_imported_for_functions;");

        for field in file.get_extension() {
            if field.get_field_type() == FieldDescriptorProto_Type::TYPE_GROUP {
                continue;
            }

            w.write_line("");
            ExtGen {
                file: file,
                root_scope: root_scope,
                field: field,
                customize: customize.clone(),
            }
            .write(w);
        }
    });
}