Make TestsInSuiteKey, TestSuiteExpansionKey implement SkyKey

PiperOrigin-RevId: 162715709

4 files changed, 36 insertions, 19 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/TestSuiteExpansionFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/TestSuiteExpansionFunction.java
index 6324b2e03e..a5a1cc2ae2 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/TestSuiteExpansionFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/TestSuiteExpansionFunction.java
@@ -21,7 +21,7 @@ import com.google.devtools.build.lib.packages.NoSuchTargetException;
 import com.google.devtools.build.lib.packages.Package;
 import com.google.devtools.build.lib.packages.Target;
 import com.google.devtools.build.lib.packages.TargetUtils;
-import com.google.devtools.build.lib.skyframe.TestSuiteExpansionValue.TestSuiteExpansion;
+import com.google.devtools.build.lib.skyframe.TestSuiteExpansionValue.TestSuiteExpansionKey;
 import com.google.devtools.build.skyframe.SkyFunction;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
@@ -40,7 +40,7 @@ import javax.annotation.Nullable;
 final class TestSuiteExpansionFunction implements SkyFunction {
   @Override
   public SkyValue compute(SkyKey key, Environment env) throws InterruptedException {
-    TestSuiteExpansion expansion = (TestSuiteExpansion) key.argument();
+    TestSuiteExpansionKey expansion = (TestSuiteExpansionKey) key.argument();
     ResolvedTargets<Target> targets = labelsToTargets(env, expansion.getTargets(), false);
     List<SkyKey> testsInSuitesKeys = new ArrayList<>();
     for (Target target : targets.getTargets()) {
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/TestSuiteExpansionValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/TestSuiteExpansionValue.java
index 9810b75333..e1d8cb7a8c 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/TestSuiteExpansionValue.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/TestSuiteExpansionValue.java
@@ -22,7 +22,7 @@ import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
 import com.google.devtools.build.lib.packages.Target;
 import com.google.devtools.build.lib.util.Preconditions;
-import com.google.devtools.build.skyframe.LegacySkyKey;
+import com.google.devtools.build.skyframe.SkyFunctionName;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 import java.io.ObjectInputStream;
@@ -70,22 +70,30 @@ public final class TestSuiteExpansionValue implements SkyValue {
    */
   @ThreadSafe
   public static SkyKey key(Collection<Label> targets) {
-    return LegacySkyKey.create(
-        SkyFunctions.TEST_SUITE_EXPANSION,
-        new TestSuiteExpansion(ImmutableSortedSet.copyOf(targets)));
+    return new TestSuiteExpansionKey(ImmutableSortedSet.copyOf(targets));
   }
 
   /**
    * A list of targets of which all test suites should be expanded.
    */
   @ThreadSafe
-  static final class TestSuiteExpansion implements Serializable {
+  static final class TestSuiteExpansionKey implements SkyKey, Serializable {
     private final ImmutableSortedSet<Label> targets;
 
-    public TestSuiteExpansion(ImmutableSortedSet<Label> targets) {
+    public TestSuiteExpansionKey(ImmutableSortedSet<Label> targets) {
       this.targets = targets;
     }
 
+    @Override
+    public SkyFunctionName functionName() {
+      return SkyFunctions.TEST_SUITE_EXPANSION;
+    }
+
+    @Override
+    public Object argument() {
+      return this;
+    }
+
     public ImmutableSet<Label> getTargets() {
       return targets;
     }
@@ -105,10 +113,10 @@ public final class TestSuiteExpansionValue implements SkyValue {
       if (this == obj) {
         return true;
       }
-      if (!(obj instanceof TestSuiteExpansion)) {
+      if (!(obj instanceof TestSuiteExpansionKey)) {
         return false;
       }
-      TestSuiteExpansion other = (TestSuiteExpansion) obj;
+      TestSuiteExpansionKey other = (TestSuiteExpansionKey) obj;
       return other.targets.equals(targets);
     }
   }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/TestsInSuiteFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/TestsInSuiteFunction.java
index 24fc760769..cf1183310e 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/TestsInSuiteFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/TestsInSuiteFunction.java
@@ -28,7 +28,7 @@ import com.google.devtools.build.lib.packages.Rule;
 import com.google.devtools.build.lib.packages.Target;
 import com.google.devtools.build.lib.packages.TargetUtils;
 import com.google.devtools.build.lib.packages.TestTargetUtils;
-import com.google.devtools.build.lib.skyframe.TestsInSuiteValue.TestsInSuite;
+import com.google.devtools.build.lib.skyframe.TestsInSuiteValue.TestsInSuiteKey;
 import com.google.devtools.build.lib.syntax.Type;
 import com.google.devtools.build.lib.util.Pair;
 import com.google.devtools.build.skyframe.SkyFunction;
@@ -53,7 +53,7 @@ import javax.annotation.Nullable;
 final class TestsInSuiteFunction implements SkyFunction {
   @Override
   public SkyValue compute(SkyKey key, Environment env) throws InterruptedException {
-    TestsInSuite expansion = (TestsInSuite) key.argument();
+    TestsInSuiteKey expansion = (TestsInSuiteKey) key.argument();
     ResolvedTargets<Target> result =
         computeTestsInSuite(env, expansion.getTestSuite(), expansion.isStrict());
     if (env.valuesMissing()) {
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/TestsInSuiteValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/TestsInSuiteValue.java
index d0ac8a4ac0..3dabf0682f 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/TestsInSuiteValue.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/TestsInSuiteValue.java
@@ -20,7 +20,7 @@ import com.google.devtools.build.lib.packages.Rule;
 import com.google.devtools.build.lib.packages.Target;
 import com.google.devtools.build.lib.packages.TargetUtils;
 import com.google.devtools.build.lib.util.Preconditions;
-import com.google.devtools.build.skyframe.LegacySkyKey;
+import com.google.devtools.build.skyframe.SkyFunctionName;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 import java.io.ObjectInputStream;
@@ -68,23 +68,32 @@ final class TestsInSuiteValue implements SkyValue {
   @ThreadSafe
   public static SkyKey key(Target testSuite, boolean strict) {
     Preconditions.checkState(TargetUtils.isTestSuiteRule(testSuite));
-    return LegacySkyKey.create(
-        SkyFunctions.TESTS_IN_SUITE, new TestsInSuite((Rule) testSuite, strict));
+    return new TestsInSuiteKey((Rule) testSuite, strict);
   }
 
   /**
    * A list of targets of which all test suites should be expanded.
    */
   @ThreadSafe
-  static final class TestsInSuite implements Serializable {
+  static final class TestsInSuiteKey implements SkyKey, Serializable {
     private final Rule testSuite;
     private final boolean strict;
 
-    public TestsInSuite(Rule testSuite, boolean strict) {
+    public TestsInSuiteKey(Rule testSuite, boolean strict) {
       this.testSuite = testSuite;
       this.strict = strict;
     }
 
+    @Override
+    public SkyFunctionName functionName() {
+      return SkyFunctions.TESTS_IN_SUITE;
+    }
+
+    @Override
+    public Object argument() {
+      return this;
+    }
+
     public Rule getTestSuite() {
       return testSuite;
     }
@@ -108,10 +117,10 @@ final class TestsInSuiteValue implements SkyValue {
       if (this == obj) {
         return true;
       }
-      if (!(obj instanceof TestsInSuite)) {
+      if (!(obj instanceof TestsInSuiteKey)) {
         return false;
       }
-      TestsInSuite other = (TestsInSuite) obj;
+      TestsInSuiteKey other = (TestsInSuiteKey) obj;
       return other.testSuite.equals(testSuite) && other.strict == strict;
     }
   }