When inlining imports, avoid visiting Skylark files multiple times that are loaded by a single package. This also fixes a bug in which Skylark code expecting reference equality for an object imported multiple ways were not getting it.

--
MOS_MIGRATED_REVID=135738789

2 files changed, 35 insertions, 19 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/PackageFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/PackageFunction.java
index 733cf5abd3..9db19e5b4e 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/PackageFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/PackageFunction.java
@@ -72,6 +72,7 @@ import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
@@ -687,9 +688,12 @@ public class PackageFunction implements SkyFunction {
         }
       } else {
         // Inlining calls to SkylarkImportLookupFunction
+        LinkedHashMap<Label, SkylarkImportLookupValue> alreadyVisitedImports =
+            Maps.newLinkedHashMapWithExpectedSize(importLookupKeys.size());
         for (SkyKey importLookupKey : importLookupKeys) {
-          SkyValue skyValue = skylarkImportLookupFunctionForInlining.computeWithInlineCalls(
-              importLookupKey, env);
+          SkyValue skyValue =
+              skylarkImportLookupFunctionForInlining.computeWithInlineCalls(
+                  importLookupKey, env, alreadyVisitedImports);
           if (skyValue == null) {
             Preconditions.checkState(
                 env.valuesMissing(), "no skylark import value for %s", importLookupKey);
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SkylarkImportLookupFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/SkylarkImportLookupFunction.java
index 951d8a78b9..8e641325cb 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SkylarkImportLookupFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SkylarkImportLookupFunction.java
@@ -47,11 +47,10 @@ import com.google.devtools.build.skyframe.SkyFunctionException.Transience;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 import com.google.devtools.build.skyframe.ValueOrException2;
-import java.util.LinkedHashSet;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
-import java.util.Set;
 import javax.annotation.Nullable;
 
 /**
@@ -88,17 +87,20 @@ public class SkylarkImportLookupFunction implements SkyFunction {
     }
   }
 
-  SkyValue computeWithInlineCalls(SkyKey skyKey, Environment env)
-      throws InconsistentFilesystemException,
-          SkylarkImportFailedException,
-          InterruptedException {
-    return computeWithInlineCallsInternal(skyKey, env, new LinkedHashSet<Label>());
+  SkyValue computeWithInlineCalls(
+      SkyKey skyKey, Environment env, LinkedHashMap<Label, SkylarkImportLookupValue> visited)
+      throws InconsistentFilesystemException, SkylarkImportFailedException, InterruptedException {
+    return computeWithInlineCallsInternal(skyKey, env, visited);
   }
 
   private SkyValue computeWithInlineCallsInternal(
-      SkyKey skyKey, Environment env, Set<Label> visited)
+      SkyKey skyKey, Environment env, LinkedHashMap<Label, SkylarkImportLookupValue> visited)
       throws InconsistentFilesystemException, SkylarkImportFailedException, InterruptedException {
     SkylarkImportLookupKey key = (SkylarkImportLookupKey) skyKey.argument();
+    SkylarkImportLookupValue precomputedResult = visited.get(key.importLabel);
+    if (precomputedResult != null) {
+      return precomputedResult;
+    }
     return computeInternal(
         key.importLabel,
         key.inWorkspace,
@@ -106,8 +108,11 @@ public class SkylarkImportLookupFunction implements SkyFunction {
         Preconditions.checkNotNull(visited, key.importLabel));
   }
 
-  SkyValue computeInternal(
-      Label fileLabel, boolean inWorkspace, Environment env, @Nullable Set<Label> visited)
+  private SkyValue computeInternal(
+      Label fileLabel,
+      boolean inWorkspace,
+      Environment env,
+      @Nullable LinkedHashMap<Label, SkylarkImportLookupValue> alreadyVisited)
       throws InconsistentFilesystemException, SkylarkImportFailedException, InterruptedException {
     PathFragment filePath = fileLabel.toPathFragment();
 
@@ -153,24 +158,26 @@ public class SkylarkImportLookupFunction implements SkyFunction {
     }
     Map<SkyKey, SkyValue> skylarkImportMap;
     boolean valuesMissing = false;
-    if (visited == null) {
+    if (alreadyVisited == null) {
       // Not inlining.
       skylarkImportMap = env.getValues(importLookupKeys);
       valuesMissing = env.valuesMissing();
     } else {
       // Inlining calls to SkylarkImportLookupFunction.
-      if (!visited.add(fileLabel)) {
+      if (alreadyVisited.containsKey(fileLabel)) {
         ImmutableList<Label> cycle =
-            CycleUtils.splitIntoPathAndChain(Predicates.equalTo(fileLabel), visited)
+            CycleUtils.splitIntoPathAndChain(Predicates.equalTo(fileLabel), alreadyVisited.keySet())
                 .second;
         if (env.getValue(SkylarkImportUniqueCycleFunction.key(cycle)) == null) {
           return null;
         }
         throw new SkylarkImportFailedException("Skylark import cycle");
       }
+      alreadyVisited.put(fileLabel, null);
       skylarkImportMap = Maps.newHashMapWithExpectedSize(imports.size());
       for (SkyKey importLookupKey : importLookupKeys) {
-        SkyValue skyValue = this.computeWithInlineCallsInternal(importLookupKey, env, visited);
+        SkyValue skyValue =
+            this.computeWithInlineCallsInternal(importLookupKey, env, alreadyVisited);
         if (skyValue == null) {
           Preconditions.checkState(
               env.valuesMissing(), "no skylark import value for %s", importLookupKey);
@@ -183,7 +190,7 @@ public class SkylarkImportLookupFunction implements SkyFunction {
         }
       }
       // All imports traversed, this key can no longer be part of a cycle.
-      visited.remove(fileLabel);
+      Preconditions.checkState(alreadyVisited.remove(fileLabel) == null, fileLabel);
     }
     if (valuesMissing) {
       // This means some imports are unavailable.
@@ -204,8 +211,13 @@ public class SkylarkImportLookupFunction implements SkyFunction {
     // Skylark UserDefinedFunction-s in that file will share this function definition Environment,
     // which will be frozen by the time it is returned by createExtension.
     Extension extension = createExtension(ast, fileLabel, extensionsForImports, env, inWorkspace);
-    return new SkylarkImportLookupValue(
-        extension, new SkylarkFileDependency(fileLabel, fileDependencies.build()));
+    SkylarkImportLookupValue result =
+        new SkylarkImportLookupValue(
+            extension, new SkylarkFileDependency(fileLabel, fileDependencies.build()));
+    if (alreadyVisited != null) {
+      alreadyVisited.put(fileLabel, result);
+    }
+    return result;
   }
 
   /**