Create a new interface to allow Starlark objects to get a thread when…

… getIndex is called. PiperOrigin-RevId: 367454604
katre · Jul 13, 2021 · 0cbb8a8 · 0cbb8a8
1 parent e376580
commit 0cbb8a8
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 4 deletions.
diff --git a/src/main/java/net/starlark/java/eval/Eval.java b/src/main/java/net/starlark/java/eval/Eval.java
@@ -403,7 +403,7 @@ private static void execAugmentedAssignment(StarlarkThread.Frame fr, AssignmentS
       IndexExpression index = (IndexExpression) lhs;
       Object object = eval(fr, index.getObject());
       Object key = eval(fr, index.getKey());
-      Object x = EvalUtils.index(fr.thread.mutability(), fr.thread.getSemantics(), object, key);
+      Object x = EvalUtils.index(fr.thread, object, key);
       // Evaluate rhs after lhs.
       Object y = eval(fr, rhs);
       Object z = inplaceBinaryOp(fr, op, x, y);
@@ -703,7 +703,7 @@ private static Object evalIndex(StarlarkThread.Frame fr, IndexExpression index)
     Object object = eval(fr, index.getObject());
     Object key = eval(fr, index.getKey());
     try {
-      return EvalUtils.index(fr.thread.mutability(), fr.thread.getSemantics(), object, key);
+      return EvalUtils.index(fr.thread, object, key);
     } catch (EvalException ex) {
       fr.setErrorLocation(index.getLbracketLocation());
       throw ex;

diff --git a/src/main/java/net/starlark/java/eval/EvalUtils.java b/src/main/java/net/starlark/java/eval/EvalUtils.java
@@ -431,9 +431,14 @@ static Object unaryOp(TokenKind op, Object x) throws EvalException {
    *
    * @throws EvalException if {@code object} is not a sequence or mapping.
    */
-  static Object index(Mutability mu, StarlarkSemantics semantics, Object object, Object key)
+  static Object index(StarlarkThread starlarkThread, Object object, Object key)
       throws EvalException {
-    if (object instanceof StarlarkIndexable) {
+    Mutability mu = starlarkThread.mutability();
+    StarlarkSemantics semantics = starlarkThread.getSemantics();
+
+    if (object instanceof StarlarkIndexable.Threaded) {
+      return ((StarlarkIndexable.Threaded) object).getIndex(starlarkThread, semantics, key);
+    } else if (object instanceof StarlarkIndexable) {
       Object result = ((StarlarkIndexable) object).getIndex(semantics, key);
       // TODO(bazel-team): We shouldn't have this fromJava call here. If it's needed at all,
       // it should go in the implementations of StarlarkIndexable#getIndex that produce non-Starlark

diff --git a/src/main/java/net/starlark/java/eval/StarlarkIndexable.java b/src/main/java/net/starlark/java/eval/StarlarkIndexable.java
@@ -30,4 +30,20 @@ public interface StarlarkIndexable extends StarlarkValue {
    * convention.
    */
   boolean containsKey(StarlarkSemantics semantics, Object key) throws EvalException;
+
+  /**
+   * A variant of {@link StarlarkIndexable} that also provides a StarlarkThread instance on method
+   * calls.
+   */
+  // TODO(brandjon): Consider replacing this subinterface by changing StarlarkIndexable's methods'
+  // signatures to take StarlarkThread in place of StarlarkSemantics.
+  interface Threaded {
+    /** {@see StarlarkIndexable.getIndex} */
+    Object getIndex(StarlarkThread starlarkThread, StarlarkSemantics semantics, Object key)
+        throws EvalException;
+
+    /** {@see StarlarkIndexable.containsKey} */
+    boolean containsKey(StarlarkThread starlarkThread, StarlarkSemantics semantics, Object key)
+        throws EvalException;
+  }
 }