From 08ef30b465e62c98730b75d7a0713ea6b8f61bb8 Mon Sep 17 00:00:00 2001
From: jcater <jcater@google.com>
Date: Fri, 9 Apr 2021 06:45:40 -0700
Subject: [PATCH] Fix toolchains to support type lookup.

Fixes #13320.

PiperOrigin-RevId: 367624002
---
 src/main/java/net/starlark/java/eval/Eval.java      | 5 ++---
 src/main/java/net/starlark/java/eval/EvalUtils.java | 9 ++++++---
 src/test/shell/bazel/toolchain_test.sh              | 2 ++
 3 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/src/main/java/net/starlark/java/eval/Eval.java b/src/main/java/net/starlark/java/eval/Eval.java
index 6d9a35fcdcf1f0..6c915b99c31206 100644
--- a/src/main/java/net/starlark/java/eval/Eval.java
+++ b/src/main/java/net/starlark/java/eval/Eval.java
@@ -430,7 +430,7 @@ private static Object inplaceBinaryOp(StarlarkThread.Frame fr, TokenKind op, Obj
       list.extend(y);
       return list;
     }
-    return EvalUtils.binaryOp(op, x, y, fr.thread.getSemantics(), fr.thread.mutability());
+    return EvalUtils.binaryOp(op, x, y, fr.thread);
   }
 
   // ---- expressions ----
@@ -500,8 +500,7 @@ private static Object evalBinaryOperator(StarlarkThread.Frame fr, BinaryOperator
       default:
         Object y = eval(fr, binop.getY());
         try {
-          return EvalUtils.binaryOp(
-              binop.getOperator(), x, y, fr.thread.getSemantics(), fr.thread.mutability());
+          return EvalUtils.binaryOp(binop.getOperator(), x, y, fr.thread);
         } catch (EvalException ex) {
           fr.setErrorLocation(binop.getOperatorLocation());
           throw ex;
diff --git a/src/main/java/net/starlark/java/eval/EvalUtils.java b/src/main/java/net/starlark/java/eval/EvalUtils.java
index 0dbcf891cbaab9..ae92efbd5234fb 100644
--- a/src/main/java/net/starlark/java/eval/EvalUtils.java
+++ b/src/main/java/net/starlark/java/eval/EvalUtils.java
@@ -73,9 +73,10 @@ static int toIndex(int index, int length) {
   }
 
   /** Evaluates an eager binary operation, {@code x op y}. (Excludes AND and OR.) */
-  static Object binaryOp(
-      TokenKind op, Object x, Object y, StarlarkSemantics semantics, Mutability mu)
+  static Object binaryOp(TokenKind op, Object x, Object y, StarlarkThread starlarkThread)
       throws EvalException {
+    StarlarkSemantics semantics = starlarkThread.getSemantics();
+    Mutability mu = starlarkThread.mutability();
     switch (op) {
       case PLUS:
         if (x instanceof StarlarkInt) {
@@ -339,6 +340,8 @@ static Object binaryOp(
       case IN:
         if (y instanceof StarlarkIndexable) {
           return ((StarlarkIndexable) y).containsKey(semantics, x);
+        } else if (y instanceof StarlarkIndexable.Threaded) {
+          return ((StarlarkIndexable.Threaded) y).containsKey(starlarkThread, semantics, x);
         } else if (y instanceof String) {
           if (!(x instanceof String)) {
             throw Starlark.errorf(
@@ -349,7 +352,7 @@ static Object binaryOp(
         break;
 
       case NOT_IN:
-        Object z = binaryOp(TokenKind.IN, x, y, semantics, mu);
+        Object z = binaryOp(TokenKind.IN, x, y, starlarkThread);
         if (z != null) {
           return !Starlark.truth(z);
         }
diff --git a/src/test/shell/bazel/toolchain_test.sh b/src/test/shell/bazel/toolchain_test.sh
index a599df7bdbfe4f..dbd4c484c0f8f4 100755
--- a/src/test/shell/bazel/toolchain_test.sh
+++ b/src/test/shell/bazel/toolchain_test.sh
@@ -78,6 +78,8 @@ function write_test_rule() {
   mkdir -p toolchain
   cat >> toolchain/rule_${rule_name}.bzl <<EOF
 def _impl(ctx):
+  if '//toolchain:${toolchain_name}' not in ctx.toolchains:
+    fail('Toolchain type //toolchain:${toolchain_name} not found')
   toolchain = ctx.toolchains['//toolchain:${toolchain_name}']
   message = ctx.attr.message
   print(