src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionRepoMappingEntriesFunction.java

// Copyright 2023 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.bazel.bzlmod;

import com.google.common.collect.ImmutableMap;
import com.google.devtools.build.lib.cmdline.RepositoryMapping;
import com.google.devtools.build.lib.cmdline.RepositoryName;
import com.google.devtools.build.skyframe.SkyFunction;
import com.google.devtools.build.skyframe.SkyKey;
import com.google.devtools.build.skyframe.SkyValue;
import javax.annotation.Nullable;

/**
 * A SkyFunction that computes the {@link RepositoryMapping#entries()} for the repos generated by a
 * module extension. While the repos differ in {@link RepositoryMapping#ownerRepo()}, the entries
 * are identical and sharing the instance ensures that memory and CPU usage are only linear in the
 * number of repos.
 */
public class ModuleExtensionRepoMappingEntriesFunction implements SkyFunction {

  @Override
  @Nullable
  public SkyValue compute(SkyKey skyKey, Environment env) throws InterruptedException {
    BazelDepGraphValue bazelDepGraphValue =
        (BazelDepGraphValue) env.getValue(BazelDepGraphValue.KEY);
    if (bazelDepGraphValue == null) {
      return null;
    }

    var moduleExtensionId = ((ModuleExtensionRepoMappingEntriesValue.Key) skyKey).argument();
    var extensionEvalValue =
        (SingleExtensionEvalValue) env.getValue(SingleExtensionEvalValue.key(moduleExtensionId));
    if (extensionEvalValue == null) {
      return null;
    }

    return computeRepoMappingEntries(moduleExtensionId, extensionEvalValue, bazelDepGraphValue);
  }

  /**
   * Calculates repo mappings for a repo generated from a module extension. Such a repo can see all
   * repos generated by the same module extension, as well as all repos that the Bazel module
   * hosting the extension can see (see above).
   */
  private ModuleExtensionRepoMappingEntriesValue computeRepoMappingEntries(
      ModuleExtensionId extensionId,
      SingleExtensionEvalValue extensionEvalValue,
      BazelDepGraphValue bazelDepGraphValue) {
    // Find the key of the module containing this extension. This will be used to compute additional
    // mappings -- any repo generated by an extension contained in the module "foo" can additionally
    // see all repos that "foo" can see.
    ModuleKey moduleKey =
        bazelDepGraphValue
            .getCanonicalRepoNameLookup()
            .get(extensionId.getBzlFileLabel().getRepository());
    // NOTE(wyv): This means that if "foo" has a bazel_dep with the repo name "bar", and the
    // extension generates an internal repo name "bar", then within a repo generated by the
    // extension, "bar" will refer to the latter. We should explore a way to differentiate between
    // the two to avoid any surprises.
    ImmutableMap.Builder<String, RepositoryName> entries = ImmutableMap.builder();
    entries.putAll(bazelDepGraphValue.getFullRepoMapping(moduleKey).entries());
    entries.putAll(extensionEvalValue.getCanonicalRepoNameToInternalNames().inverse());
    return ModuleExtensionRepoMappingEntriesValue.create(entries.buildKeepingLast(), moduleKey);
  }
}