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 }