[8.1.0] Add set data type to Starlark

src/main/java/com/google/devtools/build/lib/packages/semantics/BuildLanguageOptions.java

@@ -803,6 +803,15 @@ public final class BuildLanguageOptions extends OptionsBase {
               + " attributes of symbolic macros or attribute default values.")
   public boolean incompatibleSimplifyUnconditionalSelectsInRuleAttrs;
 
+  @Option(
+      name = "experimental_enable_starlark_set",
+      defaultValue = "true",
+      documentationCategory = OptionDocumentationCategory.STARLARK_SEMANTICS,
+      effectTags = {OptionEffectTag.BUILD_FILE_SEMANTICS},
+      metadataTags = {OptionMetadataTag.EXPERIMENTAL},
+      help = "If true, enable the set data type and set() constructor in Starlark.")
+  public boolean experimentalEnableStarlarkSet;
+
   @Option(
       name = "incompatible_locations_prefers_executable",
       defaultValue = "true",
@@ -929,6 +938,8 @@ public StarlarkSemantics toStarlarkSemantics() {
                 incompatibleSimplifyUnconditionalSelectsInRuleAttrs)
             .setBool(
                 INCOMPATIBLE_LOCATIONS_PREFERS_EXECUTABLE, incompatibleLocationsPrefersExecutable)
+            .setBool(
+                StarlarkSemantics.EXPERIMENTAL_ENABLE_STARLARK_SET, experimentalEnableStarlarkSet)
             .build();
     return INTERNER.intern(semantics);
   }

src/main/java/net/starlark/java/eval/BUILD

@@ -51,7 +51,9 @@ java_library(
         "StarlarkInt.java",
         "StarlarkIterable.java",
         "StarlarkList.java",
+        "StarlarkMembershipTestable.java",
         "StarlarkSemantics.java",
+        "StarlarkSet.java",
         "StarlarkThread.java",
         "StarlarkValue.java",
         "StringModule.java",

src/main/java/net/starlark/java/eval/Eval.java

@@ -22,6 +22,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import net.starlark.java.spelling.SpellChecker;
 import net.starlark.java.syntax.Argument;
 import net.starlark.java.syntax.AssignmentStatement;
@@ -469,20 +470,66 @@ private static void execAugmentedAssignment(StarlarkThread.Frame fr, AssignmentS
 
   private static Object inplaceBinaryOp(StarlarkThread.Frame fr, TokenKind op, Object x, Object y)
       throws EvalException {
-    // list += iterable  behaves like  list.extend(iterable)
-    // TODO(b/141263526): following Python, allow list+=iterable (but not list+iterable).
-    if (op == TokenKind.PLUS && x instanceof StarlarkList && y instanceof StarlarkList) {
-      StarlarkList<?> list = (StarlarkList) x;
-      list.extend(y);
-      return list;
-    } else if (op == TokenKind.PIPE && x instanceof Dict && y instanceof Map) {
-      // dict |= map merges the contents of the second operand (usually a dict) into the first.
-      @SuppressWarnings("unchecked")
-      Dict<Object, Object> xDict = (Dict<Object, Object>) x;
-      @SuppressWarnings("unchecked")
-      Map<Object, Object> yMap = (Map<Object, Object>) y;
-      xDict.putEntries(yMap);
-      return xDict;
+    switch (op) {
+      case PLUS:
+        // list += iterable  behaves like  list.extend(iterable)
+        // TODO(b/141263526): following Python, allow list+=iterable (but not list+iterable).
+        if (x instanceof StarlarkList && y instanceof StarlarkList) {
+          StarlarkList<?> list = (StarlarkList) x;
+          list.extend(y);
+          return list;
+        }
+        break;
+
+      case PIPE:
+        if (x instanceof Dict && y instanceof Map) {
+          // dict |= map merges the contents of the second operand (usually a dict) into the first.
+          @SuppressWarnings("unchecked")
+          Dict<Object, Object> xDict = (Dict<Object, Object>) x;
+          @SuppressWarnings("unchecked")
+          Map<Object, Object> yMap = (Map<Object, Object>) y;
+          xDict.putEntries(yMap);
+          return xDict;
+        } else if (x instanceof StarlarkSet && y instanceof Set) {
+          // set |= set merges the contents of the second operand into the first.
+          @SuppressWarnings("unchecked")
+          StarlarkSet<Object> xSet = (StarlarkSet<Object>) x;
+          xSet.update(Tuple.of(y));
+          return xSet;
+        }
+        break;
+
+      case AMPERSAND:
+        if (x instanceof StarlarkSet && y instanceof Set) {
+          // set &= set replaces the first set with the intersection of the two sets.
+          @SuppressWarnings("unchecked")
+          StarlarkSet<Object> xSet = (StarlarkSet<Object>) x;
+          xSet.intersectionUpdate(Tuple.of(y));
+          return xSet;
+        }
+        break;
+
+      case CARET:
+        if (x instanceof StarlarkSet && y instanceof Set) {
+          // set ^= set replaces the first set with the symmetric difference of the two sets.
+          @SuppressWarnings("unchecked")
+          StarlarkSet<Object> xSet = (StarlarkSet<Object>) x;
+          xSet.symmetricDifferenceUpdate(y);
+          return xSet;
+        }
+        break;
+
+      case MINUS:
+        if (x instanceof StarlarkSet && y instanceof Set) {
+          // set -= set removes all elements of the second set from the first set.
+          @SuppressWarnings("unchecked")
+          StarlarkSet<Object> xSet = (StarlarkSet<Object>) x;
+          xSet.differenceUpdate(Tuple.of(y));
+          return xSet;
+        }
+        break;
+
+      default: // fall through
     }
     return EvalUtils.binaryOp(op, x, y, fr.thread);
   }

src/main/java/net/starlark/java/eval/EvalUtils.java

@@ -15,6 +15,7 @@
 
 import java.util.IllegalFormatException;
 import java.util.Map;
+import java.util.Set;
 import javax.annotation.Nullable;
 import net.starlark.java.syntax.TokenKind;
 
@@ -133,20 +134,39 @@ static Object binaryOp(TokenKind op, Object x, Object y, StarlarkThread starlark
             // map | map (usually dicts)
             return Dict.builder().putAll((Map<?, ?>) x).putAll((Map<?, ?>) y).build(mu);
           }
+        } else if (x instanceof Set && y instanceof Set) {
+          // set | set
+          if (semantics.getBool(StarlarkSemantics.EXPERIMENTAL_ENABLE_STARLARK_SET)) {
+            return StarlarkSet.empty().union(Tuple.of(x, y), starlarkThread);
+          }
         }
         break;
 
       case AMPERSAND:
         if (x instanceof StarlarkInt && y instanceof StarlarkInt) {
           // int & int
           return StarlarkInt.and((StarlarkInt) x, (StarlarkInt) y);
+        } else if (x instanceof Set && y instanceof Set) {
+          // set & set
+          if (semantics.getBool(StarlarkSemantics.EXPERIMENTAL_ENABLE_STARLARK_SET)) {
+            StarlarkSet<?> xSet =
+                x instanceof StarlarkSet ? (StarlarkSet<?>) x : StarlarkSet.checkedCopyOf(mu, x);
+            return xSet.intersection(Tuple.of(y), starlarkThread);
+          }
         }
         break;
 
       case CARET:
         if (x instanceof StarlarkInt && y instanceof StarlarkInt) {
           // int ^ int
           return StarlarkInt.xor((StarlarkInt) x, (StarlarkInt) y);
+        } else if (x instanceof Set && y instanceof Set) {
+          // set ^ set
+          if (semantics.getBool(StarlarkSemantics.EXPERIMENTAL_ENABLE_STARLARK_SET)) {
+            StarlarkSet<?> xSet =
+                x instanceof StarlarkSet ? (StarlarkSet<?>) x : StarlarkSet.checkedCopyOf(mu, x);
+            return xSet.symmetricDifference(y, starlarkThread);
+          }
         }
         break;
 
@@ -186,6 +206,13 @@ static Object binaryOp(TokenKind op, Object x, Object y, StarlarkThread starlark
             double z = xf - ((StarlarkInt) y).toFiniteDouble();
             return StarlarkFloat.of(z);
           }
+        } else if (x instanceof Set && y instanceof Set) {
+          // set - set
+          if (semantics.getBool(StarlarkSemantics.EXPERIMENTAL_ENABLE_STARLARK_SET)) {
+            StarlarkSet<?> xSet =
+                x instanceof StarlarkSet ? (StarlarkSet<?>) x : StarlarkSet.checkedCopyOf(mu, x);
+            return xSet.difference(Tuple.of(y), starlarkThread);
+          }
         }
         break;
 
@@ -344,8 +371,8 @@ static Object binaryOp(TokenKind op, Object x, Object y, StarlarkThread starlark
         return compare(x, y) >= 0;
 
       case IN:
-        if (y instanceof StarlarkIndexable) {
-          return ((StarlarkIndexable) y).containsKey(semantics, x);
+        if (y instanceof StarlarkMembershipTestable) {
+          return ((StarlarkMembershipTestable) y).containsKey(semantics, x);
         } else if (y instanceof StarlarkIndexable.Threaded) {
           return ((StarlarkIndexable.Threaded) y).containsKey(starlarkThread, semantics, x);
         } else if (y instanceof String) {

src/main/java/net/starlark/java/eval/MethodLibrary.java

@@ -406,7 +406,7 @@ public StarlarkList<?> list(StarlarkIterable<?> x, StarlarkThread thread) throws
   @StarlarkMethod(
       name = "len",
       doc =
-          "Returns the length of a string, sequence (such as a list or tuple), dict, or other"
+          "Returns the length of a string, sequence (such as a list or tuple), dict, set, or other"
               + " iterable.",
       parameters = {@Param(name = "x", doc = "The value whose length to report.")},
       useStarlarkThread = true)
@@ -642,6 +642,38 @@ public StarlarkInt intForStarlark(Object x, Object baseO) throws EvalException {
     return dict;
   }
 
+  @StarlarkMethod(
+      name = "set",
+      doc =
+          """
+Creates a new <a href=\"../core/set.html\">set</a> containing the unique elements of a given
+iterable, preserving iteration order.
+
+<p>If called with no argument, <code>set()</code> returns a new empty set.
+
+<p>For example,
+<pre class=language-python>
+set()                          # an empty set
+set([3, 1, 1, 2])              # set([3, 1, 2]), a set of three elements
+set({"k1": "v1", "k2": "v2"})  # set(["k1", "k2"]), a set of two elements
+</pre>
+""",
+      parameters = {
+        @Param(
+            name = "elements",
+            defaultValue = "[]",
+            doc = "A set, a sequence of hashable values, or a dict."),
+      },
+      useStarlarkThread = true)
+  public StarlarkSet<Object> set(Object elements, StarlarkThread thread) throws EvalException {
+    // Ordinarily we would use StarlarkMethod#enableOnlyWithFlag, but this doesn't work for
+    // top-level symbols, so enforce it here instead.
+    if (!thread.getSemantics().getBool(StarlarkSemantics.EXPERIMENTAL_ENABLE_STARLARK_SET)) {
+      throw Starlark.errorf("Use of set() requires --experimental_enable_starlark_set");
+    }
+    return StarlarkSet.checkedCopyOf(thread.mutability(), elements);
+  }
+
   @StarlarkMethod(
       name = "enumerate",
       doc =

src/main/java/net/starlark/java/eval/Starlark.java

@@ -312,6 +312,8 @@ public static int len(Object x) {
       return ((Sequence) x).size();
     } else if (x instanceof Dict) {
       return ((Dict) x).size();
+    } else if (x instanceof StarlarkSet) {
+      return ((StarlarkSet) x).size();
     } else if (x instanceof StarlarkIterable) {
       // Iterables.size runs in constant time if x implements Collection.
       return Iterables.size((Iterable<?>) x);

src/main/java/net/starlark/java/eval/StarlarkIndexable.java

@@ -17,20 +17,25 @@
 /**
  * A Starlark value that support indexed access ({@code object[key]}) and membership tests ({@code
  * key in object}).
+ *
+ * <p>Implementations of this interface come in three flavors: map-like, sequence-like, and
+ * string-like.
+ *
+ * <ul>
+ *   <li>For map-like objects, 'x in y' should return True when 'y[x]' is valid; otherwise, it
+ *       should either be False or a failure. Examples: dict.
+ *   <li>For sequence-like objects, 'x in y' should return True when 'x == y[i]' for some integer
+ *       'i'; otherwise, it should either be False or a failure. Examples: list, tuple, and string
+ *       (which, notably, is not a {@link Sequence}).
+ *   <li>For string-like objects, 'x in y' should return True when 'x' is a substring of 'y', i.e.
+ *       'x[i] == y[i + n]' for some 'n' and all i in [0, len(x)). Examples: string.
+ * </ul>
  */
-public interface StarlarkIndexable extends StarlarkValue {
+public interface StarlarkIndexable extends StarlarkMembershipTestable {
 
   /** Returns the value associated with the given key. */
   Object getIndex(StarlarkSemantics semantics, Object key) throws EvalException;
 
-  /**
-   * Returns whether the key is in the object. New types should try to follow the semantics of dict:
-   * 'x in y' should return True when 'y[x]' is valid; otherwise, it should either be False or a
-   * failure. Note however that the builtin types string, list, and tuple do not follow this
-   * convention.
-   */
-  boolean containsKey(StarlarkSemantics semantics, Object key) throws EvalException;
-
   /**
    * A variant of {@link StarlarkIndexable} that also provides a StarlarkThread instance on method
    * calls.

src/main/java/net/starlark/java/eval/StarlarkMembershipTestable.java

@@ -0,0 +1,24 @@
+// Copyright 2024 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 net.starlark.java.eval;
+
+/**
+ * A Starlark value that support membership tests ({@code key in object} and {@code key not in
+ * object}).
+ */
+public interface StarlarkMembershipTestable extends StarlarkValue {
+  /** Returns whether the key is in the object. */
+  boolean containsKey(StarlarkSemantics semantics, Object key) throws EvalException;
+}

src/main/java/net/starlark/java/eval/StarlarkSemantics.java

@@ -256,4 +256,7 @@ public final String toString() {
    * unconditionally prohibits recursion.
    */
   public static final String ALLOW_RECURSION = "-allow_recursion";
+
+  /** Whether StarlarkSet objects may be constructed by the interpreter. */
+  public static final String EXPERIMENTAL_ENABLE_STARLARK_SET = "+experimental_enable_starlark_set";
 }