io.protostuff.compiler.parser.ProtoExtensionRegistry Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of protostuff-parser Show documentation

parser for .proto files

There is a newer version: 3.1.40

package io.protostuff.compiler.parser;

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Deque;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

import io.protostuff.compiler.model.Extension;
import io.protostuff.compiler.model.Import;
import io.protostuff.compiler.model.Proto;

/**
 * @author Kostiantyn Shchepanovskyi
 */
public final class ProtoExtensionRegistry extends AbstractExtensionRegistry {

    private final ExtensionRegistry localExtensionRegistry;
    private final ProtoContext context;
    private final ConcurrentMap> extensionCache = new ConcurrentHashMap<>();
    private Proto proto;

    public ProtoExtensionRegistry(ProtoContext context) {
        this.context = context;
        this.proto = context.getProto();
        this.localExtensionRegistry = new LocalExtensionRegistry();
    }

    @Override
    public void registerExtension(Extension extension) {
        super.registerExtension(extension);
        localExtensionRegistry.registerExtension(extension);
        String fullyQualifiedName = extension.getExtendee().getFullyQualifiedName();
        extensionCache.remove(fullyQualifiedName);
    }

    @Override
    public Collection getExtensions(String fullMessageName) {
        return extensionCache.computeIfAbsent(fullMessageName, name -> {
            Collection result = new ArrayList<>();
            result.addAll(localExtensionRegistry.getExtensions(name));
            Deque queue = new ArrayDeque<>();
            queue.addAll(proto.getImports());
            while (!queue.isEmpty()) {
                Import anImport = queue.poll();
                Proto proto = anImport.getProto();
                Collection extensions = getExtensions(proto, name);
                result.addAll(extensions);
                queue.addAll(proto.getPublicImports());
            }
            return result;
        });
    }

    private Collection getExtensions(Proto proto, String name) {
        ProtoContext context = proto.getContext();
        ExtensionRegistry extensionRegistry = context.getExtensionRegistry();
        return extensionRegistry.getExtensions(name);
    }

}