1 package edu.jiangxin.apktoolbox.convert.protobuf.supervised;
2
3 import com.google.protobuf.DescriptorProtos;
4 import com.google.protobuf.Descriptors;
5 import com.google.protobuf.InvalidProtocolBufferException;
6
7 import java.io.IOException;
8 import java.io.UncheckedIOException;
9 import java.nio.file.Files;
10 import java.nio.file.Path;
11 import java.util.*;
12 import java.util.stream.Stream;
13
14 public final class DescriptorCache {
15 private static final Descriptors.FileDescriptor[] DEPENDENCIES = new Descriptors.FileDescriptor[0];
16
17 public static DescriptorCache emptyCache() {
18 return new DescriptorCache();
19 }
20
21 public static DescriptorCache fromDirectory(final Path directory) {
22 Objects.requireNonNull(directory);
23
24 if (!Files.isDirectory(directory)) {
25 throw new IllegalArgumentException("Path must be a directory: " + directory);
26 }
27
28 final DescriptorCache cache = new DescriptorCache();
29 try (Stream<Path> walk = Files.walk(directory)) {
30 walk.filter(Files::isRegularFile)
31 .forEach(cache::addDescriptors);
32 } catch (final IOException e) {
33 throw new UncheckedIOException(e);
34 }
35 return cache;
36 }
37
38 public static DescriptorCache fromFile(final Path descriptorsFile) {
39 Objects.requireNonNull(descriptorsFile);
40
41 if (!Files.isRegularFile(descriptorsFile)) {
42 throw new IllegalArgumentException("Path must be a regular file: " + descriptorsFile);
43 }
44 final DescriptorCache cache = new DescriptorCache();
45 cache.addDescriptors(descriptorsFile);
46 return cache;
47 }
48
49 private final Map<String, Descriptors.Descriptor> typeNameToDescriptor = new HashMap<>();
50
51 private DescriptorCache() {
52
53 }
54
55 @SuppressWarnings({"WeakerAccess", "UnusedReturnValue"})
56 public Optional<Descriptors.Descriptor> addDescriptor(final Descriptors.Descriptor descriptor) {
57 Objects.requireNonNull(descriptor);
58 final String typeName = Objects.requireNonNull(descriptor.getName());
59
60 return Optional.ofNullable(typeNameToDescriptor.put(typeName, descriptor));
61 }
62
63 @SuppressWarnings("WeakerAccess")
64 public void addDescriptors(final Path descriptorsFile) {
65 Objects.requireNonNull(descriptorsFile);
66
67 if (!Files.isRegularFile(descriptorsFile)) {
68 throw new IllegalArgumentException("Path must be a regular file: " + descriptorsFile);
69 }
70
71 try {
72 addDescriptors(Files.readAllBytes(descriptorsFile));
73 } catch (final IOException e) {
74 throw new UncheckedIOException("While reading: " + descriptorsFile.toAbsolutePath(), e);
75 }
76 }
77
78 @SuppressWarnings("WeakerAccess")
79 public void addDescriptors(final byte[] descriptorsRaw) {
80 Objects.requireNonNull(descriptorsRaw);
81
82 try {
83 final DescriptorProtos.FileDescriptorSet descriptorSet =
84 DescriptorProtos.FileDescriptorSet.parseFrom(descriptorsRaw);
85 for (final DescriptorProtos.FileDescriptorProto descriptorFile : descriptorSet.getFileList()) {
86 final Descriptors.FileDescriptor fileDescriptor =
87 Descriptors.FileDescriptor.buildFrom(descriptorFile, DEPENDENCIES);
88 for (final Descriptors.Descriptor descriptor : fileDescriptor.getMessageTypes()) {
89 addDescriptor(descriptor);
90 }
91 }
92 } catch (final InvalidProtocolBufferException | Descriptors.DescriptorValidationException e) {
93 throw new RuntimeException(e);
94 }
95 }
96
97 public Optional<Descriptors.Descriptor> getByTypeName(final String typeName) {
98 Objects.requireNonNull(typeName);
99
100 return Optional.ofNullable(typeNameToDescriptor.get(typeName));
101 }
102
103 public Collection<Descriptors.Descriptor> getDescriptors() {
104 return Collections.unmodifiableCollection(typeNameToDescriptor.values());
105 }
106
107 public Collection<Map.Entry<String, Descriptors.Descriptor>> getEntries() {
108 return Collections.unmodifiableCollection(typeNameToDescriptor.entrySet());
109 }
110
111 public boolean isEmpty() {
112 return typeNameToDescriptor.isEmpty();
113 }
114
115 public int size() {
116 return typeNameToDescriptor.size();
117 }
118 }