// Copyright 2017 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.analysis.config;

import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.devtools.build.lib.analysis.config.StarlarkDefinedConfigTransition.Settings;
import com.google.devtools.build.lib.analysis.config.transitions.ConfigurationTransition;
import com.google.devtools.build.lib.analysis.config.transitions.TransitionUtil;
import com.google.devtools.build.lib.analysis.starlark.StarlarkBuildSettingsDetailsValue;
import com.google.devtools.build.lib.analysis.starlark.StarlarkBuildSettingsDetailsValue.CustomExecScopeValue;
import com.google.devtools.build.lib.analysis.starlark.StarlarkTransition;
import com.google.devtools.build.lib.analysis.starlark.StarlarkTransition.StarlarkTransitionVisitor;
import com.google.devtools.build.lib.analysis.starlark.StarlarkTransition.TransitionException;
import com.google.devtools.build.lib.cmdline.Label;
import com.google.devtools.build.lib.events.ExtendedEventHandler;
import com.google.devtools.build.lib.events.StoredEventHandler;
import java.util.Map;
import java.util.Objects;
import javax.annotation.Nullable;

/**
 * Caches the application of transitions that use Starlark.
 *
 * <p>This trivially includes {@link StarlarkTransition}s. But it also includes transitions that
 * delegate to {@link StarlarkTransition}s, like some {@link
 * com.google.devtools.build.lib.analysis.config.transitions.ComposingTransitionFactory} instances.
 *
 * <p>This cache was added to keep builds that heavily rely on Starlark transitions performant. The
 * inspiring build is a large Apple binary that heavily relies on {@code objc_library.bzl}, which
 * applies a self-transition. The build applies this transition ~600,000 times. Each application has
 * a cost, mostly from setup in translating Java objects to Starlark objects in {@link
 * com.google.devtools.build.lib.analysis.starlark.FunctionTransitionUtil#applyAndValidate}. This
 * cache saves most of that work, reducing analysis phase CPU time by 17%.
 */
public final class StarlarkTransitionCache {

  private Cache<Key, Value> cache = Caffeine.newBuilder().softValues().build();

  /**
   * Cache of the set of Starlark build settings referenced by a {@link ConfigurationTransition}.
   *
   * <p>This is a separate cache because even if a transition value is evaluated, its Starlark build
   * settings are computed multiple times by {@link TransitionApplier}.
   *
   * <p>Since `--flag_alias` is non-configurable, we can assume that during a single build, the set
   * of starlark build settings for a given transition won't change.
   */
  private Cache<ConfigurationTransition, ImmutableSet<Label>> starlarkBuildSettingsCache =
      Caffeine.newBuilder().softValues().build();

  /**
   * Given a {@link ConfigurationTransition}, decompose (if possible) and find all referenced
   * Starlark build settings.
   *
   * <p>If a transition references a build setting via an alias, this set includes the alias' label
   * and *does not* include the actual label i.e. this method returns all referenced labels exactly
   * as they are.
   *
   * <p>If a flag alias (defined via --flag_alias) is used in the transition, include the starlark
   * flag mapped to this alias.
   */
  public ImmutableSet<Label> getAllStarlarkBuildSettings(
      ConfigurationTransition root, ImmutableMap<String, Label> flagsAliases) {
    var cachedValue = starlarkBuildSettingsCache.getIfPresent(root);
    if (cachedValue != null) {
      return cachedValue;
    }

    ImmutableSet.Builder<Label> keyBuilder = new ImmutableSet.Builder<>();
    try {
      root.visit(
          (StarlarkTransitionVisitor)
              transition ->
                  keyBuilder.addAll(
                      StarlarkTransition.getRelevantStarlarkSettingsFromTransition(
                          transition, flagsAliases, Settings.INPUTS_AND_OUTPUTS)));
    } catch (TransitionException e) {
      // Not actually thrown in the visitor, but declared.
    }

    ImmutableSet<Label> result = keyBuilder.build();
    starlarkBuildSettingsCache.put(root, result);
    return result;
  }

  /** Adds the default values for a transition's input build settings to its input build options. */
  private BuildOptions addDefaultStarlarkOptions(
      BuildOptions fromOptions,
      ImmutableMap<String, Label> flagsAliases,
      ConfigurationTransition transition,
      StarlarkBuildSettingsDetailsValue details) {
    if (details.buildSettingToDefault().isEmpty() && details.customExecScopeValues().isEmpty()) {
      // No need to traverse the transition to find its Starlark flag inputs. There are none.
      return fromOptions;
    }

    BuildOptions.Builder optionsWithDefaults = null;
    for (Label maybeAliasSetting : getAllStarlarkBuildSettings(transition, flagsAliases)) {
      // details will only have the defaults of the actual setting so must unalias
      Label setting = details.aliasToActual().getOrDefault(maybeAliasSetting, maybeAliasSetting);
      if (!fromOptions.getStarlarkOptions().containsKey(maybeAliasSetting)) {
        if (optionsWithDefaults == null) {
          optionsWithDefaults = fromOptions.toBuilder();
        }
        optionsWithDefaults.addStarlarkOption(
            maybeAliasSetting, details.buildSettingToDefault().get(setting));
      }
    }

    return getDefaultStarlarkOptionsForCustomExec(optionsWithDefaults, details, fromOptions);
  }

  public static BuildOptions getDefaultStarlarkOptionsForCustomExec(
      BuildOptions.Builder optionsWithDefaults,
      StarlarkBuildSettingsDetailsValue details,
      BuildOptions fromOptions) {
    for (Label customExecSetting : details.customExecScopeValues().keySet()) {
      if (optionsWithDefaults == null) {
        optionsWithDefaults = fromOptions.toBuilder();
      }

      CustomExecScopeValue customExecScopeValue =
          details.customExecScopeValues().get(customExecSetting);

      if (!fromOptions.getStarlarkOptions().containsKey(customExecScopeValue.hostFlag())) {
        optionsWithDefaults.addStarlarkOption(
            customExecScopeValue.hostFlag(), customExecScopeValue.hostFlagDefault());
        optionsWithDefaults.addScopeType(
            customExecScopeValue.hostFlag(),
            new Scope.ScopeType(customExecScopeValue.hostFlagScopeType()));
      }
      if (!fromOptions.getStarlarkOptions().containsKey(customExecScopeValue.flag())) {
        optionsWithDefaults.addStarlarkOption(
            customExecScopeValue.flag(), customExecScopeValue.flagDefault());
        optionsWithDefaults.addScopeType(
            customExecScopeValue.flag(), new Scope.ScopeType(customExecScopeValue.flagScopeType()));
      }
    }

    return optionsWithDefaults == null
        ? fromOptions
        : optionsWithDefaults.addScopeTypeMap(fromOptions.getScopeTypeMap()).build();
  }

  /**
   * Applies a Starlark transition, possibly returning a cached result.
   *
   * @param fromOptions source options before the transition
   * @param transition the transition itself
   * @param details information from packages about Starlark build settings needed by transition
   * @param eventHandler handler for errors evaluating the transition.
   * @return transition output
   */
  public Map<String, BuildOptions> computeIfAbsent(
      BuildOptions fromOptions,
      ConfigurationTransition transition,
      StarlarkBuildSettingsDetailsValue details,
      ExtendedEventHandler eventHandler)
      throws TransitionException, InterruptedException {
    Key cacheKey = new Key(transition, fromOptions, details);
    Value cachedResult = cache.getIfPresent(cacheKey);
    if (cachedResult != null) {
      if (cachedResult.nonErrorEvents != null) {
        cachedResult.nonErrorEvents.replayOn(eventHandler);
      }
      return cachedResult.result;
    }

    ImmutableMap<String, Label> flagsAliases =
        fromOptions.get(CoreOptions.class).getCommandLineFlagAliasesMap();

    // All code below here only executes on a cache miss and thus should rely only on values that
    // are part of the above cache key or constants that exist throughout the lifetime of the
    // Blaze server instance.
    BuildOptions adjustedOptions =
        addDefaultStarlarkOptions(fromOptions, flagsAliases, transition, details);
    // TODO(bazel-team): Add safety-check that this never mutates fromOptions.
    StoredEventHandler handlerWithErrorStatus = new StoredEventHandler();
    Map<String, BuildOptions> result =
        transition.apply(
            TransitionUtil.restrict(transition, adjustedOptions), handlerWithErrorStatus);

    // We use a temporary StoredEventHandler instead of the caller's event handler because
    // StarlarkTransition.validate assumes no errors occurred. We need a StoredEventHandler to be
    // able to check that, and fail out early if there are errors.
    //
    // TODO(bazel-team): harden StarlarkTransition.validate so we can eliminate this step.
    // StarlarkRuleTransitionProviderTest#testAliasedBuildSetting_outputReturnMismatch shows the
    // effect.
    handlerWithErrorStatus.replayOn(eventHandler);
    if (handlerWithErrorStatus.hasErrors()) {
      throw new TransitionException("Errors encountered while applying Starlark transition");
    }
    result = StarlarkTransition.validate(transition, details, flagsAliases, result);
    // If the transition errored (like bad Starlark code), this method already exited with an
    // exception so the results won't go into the cache. We still want to collect non-error events
    // like print() output.
    StoredEventHandler nonErrorEvents =
        !handlerWithErrorStatus.isEmpty() ? handlerWithErrorStatus : null;
    cache.put(cacheKey, new Value(result, nonErrorEvents));
    return result;
  }

  public void clear() {
    cache = Caffeine.newBuilder().softValues().build();
    starlarkBuildSettingsCache = Caffeine.newBuilder().softValues().build();
  }

  private static final class Key {
    private final ConfigurationTransition transition;
    private final BuildOptions fromOptions;
    private final StarlarkBuildSettingsDetailsValue details;
    private final int hashCode;

    private Key(
        ConfigurationTransition transition,
        BuildOptions fromOptions,
        StarlarkBuildSettingsDetailsValue details) {
      // For rule self-transitions, the transition instance encapsulates both the transition logic
      // and attributes of the target it's attached to. This is important: the same transition in
      // the same configuration applied to distinct targets may produce different outputs. See
      // StarlarkRuleTransitionProvider.FunctionPatchTransition for details.
      this.transition = transition;
      this.fromOptions = fromOptions;
      this.details = details;
      this.hashCode = Objects.hash(transition, fromOptions, details);
    }

    @Override
    public boolean equals(Object other) {
      if (other == this) {
        return true;
      }
      if (!(other instanceof Key otherKey)) {
        return false;
      }
      return this.transition.equals(otherKey.transition)
          && this.fromOptions.equals(otherKey.fromOptions)
          && this.details.equals(otherKey.details);
    }

    @Override
    public int hashCode() {
      return hashCode;
    }
  }

  private static final class Value {
    private final Map<String, BuildOptions> result;

    /**
     * Stores events for successful transitions. Transitions that fail aren't added to the cache.
     * This is meant for non-error events like Starlark {@code print()} output. See {@link
     * com.google.devtools.build.lib.starlark.StarlarkIntegrationTest#testPrintFromTransitionImpl}
     * for a test that covers this.
     *
     * <p>This is null if the transition lacks non-error events.
     */
    @Nullable private final StoredEventHandler nonErrorEvents;

    Value(Map<String, BuildOptions> result, @Nullable StoredEventHandler nonErrorEvents) {
      this.result = result;
      this.nonErrorEvents = nonErrorEvents;
    }
  }
}
