// 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 com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ClassToInstanceMap;
import com.google.common.collect.ImmutableClassToInstanceMap;
import com.google.common.collect.Maps;
import com.google.devtools.build.lib.skyframe.serialization.ObjectCodec.MemoizationStrategy;
import com.google.devtools.build.lib.skyframe.serialization.ObjectCodecRegistry.CodecDescriptor;
import com.google.protobuf.CodedInputStream;
import java.io.IOException;
import javax.annotation.CheckReturnValue;
import javax.annotation.Nullable;

/**
 * Stateful class for providing additional context to a single deserialization "session". This class
 * is thread-safe so long as {@link #deserializer} is null. If it is not null, this class is not
 * thread-safe and should only be accessed on a single thread for deserializing one serialized
 * object (that may contain other serialized objects inside it).
 */
public class DeserializationContext implements SerializationDependencyProvider {
  private final ObjectCodecRegistry registry;
  private final ImmutableClassToInstanceMap<Object> dependencies;
  @Nullable private final Memoizer.Deserializer deserializer;

  private DeserializationContext(
      ObjectCodecRegistry registry,
      ImmutableClassToInstanceMap<Object> dependencies,
      @Nullable Memoizer.Deserializer deserializer) {
    this.registry = registry;
    this.dependencies = dependencies;
    this.deserializer = deserializer;
  }

  @VisibleForTesting
  public DeserializationContext(
      ObjectCodecRegistry registry, ImmutableClassToInstanceMap<Object> dependencies) {
    this(registry, dependencies, /*deserializer=*/ null);
  }

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

  // TODO(shahan): consider making codedIn a member of this class.
  @SuppressWarnings({"TypeParameterUnusedInFormals"})
  public <T> T deserialize(CodedInputStream codedIn) throws IOException, SerializationException {
    return deserializeInternal(codedIn, /*customMemoizationStrategy=*/ null);
  }

  @SuppressWarnings({"TypeParameterUnusedInFormals"})
  public <T> T deserializeWithAdHocMemoizationStrategy(
      CodedInputStream codedIn, MemoizationStrategy memoizationStrategy)
      throws IOException, SerializationException {
    return deserializeInternal(codedIn, memoizationStrategy);
  }

  @SuppressWarnings({"TypeParameterUnusedInFormals", "unchecked"})
  private <T> T deserializeInternal(
      CodedInputStream codedIn, @Nullable MemoizationStrategy customMemoizationStrategy)
      throws IOException, SerializationException {
    int tag = codedIn.readSInt32();
    if (tag == 0) {
      return null;
    }
    if (tag < 0) {
      // Subtract 1 to undo transformation from SerializationContext to avoid null.
      return (T) deserializer.getMemoized(-tag - 1); // unchecked cast
    }
    T constant = (T) registry.maybeGetConstantByTag(tag);
    if (constant != null) {
      return constant;
    }
    CodecDescriptor codecDescriptor = registry.getCodecDescriptorByTag(tag);
    if (deserializer == null) {
      return (T) codecDescriptor.deserialize(this, codedIn); // unchecked cast
    } else {
      @SuppressWarnings("unchecked")
      ObjectCodec<T> castCodec = (ObjectCodec<T>) codecDescriptor.getCodec();
      MemoizationStrategy memoizationStrategy =
          customMemoizationStrategy != null ? customMemoizationStrategy : castCodec.getStrategy();
      return deserializer.deserialize(this, castCodec, memoizationStrategy, codedIn);
    }
  }


  /**
   * Register an initial value for the currently deserializing value, for use by child objects that
   * may have references to it.
   *
   * <p>This is a noop when memoization is disabled.
   */
  public <T> void registerInitialValue(T initialValue) {
    if (deserializer != null) {
      deserializer.registerInitialValue(initialValue);
    }
  }

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

  /**
   * Returns a {@link DeserializationContext} that will memoize values it encounters (using
   * reference equality), the inverse of the memoization performed by a {@link SerializationContext}
   * returned by {@link SerializationContext#getMemoizingContext}. The context returned here should
   * be used instead of the original: memoization may only occur when using the returned context.
   *
   * <p>This method is idempotent: calling it on an already memoizing context will return the same
   * context.
   */
  @CheckReturnValue
  public DeserializationContext getMemoizingContext() {
    if (deserializer != null) {
      return this;
    }
    return getNewMemoizingContext();
  }

  /**
   * Returns a new memoizing {@link DeserializationContext}, as {@link #getMemoizingContext}. Unlike
   * {@link #getMemoizingContext}, this method is not idempotent - the returned context will always
   * be fresh.
   */
  public DeserializationContext getNewMemoizingContext() {
    return new DeserializationContext(registry, dependencies, new Memoizer.Deserializer());
  }

  /**
   * Returns a new {@link DeserializationContext} 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 is replaced) or new keys (in which case the dependency is added).
   *
   * <p>Must only be called on a base context (no memoization state), since changing dependencies
   * may change deserialization semantics.
   */
  @CheckReturnValue
  public DeserializationContext withDependencyOverrides(ClassToInstanceMap<?> dependencyOverrides) {
    checkState(deserializer == null, "Must only be called on base DeserializationContext");
    return new DeserializationContext(
        registry,
        ImmutableClassToInstanceMap.builder()
            .putAll(Maps.filterKeys(dependencies, k -> !dependencyOverrides.containsKey(k)))
            .putAll(dependencyOverrides)
            .build(),
        /*deserializer=*/ null);
  }
}