ProtoTypeRegistry.java

package net.morimekta.proto.utils;

import com.google.protobuf.DescriptorProtos;
import com.google.protobuf.Descriptors;
import net.morimekta.collect.UnmodifiableMap;
import net.morimekta.collect.util.Pair;

import java.util.HashMap;
import java.util.Map;

import static net.morimekta.collect.util.Pair.pairOf;

/**
 * A type registry that can look up various types, extensions and values.
 */
public class ProtoTypeRegistry {
    /**
     * @param name Fully qualified type name.
     * @return Message descriptor or null if not found.
     */
    public Descriptors.Descriptor messageTypeByName(String name) {
        return messageTypes.get(name);
    }

    /**
     * @param typeUrl A type URL.
     * @return Message descriptor or null if not found.
     */
    public Descriptors.Descriptor messageTypeByTypeUrl(String typeUrl) {
        return messageTypeByName(getTypeNameFromTypeUrl(typeUrl));
    }

    /**
     * @param name Fully qualified type name.
     * @return Enum type or null if not found.
     */
    public Descriptors.EnumDescriptor enumTypeByName(String name) {
        return enumTypes.get(name);
    }

    /**
     * @param identifier Fully qualified enum type and value.
     * @return The enum value descriptor.
     */
    public Descriptors.EnumValueDescriptor enumValueByQualifiedIdentifier(String identifier) {
        return enumValues.get(identifier);
    }

    /**
     * @param scope Message scope to look for extension for.
     * @param name  The extension qualified name.
     * @return The extension field or null if not found.
     */
    public Descriptors.FieldDescriptor extensionByScopeAndName(Descriptors.Descriptor scope, String name) {
        return extensions.get(pairOf(scope.getFullName(), name));
    }

    /**
     * @param scope  Message scope to look for extension for.
     * @param number The extension (field) number.
     * @return The extension field or null if not found.
     */
    public Descriptors.FieldDescriptor extensionByScopeAndNumber(Descriptors.Descriptor scope, int number) {
        return extensionByNumber.get(pairOf(scope.getFullName(), number));
    }

    /**
     * @param typeUrl A full type URL.
     * @return The qualified type name for the URL.
     */
    public static String getTypeNameFromTypeUrl(String typeUrl) {
        int pos = typeUrl.lastIndexOf('/');
        return pos == -1 ? typeUrl : typeUrl.substring(pos + 1);
    }

    /**
     * @param descriptor A message descriptor.
     * @return The type URL for the type.
     */
    public static String getTypeUrl(Descriptors.Descriptor descriptor) {
        return getTypeUrl("type.googleapis.com", descriptor);
    }

    /**
     * @param typeUrlPrefix The type URL prefix to use.
     * @param descriptor    A message descriptor.
     * @return The type URL for the type.
     */
    public static String getTypeUrl(String typeUrlPrefix, Descriptors.Descriptor descriptor) {
        return typeUrlPrefix.endsWith("/") || typeUrlPrefix.isEmpty()
               ? typeUrlPrefix + descriptor.getFullName()
               : typeUrlPrefix + "/" + descriptor.getFullName();
    }

    /**
     * @return A prot type registry builder.
     */
    public static Builder newBuilder() {
        return new Builder();
    }

    /**
     * The proto type registry builder.
     */
    public static class Builder {
        private final Map<String, Descriptors.Descriptor>                     messageTypes;
        private final Map<String, Descriptors.EnumDescriptor>                 enumTypes;
        private final Map<String, Descriptors.EnumValueDescriptor>            enumValues;
        private final Map<Pair<String, String>, Descriptors.FieldDescriptor>  extensions;
        private final Map<Pair<String, Integer>, Descriptors.FieldDescriptor> extensionByNumber;

        private Builder() {
            messageTypes = new HashMap<>();
            enumTypes = new HashMap<>();
            enumValues = new HashMap<>();
            extensions = new HashMap<>();
            extensionByNumber = new HashMap<>();
        }

        /**
         * Register a message type, and all nested types, to the registry.
         *
         * @param descriptor The message type.
         * @return The registry.
         */
        public Builder register(Descriptors.Descriptor descriptor) {
            if (messageTypes.containsKey(descriptor.getFullName())) {
                return this;
            }
            messageTypes.put(descriptor.getFullName(), descriptor);
            descriptor.getNestedTypes().forEach(this::register);
            descriptor.getEnumTypes().forEach(this::register);
            descriptor.getExtensions().forEach(this::register);
            return this;
        }

        /**
         * Register an extension to the registry.
         *
         * @param descriptor The extension field descriptor.
         * @return The registry.
         */
        public Builder register(Descriptors.FieldDescriptor descriptor) {
            if (!descriptor.isExtension()) {
                return this;
            }
            if (descriptor.getContainingType().getFile().equals(DescriptorProtos.getDescriptor())) {
                return this;
            }
            extensions.put(pairOf(descriptor.getContainingType().getFullName(), descriptor.getFullName()),
                           descriptor);
            extensionByNumber.put(pairOf(descriptor.getContainingType().getFullName(), descriptor.getNumber()),
                                  descriptor);
            return this;
        }

        /**
         * Register a enum type to the registry.
         *
         * @param descriptor The enum type.
         * @return The registry.
         */
        public Builder register(Descriptors.EnumDescriptor descriptor) {
            enumTypes.put(descriptor.getFullName(), descriptor);
            for (var value : descriptor.getValues()) {
                enumValues.put(descriptor.getFullName() + "." + value.getName(), value);
            }
            return this;
        }

        /**
         * Register a file and all containing types to the registry.
         *
         * @param descriptor The file descriptor.
         * @return The registry.
         */
        public Builder register(Descriptors.FileDescriptor descriptor) {
            descriptor.getDependencies().forEach(this::register);
            descriptor.getMessageTypes().forEach(this::register);
            descriptor.getEnumTypes().forEach(this::register);
            descriptor.getExtensions().forEach(this::register);
            return this;
        }

        /**
         * @return The built type registry.
         */
        public ProtoTypeRegistry build() {
            return new ProtoTypeRegistry(messageTypes, enumTypes, enumValues, extensions, extensionByNumber);
        }
    }

    // ---- Private ----

    private final Map<String, Descriptors.Descriptor>                     messageTypes;
    private final Map<String, Descriptors.EnumDescriptor>                 enumTypes;
    private final Map<String, Descriptors.EnumValueDescriptor>            enumValues;
    private final Map<Pair<String, String>, Descriptors.FieldDescriptor>  extensions;
    private final Map<Pair<String, Integer>, Descriptors.FieldDescriptor> extensionByNumber;

    private ProtoTypeRegistry(Map<String, Descriptors.Descriptor> messageTypes,
                              Map<String, Descriptors.EnumDescriptor> enumTypes,
                              Map<String, Descriptors.EnumValueDescriptor> enumValues,
                              Map<Pair<String, String>, Descriptors.FieldDescriptor> extensions,
                              Map<Pair<String, Integer>, Descriptors.FieldDescriptor> extensionByNumber) {
        this.messageTypes = UnmodifiableMap.asMap(messageTypes);
        this.enumTypes = UnmodifiableMap.asMap(enumTypes);
        this.enumValues = UnmodifiableMap.asMap(enumValues);
        this.extensions = UnmodifiableMap.asMap(extensions);
        this.extensionByNumber = UnmodifiableMap.asMap(extensionByNumber);
    }

}