// Copyright 2019 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.collect.nestedset;

import static com.google.common.base.Preconditions.checkArgument;

import com.google.common.collect.ImmutableList;
import com.google.common.graph.MutableGraph;
import com.google.devtools.build.lib.actions.ActionInput;
import com.google.devtools.build.lib.actions.Artifact;
import com.google.devtools.build.lib.skyframe.SkyFunctions;
import com.google.devtools.build.skyframe.ExecutionPhaseSkyKey;
import com.google.devtools.build.skyframe.SkyFunctionName;
import com.google.devtools.build.skyframe.SkyKey;
import java.util.HashSet;
import java.util.Set;

/**
 * {@link SkyKey} for requesting all artifacts in a {@link NestedSet}.
 *
 * <p>{@link com.google.devtools.build.lib.skyframe.ArtifactNestedSetFunction} requests the keys
 * returned by {@link #getDirectDepKeys} to ensure that the Skyframe graph mirrors the {@link
 * NestedSet} structure. This class is only declared in the {@code nestedset} package to give it
 * low-level access to the backing {@code Object[]}.
 *
 * <p>An {@link ArtifactNestedSetKey} must only be created from a {@link NestedSet} with multiple
 * elements. Singletons should be translated into a direct request for the corresponding {@link
 * Artifact#key}.
 *
 * <p>Instances are compared using identity equality on the corresponding {@code Object[]} (note
 * that to save memory, this class does not retain a {@link NestedSet.Node}). This means that two
 * instances are not guaranteed to compare equal even if they contain the same set of artifacts.
 */
public final class ArtifactNestedSetKey implements ExecutionPhaseSkyKey {

  private static final SkyKeyInterner<ArtifactNestedSetKey> interner = SkyKey.newInterner();

  public static ArtifactNestedSetKey create(NestedSet<Artifact> set) {
    Object children = set.getChildren();
    checkArgument(
        children instanceof Object[],
        "ArtifactNestedSetKey cannot represent empty or singleton set: %s",
        set);
    return createInternal((Object[]) children);
  }

  private static ArtifactNestedSetKey createInternal(Object[] children) {
    return interner.intern(new ArtifactNestedSetKey(children));
  }

  private final Object[] children;

  private ArtifactNestedSetKey(Object[] children) {
    this.children = children;
  }

  /**
   * Returns a list of this key's direct dependencies, including {@link Artifact#key} for leaves and
   * {@link ArtifactNestedSetKey} for non-leaves.
   */
  public ImmutableList<SkyKey> getDirectDepKeys() {
    ImmutableList.Builder<SkyKey> depKeys = ImmutableList.builderWithExpectedSize(children.length);
    for (Object child : children) {
      if (child instanceof Artifact artifact) {
        depKeys.add(Artifact.key(artifact));
      } else {
        depKeys.add(createInternal((Object[]) child));
      }
    }
    return depKeys.build();
  }

  /** Applies a consumer function to the direct artifacts of this nested set. */
  public void applyToDirectArtifacts(DirectArtifactConsumer function) throws InterruptedException {
    for (Object child : children) {
      if (child instanceof Artifact artifact) {
        function.accept(artifact);
      }
    }
  }

  @Override
  public SkyFunctionName functionName() {
    return SkyFunctions.ARTIFACT_NESTED_SET;
  }

  /**
   * Like {@link NestedSet#toList}, returns a duplicate-free list of the artifacts contained beneath
   * the node represented by this key.
   *
   * <p>Order of the returned list is undefined.
   */
  public ImmutableList<Artifact> expandToArtifacts() {
    // Depth is not accurate, but doesn't matter.
    return NestedSet.<Artifact>create(Order.STABLE_ORDER, /* depth= */ 3, children).toList();
  }

  /**
   * Augments the given rewind graph with the entire nested set structure reachable from {@code
   * key}, including all child {@link ArtifactNestedSetKey} nodes and non-source artifacts.
   *
   * <p>This is used in the imprecise/legacy rewinding case where any lost input in a nested set
   * results in rewinding all artifacts within it. The walk is terminated when a node is already in
   * the rewind graph.
   */
  public static void addEntireNestedSetToRewindGraph(
      MutableGraph<SkyKey> rewindGraph, ArtifactNestedSetKey key) {
    if (rewindGraph.nodes().contains(key)) {
      return;
    }
    for (Object child : key.children) {
      if (child instanceof Artifact artifact) {
        if (!artifact.isSourceArtifact()) {
          rewindGraph.putEdge(key, Artifact.key(artifact));
        }
      } else {
        ArtifactNestedSetKey nextNode = createInternal((Object[]) child);
        addEntireNestedSetToRewindGraph(rewindGraph, nextNode);
        rewindGraph.putEdge(key, nextNode);
      }
    }
  }

  /**
   * Augments the given rewind graph only with the specific paths from {@code failedKey} to {@code
   * lostArtifacts} discoverable by following the {@link ArtifactNestedSetKey} nodes in {@code
   * failedKeyDeps}. This avoids dirtying unrelated inputs in the nested set.
   *
   * <p>If {@code rewindGraph} already contains {@link ArtifactNestedSetKey} nodes prior to calling
   * this method, those nodes are assumed to contain lost artifacts. Paths between {@code failedKey}
   * and such nodes are also added to the rewind graph.
   */
  public static void addNestedSetPathsToRewindGraph(
      MutableGraph<SkyKey> rewindGraph,
      SkyKey failedKey,
      Iterable<? extends SkyKey> failedKeyDeps,
      Set<? extends ActionInput> lostArtifacts) {
    Set<ArtifactNestedSetKey> seen = new HashSet<>();
    for (SkyKey dep : failedKeyDeps) {
      if (dep instanceof ArtifactNestedSetKey nestedSetKey) {
        addNestedSetPathsToRewindGraph(rewindGraph, failedKey, nestedSetKey, lostArtifacts, seen);
      }
    }
  }

  /**
   * Augments the given rewind graph only with the specific paths from {@code failedKey} to {@code
   * lostArtifacts} discoverable by following {@code nestedSetKey}.
   *
   * <p>{@code seen} is used to cache traversals of shared nested set nodes across multiple calls,
   * which is useful when processing multiple root nested sets for propagating actions.
   *
   * <p>If {@code rewindGraph} already contains {@link ArtifactNestedSetKey} nodes prior to calling
   * this method, those nodes are assumed to contain lost artifacts. Paths between {@code failedKey}
   * and such nodes are also added to the rewind graph.
   */
  public static void addNestedSetPathsToRewindGraph(
      MutableGraph<SkyKey> rewindGraph,
      SkyKey failedKey,
      ArtifactNestedSetKey nestedSetKey,
      Set<? extends ActionInput> lostArtifacts,
      Set<ArtifactNestedSetKey> seen) {
    if (searchForLostArtifacts(nestedSetKey, rewindGraph, lostArtifacts, seen)) {
      rewindGraph.putEdge(failedKey, nestedSetKey);
    }
  }

  private static boolean searchForLostArtifacts(
      ArtifactNestedSetKey node,
      MutableGraph<SkyKey> rewindGraph,
      Set<? extends ActionInput> lostArtifacts,
      Set<ArtifactNestedSetKey> seen) {
    if (rewindGraph.nodes().contains(node)) {
      return true;
    }
    if (!seen.add(node)) {
      return false;
    }
    boolean anyFound = false;
    for (Object child : node.children) {
      if (child instanceof Artifact artifact) {
        if (lostArtifacts.contains(artifact)) {
          rewindGraph.putEdge(node, Artifact.key(artifact));
          anyFound = true;
        }
      } else {
        ArtifactNestedSetKey nextNode = createInternal((Object[]) child);
        if (searchForLostArtifacts(nextNode, rewindGraph, lostArtifacts, seen)) {
          rewindGraph.putEdge(node, nextNode);
          anyFound = true;
        }
      }
    }
    return anyFound;
  }

  @Override
  public boolean valueIsShareable() {
    // ArtifactNestedSetValue is just a promise that data is available in memory. Not meant for
    // cross-server sharing.
    return false;
  }

  @Override
  public SkyKeyInterner<?> getSkyKeyInterner() {
    return interner;
  }

  @Override
  public int hashCode() {
    return System.identityHashCode(children);
  }

  @Override
  public boolean equals(Object that) {
    if (this == that) {
      return true;
    }
    return that instanceof ArtifactNestedSetKey artifactNestedSetKey
        && children == artifactNestedSetKey.children;
  }

  @Override
  public String toString() {
    return String.format("ArtifactNestedSetKey[%s]@%s", children.length, hashCode());
  }

  /** A consumer to be applied to each direct artifact. */
  @FunctionalInterface
  public interface DirectArtifactConsumer {
    void accept(Artifact artifact) throws InterruptedException;
  }
}
