bazel packages: record load graph (DAG over Modules) in Module and Package

This is a preparatory step for a new, rich query output format that will provide
all the information needed to implement stardoc as a client of blaze, thus allowing
us to delete the fakebuildapi implementations and the starlarkbuildapi abstractions.

In some situations this could increase live heap usage by causing a package
to keep its loaded modules live for longer. Specifically, if packages P and Q
both load module M, then M is changed and package Q is reloaded, the stale
package P remains live, and it holds a reference to the stale M, whereas before
it did not. The next time P is loaded, the stale P and M are evicted.
Note that today, P already keeps alive any modules required by rules instantiated
by P, so an additional cost will be incurred only if P uses functions ("macros")
or data defined in M, but no rules. I expect the effect to be quite marginal.

A number of opportunities for simplification (and memory saving) have been
identified and will be acted on in a follow-up.

Also:
- define BazelModuleContext.of helper function.

RELNOTES: N/A
PiperOrigin-RevId: 320283510
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/PackageFunctionTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/PackageFunctionTest.java
index 8b3a2b2..b8e545a 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/PackageFunctionTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/PackageFunctionTest.java
@@ -34,6 +34,7 @@
 import com.google.devtools.build.lib.cmdline.PackageIdentifier;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.ExtendedEventHandler;
+import com.google.devtools.build.lib.packages.BazelModuleContext;
 import com.google.devtools.build.lib.packages.BuildFileNotFoundException;
 import com.google.devtools.build.lib.packages.ConstantRuleVisibility;
 import com.google.devtools.build.lib.packages.NoSuchPackageException;
@@ -47,6 +48,7 @@
 import com.google.devtools.build.lib.rules.repository.RepositoryDelegatorFunction;
 import com.google.devtools.build.lib.server.FailureDetails.PackageLoading;
 import com.google.devtools.build.lib.skyframe.util.SkyframeExecutorTestUtils;
+import com.google.devtools.build.lib.syntax.Module;
 import com.google.devtools.build.lib.testutil.ManualClock;
 import com.google.devtools.build.lib.testutil.MoreAsserts;
 import com.google.devtools.build.lib.util.DetailedExitCode;
@@ -1232,6 +1234,41 @@
   }
 
   @Test
+  public void testPackageRecordsLoadedModules() throws Exception {
+    scratch.file("p/BUILD", "load('a.bzl', 'a'); load(':b.bzl', 'b')");
+    scratch.file("p/a.bzl", "load('c.bzl', 'c'); a = c");
+    scratch.file("p/b.bzl", "load(':c.bzl', 'c'); b = c");
+    scratch.file("p/c.bzl", "c = 0");
+
+    // load p
+    preparePackageLoading(rootDirectory);
+    SkyKey skyKey = PackageValue.key(PackageIdentifier.parse("@//p"));
+    Package p = validPackageWithoutErrors(skyKey);
+
+    // Keys are load strings as they appear in the source (notice ":" in one of them).
+    Map<String, Module> pLoads = p.getLoads();
+    assertThat(pLoads.keySet().toString()).isEqualTo("[a.bzl, :b.bzl]");
+
+    // subgraph a
+    Module a = pLoads.get("a.bzl");
+    assertThat(a.toString()).isEqualTo("<module //p:a.bzl>");
+    Map<String, Module> aLoads = BazelModuleContext.of(a).loads();
+    assertThat(aLoads.keySet().toString()).isEqualTo("[c.bzl]");
+    Module cViaA = aLoads.get("c.bzl");
+    assertThat(cViaA.toString()).isEqualTo("<module //p:c.bzl>");
+
+    // subgraph b
+    Module b = pLoads.get(":b.bzl");
+    assertThat(b.toString()).isEqualTo("<module //p:b.bzl>");
+    Map<String, Module> bLoads = BazelModuleContext.of(b).loads();
+    assertThat(bLoads.keySet().toString()).isEqualTo("[:c.bzl]");
+    Module cViaB = bLoads.get(":c.bzl");
+    assertThat(cViaB).isSameInstanceAs(cViaA);
+
+    assertThat(cViaA.getGlobal("c")).isEqualTo(0);
+  }
+
+  @Test
   public void veryBrokenPackagePostsDoneToProgressReceiver() throws Exception {
     reporter.removeHandler(failFastHandler);