Make TestsInSuiteKey, TestSuiteExpansionKey implement SkyKey

PiperOrigin-RevId: 162715709

1 files changed, 16 insertions, 7 deletions

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;
     }
   }