SimplePacketManager.java

package net.anxuiz.protobuf.packet.base;

import com.google.common.base.Preconditions;
import com.google.common.collect.BiMap;
import com.google.common.collect.HashBiMap;
import com.google.protobuf.Descriptors.Descriptor;
import com.google.protobuf.Descriptors.FieldDescriptor;
import com.google.protobuf.Descriptors.FieldDescriptor.Type;
import com.google.protobuf.ExtensionRegistry;
import com.google.protobuf.Message;
import com.google.protobuf.Message.Builder;
import net.anxuiz.protobuf.packet.MessageHandlerRegistry;
import net.anxuiz.protobuf.packet.PacketManager;

import javax.annotation.Nonnull;
import java.util.Map;

public final class SimplePacketManager<T extends Message> implements PacketManager<T> {
    private final Builder factory;
    /**
     * Maps {@link Descriptor}s of Messages containing fields extended from this manager's Message to those fields.
     */
    private final BiMap<Descriptor, FieldDescriptor> descriptorMapping = HashBiMap.create();
    /**
     * Registry of all of the extended fields.
     */
    private final ExtensionRegistry reg = ExtensionRegistry.newInstance();
    /**
     * The packet.
     */
    @Nonnull
    private final T packet;

    /**
     * Creates a new manager for specified packet and maps all of its Message's extended fields to their containing
     * Messages.
     *
     * @param packet The packet.
     */
    public SimplePacketManager(@Nonnull final T packet) {
        this.packet = Preconditions.checkNotNull(packet, "packet");

        //  create a builder to make the message mutable
        this.factory = packet.newBuilderForType();

        //  for each extended field of the provided packet's Message...
        for (FieldDescriptor desc : this.factory.getDescriptorForType().getExtensions()) {
            //  check if the extended field is another message
            if (desc.getType().equals(Type.MESSAGE) && !descriptorMapping.containsKey(desc.getExtensionScope())) {
                //  if it is, map it's outer message to itself
                //  this assumes that extensions are always nested within messages
                this.descriptorMapping.put(desc.getExtensionScope(), desc);
                //  and add it to the extension registry so we can parse it later
                this.reg.add(desc);
            }
        }
    }

    /**
     * Creates and handles all of the packet's extended Message fields.
     *
     * @param registry The registry to use to handle the messages.
     * @return The number of Message fields handled.
     */
    public int parse(@Nonnull MessageHandlerRegistry registry) {
        Preconditions.checkNotNull(registry, "handler registry");

        int numParsed = 0;
        for (Map.Entry<FieldDescriptor, Object> entry : packet.getAllFields().entrySet()) {
            //  if we know about this field (extension)
            if (this.descriptorMapping.containsValue(entry.getKey())) {
                //  create and handle the message
                Message msg = (Message) entry.getValue();
                registry.handle(msg);
                numParsed++;
            }
        }

        return numParsed;
    }

    /**
     * Re-builds the packet from the provided Message, setting the extended field back to the Message.
     *
     * @param msg The message.
     * @return The packet.
     */
    @SuppressWarnings("unchecked")
    @Nonnull
    public T build(@Nonnull final Message msg) {
        Preconditions.checkNotNull(msg, "message");

        Builder packet = this.factory.clone();
        FieldDescriptor desc = this.descriptorMapping.get(msg.getDescriptorForType());
        packet.setField(desc, msg);

        return (T) packet.buildPartial();
    }
}