// Copyright 2018 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.skyframe.serialization;

import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.util.concurrent.MoreExecutors.directExecutor;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ClassToInstanceMap;
import com.google.common.collect.ImmutableClassToInstanceMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.devtools.build.lib.skyframe.serialization.Memoizer.Serializer;
import com.google.devtools.build.lib.skyframe.serialization.ObjectCodec.MemoizationStrategy;
import com.google.devtools.build.lib.skyframe.serialization.SerializationException.NoCodecException;
import com.google.errorprone.annotations.CheckReturnValue;
import com.google.protobuf.CodedOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import javax.annotation.Nullable;

/**
 * Stateful class for providing additional context to a single serialization "session". This class
 * is thread-safe so long as {@link #serializer} is null (which also implies that {@link
 * #allowFuturesToBlockWritingOn}) is false). If it is not null, this class is not thread-safe and
 * should only be accessed on a single thread for serializing one object (that may involve
 * serializing other objects contained in it).
 */
public class SerializationContext implements SerializationDependencyProvider {
  private final ObjectCodecRegistry registry;
  private final ImmutableClassToInstanceMap<Object> dependencies;
  @Nullable private final Memoizer.Serializer serializer;
  private final Set<Class<?>> explicitlyAllowedClasses;
  /** Initialized lazily. */
  @Nullable private List<ListenableFuture<Void>> futuresToBlockWritingOn;

  private final boolean allowFuturesToBlockWritingOn;

  private SerializationContext(
      ObjectCodecRegistry registry,
      ImmutableClassToInstanceMap<Object> dependencies,
      @Nullable Serializer serializer,
      boolean allowFuturesToBlockWritingOn) {
    this.registry = registry;
    this.dependencies = dependencies;
    this.serializer = serializer;
    this.allowFuturesToBlockWritingOn = allowFuturesToBlockWritingOn;
    explicitlyAllowedClasses = serializer != null ? new HashSet<>() : ImmutableSet.of();
  }

  @VisibleForTesting
  public SerializationContext(
      ObjectCodecRegistry registry, ImmutableClassToInstanceMap<Object> dependencies) {
    this(registry, dependencies, /*serializer=*/ null, /*allowFuturesToBlockWritingOn=*/ false);
  }

  @VisibleForTesting
  public SerializationContext(ImmutableClassToInstanceMap<Object> dependencies) {
    this(AutoRegistry.get(), dependencies);
  }

  // TODO(shahan): consider making codedOut a member of this class.
  public void serialize(Object object, CodedOutputStream codedOut)
      throws IOException, SerializationException {
    serializeInternal(object, /*customMemoizationStrategy=*/ null, codedOut);
  }

  void serializeWithAdHocMemoizationStrategy(
      Object object, MemoizationStrategy memoizationStrategy, CodedOutputStream codedOut)
      throws IOException, SerializationException {
    serializeInternal(object, memoizationStrategy, codedOut);
  }

  private void serializeInternal(
      Object object,
      @Nullable MemoizationStrategy customMemoizationStrategy,
      CodedOutputStream codedOut)
      throws IOException, SerializationException {
    ObjectCodecRegistry.CodecDescriptor descriptor =
        recordAndGetDescriptorIfNotConstantMemoizedOrNull(object, codedOut);
    if (descriptor != null) {
      if (serializer == null) {
        descriptor.serialize(this, object, codedOut);
      } else {
        @SuppressWarnings("unchecked")
        ObjectCodec<Object> castCodec = (ObjectCodec<Object>) descriptor.getCodec();
        MemoizationStrategy memoizationStrategy =
            customMemoizationStrategy != null ? customMemoizationStrategy : castCodec.getStrategy();
        serializer.serialize(this, object, castCodec, codedOut, memoizationStrategy);
      }
    }
  }

  @Override
  public <T> T getDependency(Class<T> type) {
    return checkNotNull(dependencies.getInstance(type), "Missing dependency of type %s", type);
  }

  /**
   * Returns a {@link SerializationContext} that will memoize values it encounters (using reference
   * equality) in a new memoization table. The returned context should be used instead of the
   * original: memoization may only occur when using the returned context. Calls must be in pairs
   * with {@link DeserializationContext#getMemoizingContext} in the corresponding deserialization
   * code.
   *
   * <p>This method is idempotent: calling it on an already memoizing context will return the same
   * context.
   */
  @CheckReturnValue
  public SerializationContext getMemoizingContext() {
    if (serializer != null) {
      return this;
    }
    return getNewMemoizingContext(/*allowFuturesToBlockWritingOn=*/ false);
  }

  /**
   * Returns a {@link SerializationContext} that will memoize values as described in {@link
   * #getMemoizingContext} and additionally permits attaching futures through {@link
   * #addFutureToBlockWritingOn}.
   */
  @CheckReturnValue
  public SerializationContext getMemoizingAndBlockingOnWriteContext() {
    checkState(serializer == null, "Should only be called on base serializationContext");
    checkState(!allowFuturesToBlockWritingOn, "Should only be called on base serializationContext");
    return getNewMemoizingContext(/*allowFuturesToBlockWritingOn=*/ true);
  }

  /**
   * Returns a memoizing {@link SerializationContext}, as getMemoizingContext above. Unlike
   * getMemoizingContext, this method is not idempotent - the returned context will always be fresh.
   */
  public SerializationContext getNewMemoizingContext() {
    return getNewMemoizingContext(allowFuturesToBlockWritingOn);
  }

  private SerializationContext getNewMemoizingContext(boolean allowFuturesToBlockWritingOn) {
    return new SerializationContext(
        registry, dependencies, new Memoizer.Serializer(), allowFuturesToBlockWritingOn);
  }

  /**
   * Returns a new {@link SerializationContext} mostly identical to this one, but with a dependency
   * map composed by applying overrides to this context's dependencies.
   *
   * <p>The given {@code dependencyOverrides} may contain keys already present (in which case the
   * dependency will be replaced) or new keys (in which case the dependency will be added).
   *
   * <p>Must only be called on a base context (no memoization state), since changing dependencies
   * may change deserialization semantics.
   */
  @CheckReturnValue
  public SerializationContext withDependencyOverrides(ClassToInstanceMap<?> dependencyOverrides) {
    checkState(serializer == null, "Must only be called on base SerializationContext");
    return new SerializationContext(
        registry,
        ImmutableClassToInstanceMap.builder()
            .putAll(Maps.filterKeys(dependencies, k -> !dependencyOverrides.containsKey(k)))
            .putAll(dependencyOverrides)
            .build(),
        /*serializer=*/ null,
        allowFuturesToBlockWritingOn);
  }

  /**
   * Registers a {@link ListenableFuture} that must complete successfully before the serialized
   * bytes generated using this context can be written remotely.
   */
  public void addFutureToBlockWritingOn(ListenableFuture<Void> future) {
    checkState(allowFuturesToBlockWritingOn, "This context cannot block on a future");
    if (futuresToBlockWritingOn == null) {
      futuresToBlockWritingOn = new ArrayList<>();
    }
    futuresToBlockWritingOn.add(future);
  }

  /**
   * Creates a future that succeeds when all futures stored in this context via {@link
   * #addFutureToBlockWritingOn} have succeeded, or null if no such futures were stored.
   */
  @Nullable
  public ListenableFuture<Void> createFutureToBlockWritingOn() {
    return futuresToBlockWritingOn != null
        ? Futures.whenAllSucceed(futuresToBlockWritingOn).call(() -> null, directExecutor())
        : null;
  }

  /**
   * Asserts during serialization that the encoded class of this codec has been explicitly
   * whitelisted for serialization (using {@link #addExplicitlyAllowedClass}). Codecs for objects
   * that are expensive to serialize and that should only be encountered in a limited number of
   * types of {@link com.google.devtools.build.skyframe.SkyValue}s should call this method to check
   * that the object is being serialized as part of an expected {@link
   * com.google.devtools.build.skyframe.SkyValue}, like {@link
   * com.google.devtools.build.lib.packages.Package} inside {@link
   * com.google.devtools.build.lib.skyframe.PackageValue}.
   */
  public <T> void checkClassExplicitlyAllowed(Class<T> allowedClass, T objectForDebugging)
      throws SerializationException {
    if (serializer == null) {
      throw new SerializationException(
          "Cannot check explicitly allowed class "
              + allowedClass
              + " without memoization ("
              + objectForDebugging
              + ")");
    }
    if (!explicitlyAllowedClasses.contains(allowedClass)) {
      throw new SerializationException(
          allowedClass
              + " not explicitly allowed (allowed classes were: "
              + explicitlyAllowedClasses
              + ") and object is "
              + objectForDebugging);
    }
  }

  /**
   * Adds an explicitly allowed class for this serialization context, which must be a memoizing
   * context. Must be called by any codec that transitively serializes an object whose codec calls
   * {@link #checkClassExplicitlyAllowed}.
   *
   * <p>Normally called by codecs for {@link com.google.devtools.build.skyframe.SkyValue} subclasses
   * that know they may encounter an object that is expensive to serialize, like {@link
   * com.google.devtools.build.lib.skyframe.PackageValue} and {@link
   * com.google.devtools.build.lib.packages.Package} or {@link
   * com.google.devtools.build.lib.analysis.ConfiguredTargetValue} and {@link
   * com.google.devtools.build.lib.analysis.configuredtargets.RuleConfiguredTarget}.
   *
   * <p>In case of an unexpected failure from {@link #checkClassExplicitlyAllowed}, it should first
   * be determined if the inclusion of the expensive object is legitimate, before it is whitelisted
   * using this method.
   */
  public void addExplicitlyAllowedClass(Class<?> allowedClass) throws SerializationException {
    if (serializer == null) {
      throw new SerializationException(
          "Cannot add explicitly allowed class %s without memoization: " + allowedClass);
    }
    explicitlyAllowedClasses.add(allowedClass);
  }

  private boolean writeNullOrConstant(@Nullable Object object, CodedOutputStream codedOut)
      throws IOException {
    if (object == null) {
      codedOut.writeSInt32NoTag(0);
      return true;
    }
    Integer tag = registry.maybeGetTagForConstant(object);
    if (tag != null) {
      codedOut.writeSInt32NoTag(tag);
      return true;
    }
    return false;
  }

  @Nullable
  private ObjectCodecRegistry.CodecDescriptor recordAndGetDescriptorIfNotConstantMemoizedOrNull(
      @Nullable Object object, CodedOutputStream codedOut) throws IOException, NoCodecException {
    if (writeNullOrConstant(object, codedOut)) {
      return null;
    }
    if (serializer != null) {
      Integer memoizedIndex = serializer.getMemoizedIndex(object);
      if (memoizedIndex != null) {
        // Subtract 1 so it will be negative and not collide with null.
        codedOut.writeSInt32NoTag(-memoizedIndex - 1);
        return null;
      }
    }
    ObjectCodecRegistry.CodecDescriptor descriptor = registry.getCodecDescriptorForObject(object);
    codedOut.writeSInt32NoTag(descriptor.getTag());
    return descriptor;
  }
}