001/*
002 * SPDX-License-Identifier: Apache-2.0
003 *
004 * Copyright 2023-2026 The Enola <https://enola.dev> Authors
005 *
006 * Licensed under the Apache License, Version 2.0 (the "License");
007 * you may not use this file except in compliance with the License.
008 * You may obtain a copy of the License at
009 *
010 *     https://www.apache.org/licenses/LICENSE-2.0
011 *
012 * Unless required by applicable law or agreed to in writing, software
013 * distributed under the License is distributed on an "AS IS" BASIS,
014 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015 * See the License for the specific language governing permissions and
016 * limitations under the License.
017 */
018package dev.enola.common.protobuf;
019
020import com.google.common.collect.ImmutableMap;
021import com.google.protobuf.DescriptorProtos.FileDescriptorSet;
022import com.google.protobuf.Descriptors.Descriptor;
023import com.google.protobuf.Descriptors.DescriptorValidationException;
024import com.google.protobuf.Descriptors.EnumDescriptor;
025import com.google.protobuf.Descriptors.FieldDescriptor;
026import com.google.protobuf.Descriptors.FileDescriptor;
027import com.google.protobuf.Descriptors.GenericDescriptor;
028import com.google.protobuf.Descriptors.ServiceDescriptor;
029import com.google.protobuf.TypeRegistry;
030
031import java.util.HashSet;
032import java.util.Set;
033
034/**
035 * TypeRegistryWrapper is a registry of ProtoBuf type descriptors.
036 *
037 * <p>While it's similar to ProtoBuf's own {@link TypeRegistry}, this one (a) has a {@link #names()}
038 * method to enumerate all registered types' names, and (b) includes not just "message" but also
039 * "enum" and "service".
040 */
041// TODO Rename to drop the *Wrapper suffix (it used to wrap TypeRegistry, but now does not anymore)
042// TODO Optimization: This should allow clients like CLI to fetch as Map of Protos!
043public class TypeRegistryWrapper implements DescriptorProvider {
044
045    private final TypeRegistry originalTypeRegistry;
046    private final ImmutableMap<String, GenericDescriptor> types;
047    private final FileDescriptorSet fileDescriptorSet;
048
049    private TypeRegistryWrapper(
050            TypeRegistry typeRegistry,
051            ImmutableMap<String, GenericDescriptor> types,
052            FileDescriptorSet fileDescriptorSet) {
053        this.types = types;
054        this.originalTypeRegistry = typeRegistry;
055        this.fileDescriptorSet = fileDescriptorSet;
056    }
057
058    public static Builder newBuilder() {
059        return new Builder();
060    }
061
062    public static TypeRegistryWrapper from(FileDescriptorSet fileDescriptorSet)
063            throws DescriptorValidationException {
064        var builder = newBuilder();
065        for (var fileDescriptorProto : fileDescriptorSet.getFileList()) {
066            FileDescriptor[] noDependencies = new FileDescriptor[0];
067            FileDescriptor fileDescriptor =
068                    FileDescriptor.buildFrom(fileDescriptorProto, noDependencies, true);
069            builder.add(fileDescriptor.getMessageTypes());
070        }
071        return builder.build();
072    }
073
074    public TypeRegistry get() {
075        return originalTypeRegistry;
076    }
077
078    public FileDescriptorSet fileDescriptorSet() {
079        return fileDescriptorSet;
080    }
081
082    public Set<String> names() {
083        return types.keySet();
084    }
085
086    @Override
087    public GenericDescriptor findByName(String name) {
088        if (name == null) throw new IllegalArgumentException("name == null");
089        if (name.isEmpty()) throw new IllegalArgumentException("name is empty");
090        var descriptor = types.get(name);
091        if (descriptor == null) {
092            throw new IllegalArgumentException(
093                    "Proto unknown: " + name + "; only knows: " + names());
094        }
095        return descriptor;
096    }
097
098    @Override
099    public Descriptor getDescriptorForTypeUrl(String typeURL) {
100        return (Descriptor) findByName(getTypeName(typeURL));
101    }
102
103    // This method is copy/pasted from com.google.protobuf.TypeRegistry
104    private static String getTypeName(String typeUrl) throws IllegalArgumentException {
105        String[] parts = typeUrl.split("/");
106        if (parts.length == 1) {
107            throw new IllegalArgumentException("Invalid type url found: " + typeUrl);
108        }
109        return parts[parts.length - 1];
110    }
111
112    // skipcq: JAVA-E0169
113    public static final class Builder implements dev.enola.common.Builder<TypeRegistryWrapper> {
114        private final Set<String> files = new HashSet<>();
115        private ImmutableMap.Builder<String, GenericDescriptor> typesBuilder =
116                ImmutableMap.builder();
117        private final TypeRegistry.Builder typeRegistryBuilder = TypeRegistry.newBuilder();
118
119        private FileDescriptorSet.Builder fileDescriptorBuilder = FileDescriptorSet.newBuilder();
120
121        private Builder() {}
122
123        public Builder add(Descriptor descriptor) {
124            typeRegistryBuilder.add(descriptor);
125            addFile(descriptor.getFile());
126            return this;
127        }
128
129        public Builder add(Iterable<Descriptor> descriptors) {
130            for (Descriptor descriptor : descriptors) {
131                add(descriptor);
132            }
133            return this;
134        }
135
136        private void addFile(FileDescriptor file) {
137            if (!files.add(file.getFullName())) {
138                return;
139            }
140            for (FileDescriptor dependency : file.getDependencies()) {
141                addFile(dependency);
142            }
143            fileDescriptorBuilder.addFile(file.toProto());
144
145            for (Descriptor messageType : file.getMessageTypes()) {
146                addDescriptor(messageType);
147            }
148            for (var enumType : file.getEnumTypes()) {
149                addDescriptor(enumType);
150            }
151            for (var fieldType : file.getExtensions()) {
152                addDescriptor(fieldType);
153            }
154            for (var serviceType : file.getServices()) {
155                addDescriptor(serviceType);
156            }
157        }
158
159        private void addDescriptor(Descriptor descriptor) {
160            for (var nestedType : descriptor.getNestedTypes()) {
161                addDescriptor(nestedType);
162            }
163            for (var nestedType : descriptor.getEnumTypes()) {
164                addDescriptor(nestedType);
165            }
166            for (var nestedType : descriptor.getExtensions()) {
167                addDescriptor(nestedType);
168            }
169            typesBuilder.put(descriptor.getFullName(), descriptor);
170        }
171
172        private void addDescriptor(EnumDescriptor descriptor) {
173            typesBuilder.put(descriptor.getFullName(), descriptor);
174        }
175
176        private void addDescriptor(FieldDescriptor descriptor) {
177            typesBuilder.put(descriptor.getFullName(), descriptor);
178        }
179
180        private void addDescriptor(ServiceDescriptor descriptor) {
181            typesBuilder.put(descriptor.getFullName(), descriptor);
182        }
183
184        @Override
185        public TypeRegistryWrapper build() {
186            var types = typesBuilder.build();
187            var typeRegistry = typeRegistryBuilder.build();
188            var fileDescriptor = fileDescriptorBuilder.build();
189            var wrapper = new TypeRegistryWrapper(typeRegistry, types, fileDescriptor);
190            typesBuilder = null;
191            fileDescriptorBuilder = null;
192            return wrapper;
193        }
194    }
195}