// Copyright 2015 The Bazel Authors. All rights reserved.
//
// 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.devtools.build.lib.packages;

import com.google.common.base.Preconditions;
import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
import com.google.devtools.build.lib.skyframe.serialization.DeserializationContext;
import com.google.devtools.build.lib.skyframe.serialization.ObjectCodec;
import com.google.devtools.build.lib.skyframe.serialization.SerializationContext;
import com.google.devtools.build.lib.skyframe.serialization.SerializationException;
import com.google.protobuf.CodedInputStream;
import com.google.protobuf.CodedOutputStream;
import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * An instance of a given {@code AspectClass} with loaded definition and parameters.
 *
 * This is an aspect equivalent of {@link Rule} class for build rules.
 *
 * Note: this class does not have {@code equals()} and {@code hashCode()} redefined, so should
 * not be used in SkyKeys.
 */
@Immutable
public final class Aspect implements DependencyFilter.AttributeInfoProvider {

  /**
   * The aspect definition is a function of the aspect class + its parameters, so we can cache that.
   *
   * <p>The native aspects are loaded with blaze and are not stateful. Reference equality works fine
   * in this case.
   *
   * <p>Caching of Skylark aspects is not yet implemented.
   */
  private static final Map<NativeAspectClass, Map<AspectParameters, AspectDefinition>>
      definitionCache = new ConcurrentHashMap<>();

  private final AspectDescriptor aspectDescriptor;
  private final AspectDefinition aspectDefinition;

  private Aspect(
      AspectClass aspectClass,
      AspectDefinition aspectDefinition,
      AspectParameters parameters) {
    this.aspectDescriptor = new AspectDescriptor(
        Preconditions.checkNotNull(aspectClass),
        Preconditions.checkNotNull(parameters));
    this.aspectDefinition = Preconditions.checkNotNull(aspectDefinition);
  }

  public static Aspect forNative(
      NativeAspectClass nativeAspectClass, AspectParameters parameters) {
    AspectDefinition definition =
        definitionCache
            .computeIfAbsent(nativeAspectClass, key -> new ConcurrentHashMap<>())
            .computeIfAbsent(parameters, nativeAspectClass::getDefinition);
    return new Aspect(nativeAspectClass, definition, parameters);
  }

  public static Aspect forNative(NativeAspectClass nativeAspectClass) {
    return forNative(nativeAspectClass, AspectParameters.EMPTY);
  }

  public static Aspect forSkylark(
      SkylarkAspectClass skylarkAspectClass,
      AspectDefinition aspectDefinition,
      AspectParameters parameters) {
    return new Aspect(skylarkAspectClass, aspectDefinition, parameters);
  }

  /**
   * Returns the aspectClass required for building the aspect.
   */
  public AspectClass getAspectClass() {
    return aspectDescriptor.getAspectClass();
  }

  /**
   * Returns parameters for evaluation of the aspect.
   */
  public AspectParameters getParameters() {
    return aspectDescriptor.getParameters();
  }

  public AspectDescriptor getDescriptor() {
    return aspectDescriptor;
  }

  @Override
  public String toString() {
    return String.format("Aspect %s", aspectDescriptor.toString());
  }

  public AspectDefinition getDefinition() {
    return aspectDefinition;
  }

  @Override
  public boolean isAttributeValueExplicitlySpecified(Attribute attribute) {
    // All aspect attributes are implicit.
    return false;
  }

  /** {@link ObjectCodec} for {@link Aspect}. */
  static class AspectCodec implements ObjectCodec<Aspect> {
    @Override
    public Class<Aspect> getEncodedClass() {
      return Aspect.class;
    }

    @Override
    public void serialize(SerializationContext context, Aspect obj, CodedOutputStream codedOut)
        throws SerializationException, IOException {
      context.serialize(obj.getDescriptor(), codedOut);
      boolean nativeAspect = obj.getDescriptor().getAspectClass() instanceof NativeAspectClass;
      codedOut.writeBoolNoTag(nativeAspect);
      if (!nativeAspect) {
        context.serialize(obj.getDefinition(), codedOut);
      }
    }

    @Override
    public Aspect deserialize(DeserializationContext context, CodedInputStream codedIn)
        throws SerializationException, IOException {
      AspectDescriptor aspectDescriptor = context.deserialize(codedIn);
      if (codedIn.readBool()) {
        return forNative(
            (NativeAspectClass) aspectDescriptor.getAspectClass(),
            aspectDescriptor.getParameters());
      } else {
        AspectDefinition aspectDefinition = context.deserialize(codedIn);
        return forSkylark(
            (SkylarkAspectClass) aspectDescriptor.getAspectClass(),
            aspectDefinition,
            aspectDescriptor.getParameters());
      }
    }
  }
}

@@ -702,16 +702,150 @@
 itions;%0A
+import com.google.common.cache.CacheBuilder;%0Aimport com.google.common.cache.CacheLoader;%0Aimport com.google.common.cache.LoadingCache;%0A
 import c
@@ -908,16 +908,16 @@
 utable;%0A
-
 import c
@@ -1353,77 +1353,8 @@
 ion;
-%0Aimport java.util.Map;%0Aimport java.util.concurrent.ConcurrentHashMap;
 %0A%0A/*
@@ -2058,20 +2058,40 @@
 c final 
-Map%3C
+LoadingCache%3C%0A          
 NativeAs
@@ -2101,19 +2101,28 @@
 tClass, 
-Map
+LoadingCache
 %3CAspectP
@@ -2178,33 +2178,440 @@
 he =
- new ConcurrentHashMap%3C%3E(
+%0A          CacheBuilder.newBuilder()%0A              .build(%0A                  CacheLoader.from(%0A                      nativeAspectClass -%3E%0A                          CacheBuilder.newBuilder()%0A                              .build(%0A                                  CacheLoader.from(%0A                                      aspectParameters -%3E%0A                                          nativeAspectClass.getDefinition(aspectParameters))))
 );%0A%0A
@@ -3236,164 +3236,64 @@
 ache
-%0A            .computeIfAbsent(nativeAspectClass, key -%3E new ConcurrentHashMap%3C%3E())%0A            .computeIfAbsent(parameters, nativeAspectClass::getDefinition
+.getUnchecked(nativeAspectClass).getUnchecked(parameters
 );%0A 

Use two layers of LoadingCaches, rather than two layers of ConcurrentHashMaps, to implement Aspect#forNative.