com.google.api.tools.framework.model.ExtensionPool Maven / Gradle / Ivy

/*
 * Copyright (C) 2016 Google Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.google.api.tools.framework.model;

import com.google.common.base.Function;
import com.google.common.base.Joiner;
import com.google.common.base.Preconditions;
import com.google.common.base.Strings;
import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.ImmutableListMultimap;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableMultimap;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Multimap;
import com.google.common.collect.Multimaps;
import com.google.common.collect.Ordering;
import com.google.common.collect.Queues;
import com.google.protobuf.DescriptorProtos;
import com.google.protobuf.DescriptorProtos.DescriptorProto;
import com.google.protobuf.DescriptorProtos.FieldDescriptorProto;
import com.google.protobuf.DescriptorProtos.FileDescriptorProto;
import com.google.protobuf.DescriptorProtos.FileDescriptorSet;

import java.util.Deque;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

/**
 * Facilitates easy look up of extensions by name.
 */
public class ExtensionPool {
  public static final ExtensionPool EMPTY =
      new ExtensionPool(null, ImmutableMap.>of());

  private static final Ordering> FIELD_NUMBER_ORDERING =
      new Ordering>() {
        @Override
        public int compare(
            Entry left, Entry right) {
          return left.getValue().getProto().getNumber() - right.getValue().getProto().getNumber();
        }
      };

  private final ImmutableMap> extensions;
  private final FileDescriptorSet descriptor;

  public static final ExtensionPool create(FileDescriptorSet extensionDescriptor) {
    return new Builder().setFileDescriptorSet(extensionDescriptor).build();
  }

  private ExtensionPool(
      FileDescriptorSet descriptor,
      ImmutableMap> extensions) {
    this.descriptor = descriptor;
    this.extensions = extensions;
  }

  /**
   * Returns an {@link Iterable} of extension names to {@link Extension}. Extension names
   * are a fully qualified name surrounded by parentheses. The path of the name is determined by
   * where the extension is defined. For example, given the extension defined below, the field name
   * is "(foo.bar.baz)". The returned {@link Iterable} is sorted by field number.
   *
   *    * syntax = "proto2";
   *
   * package = foo.bar;
   *
   * import "extendee.proto";
   *
   * extend extendee.ExtendeeMessage {
   *   optional baz = 1;
   * }
   * 
   */
  public Iterable> getSortedExtensionsByTypeName(String name) {
    ImmutableMultimap messageExtensions = extensions.get(name);
    if (messageExtensions == null) {
      return Lists.newArrayList();
    }
    return FIELD_NUMBER_ORDERING.immutableSortedCopy(messageExtensions.entries());
  }

  public FileDescriptorSet getDescriptor() {
    return descriptor;
  }

  /**
   * Encapsulates source information for an extension field.
   */
  public static class Extension {
    private final FieldDescriptorProto proto;
    private final DescriptorProtos.SourceCodeInfo.Location location;
    private final String path;
    private final Location fileLocation;

    private Extension(FieldDescriptorProto proto, DescriptorProtos.SourceCodeInfo.Location location,
        String path, Location fileLocation) {
      this.proto = proto;
      this.location = location;
      this.path = path;
      this.fileLocation = fileLocation;
    }

    public FieldDescriptorProto getProto() {
      return proto;
    }

    public DescriptorProtos.SourceCodeInfo.Location getLocation() {
      return location;
    }

    public String getPath() {
      return path;
    }

    public Location getFileLocation() {
      return fileLocation;
    }
  }

  private static class Builder {
    private static final Joiner DOT_JOINER = Joiner.on('.');
    private final Map> builder = Maps.newHashMap();
    private final Deque fullNameSegments = Queues.newArrayDeque();
    private final Deque pathSegments = Queues.newArrayDeque();
    private FileDescriptorSet descriptor;
    private ImmutableListMultimap locationMap;
    private FileDescriptorProto currentFile;

    public ExtensionPool build() {
      ImmutableMap.Builder> builder =
          ImmutableMap.builder();
      for (Entry> entry : this.builder.entrySet()) {
        builder.put(entry.getKey(), ImmutableMultimap.copyOf(entry.getValue()));
      }
      return new ExtensionPool(descriptor, builder.build());
    }

    public Builder setFileDescriptorSet(FileDescriptorSet descriptorSet) {
      Preconditions.checkState(this.descriptor == null, "can only add one FileDescriptorSet");
      this.descriptor = descriptorSet;
      for (FileDescriptorProto fileDescriptor : descriptorSet.getFileList()) {
        add(fileDescriptor);
      }
      return this;
    }

    private void add(FileDescriptorProto file) {
      currentFile = file;
      fullNameSegments.push(file.getPackage());
      locationMap = buildLocationMap(file);
      pathSegments.push(FileDescriptorProto.EXTENSION_FIELD_NUMBER);
      add(file.getExtensionList());
      pathSegments.pop();
      pathSegments.push(FileDescriptorProto.MESSAGE_TYPE_FIELD_NUMBER);
      for (int i = 0; i < file.getMessageTypeCount(); i++) {
        pathSegments.push(i);
        add(file.getMessageType(i));
        pathSegments.pop();
      }
      pathSegments.pop();
      fullNameSegments.pop();
    }

    private void add(DescriptorProto message) {
      fullNameSegments.push(message.getName());
      pathSegments.push(DescriptorProto.EXTENSION_FIELD_NUMBER);
      add(message.getExtensionList());
      pathSegments.pop();
      pathSegments.push(DescriptorProto.NESTED_TYPE_FIELD_NUMBER);
      for (int i = 0; i < message.getNestedTypeCount(); i++) {
        pathSegments.push(i);
        DescriptorProto nested = message.getNestedType(i);
        add(nested);
        pathSegments.pop();
      }
      pathSegments.pop();
      fullNameSegments.pop();
    }

    private void add(List extensions) {
      for (int i = 0; i < extensions.size(); i++) {
        pathSegments.push(i);
        FieldDescriptorProto extensionProto = extensions.get(i);
        String extendee = resolve(extensionProto.getExtendee());
        Multimap messageExtensions = builder.get(extendee);
        if (messageExtensions == null) {
          messageExtensions = ArrayListMultimap.create();
          builder.put(extendee, messageExtensions);
        }
        String path = DOT_JOINER.join(pathSegments.descendingIterator());
        DescriptorProtos.SourceCodeInfo.Location location = locationMap.get(path).get(0);
        // Since paths are only unique within a file, we need a synthetic path to make them unique,
        // given that paths are used to uniquely identify elements in a ProtoFile, and we're
        // stuffing elements from another file into it.
        path = currentFile.getName() + ":" + path;
        Location fileLocation = new SimpleLocation(String.format(
            "%s:%d:%d", currentFile.getName(), location.getSpan(0) + 1, location.getSpan(1) + 1));
        Extension extension = new Extension(extensionProto, location, path, fileLocation);
        messageExtensions.put(getExtensionFieldName(extensionProto.getName()), extension);
        pathSegments.pop();
      }
    }

    private String resolve(String name) {
      if (name.startsWith(".")) {
        return name.substring(1);
      }
      // TODO(user): implement relative extendee naming. Haven't seen it used in protoc output.
      throw new IllegalStateException("ExtensionPool relative name resolution not implemented");
    }

    private String getExtensionFieldName(String shortName) {
      String prefix = DOT_JOINER.join(fullNameSegments.descendingIterator());
      // It's technically possible not to define a package name.
      if (Strings.isNullOrEmpty(prefix)) {
        return String.format("(%s)", shortName);
      }
      return String.format("(%s.%s)", prefix, shortName);
    }

    private static ImmutableListMultimap
        buildLocationMap(FileDescriptorProto file) {
      return Multimaps.index(
          file.getSourceCodeInfo().getLocationList(),
          new Function() {
            @Override
            public String apply(DescriptorProtos.SourceCodeInfo.Location location) {
              return DOT_JOINER.join(location.getPathList());
            }
          });
    }
  }
}