Disallow mutation of values being iterated by a for loop or comprehension

This entails adding a read-locking mechanism to Mutability contexts.

RELNOTES[INC]: Updating list/dicts while they are being looped over is not allowed. Use an explicit copy if needed ("for x in list(mylist):").

--
MOS_MIGRATED_REVID=134442701

3 files changed, 162 insertions, 25 deletions

diff --git a/src/test/java/com/google/devtools/build/lib/syntax/EnvironmentTest.java b/src/test/java/com/google/devtools/build/lib/syntax/EnvironmentTest.java
index c9f6e1275e..3d6dc54be3 100644
--- a/src/test/java/com/google/devtools/build/lib/syntax/EnvironmentTest.java
+++ b/src/test/java/com/google/devtools/build/lib/syntax/EnvironmentTest.java
@@ -20,7 +20,9 @@ import static org.junit.Assert.assertNull;
 import static org.junit.Assert.fail;
 
 import com.google.common.collect.Sets;
+import com.google.devtools.build.lib.events.Location;
 import com.google.devtools.build.lib.syntax.util.EvaluationTestCase;
+import com.google.devtools.build.lib.vfs.PathFragment;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -171,6 +173,60 @@ public class EnvironmentTest extends EvaluationTestCase {
   }
 
   @Test
+  public void testLocked() throws Exception {
+    Mutability mutability = Mutability.create("testLocked");
+    class DummyFreezable implements Mutability.Freezable {
+      @Override
+      public Mutability mutability() {
+        return mutability;
+      }
+    }
+    DummyFreezable dummy = new DummyFreezable();
+    Location locA = Location.fromPathFragment(new PathFragment("/a"));
+    Location locB = Location.fromPathFragment(new PathFragment("/b"));
+    Environment env = Environment.builder(mutability).build();
+
+    // Acquire two locks, release two locks, check along the way.
+    assertThat(mutability.isLocked(dummy)).isFalse();
+    mutability.lock(dummy, locA);
+    assertThat(mutability.isLocked(dummy)).isTrue();
+    mutability.lock(dummy, locB);
+    assertThat(mutability.isLocked(dummy)).isTrue();
+    assertThat(mutability.getLockLocations(dummy)).containsExactly(locA, locB);
+    mutability.unlock(dummy, locA);
+    assertThat(mutability.isLocked(dummy)).isTrue();
+    try {
+      Mutability.checkMutable(dummy, env);
+      fail("Able to mutate locked object");
+    } catch (Mutability.MutabilityException e) {
+      assertThat(e).hasMessage("trying to mutate a locked object (is it currently being iterated "
+          + "over by a for loop or comprehension?)\n"
+          + "Object locked at the following locations: /b:1");
+    }
+    try {
+      mutability.unlock(dummy, locA);
+      fail("Able to unlock object with wrong location");
+    } catch (AssertionError e) {
+      assertThat(e).hasMessage("trying to unlock an object for a location at which "
+          + "it was not locked (/a:1)");
+    }
+    mutability.unlock(dummy, locB);
+    assertThat(mutability.isLocked(dummy)).isFalse();
+    Mutability.checkMutable(dummy, env);
+
+    // Acquire, then freeze.
+    mutability.lock(dummy, locA);
+    mutability.freeze();
+    assertThat(mutability.isLocked(dummy)).isFalse();
+    try {
+      Mutability.checkMutable(dummy, env);
+      fail("Able to mutate locked object");
+    } catch (Mutability.MutabilityException e) {
+      assertThat(e).hasMessage("trying to mutate a frozen object");
+    }
+  }
+
+  @Test
   public void testReadOnly() throws Exception {
     Environment env = newSkylarkEnvironment()
         .setup("special_var", 42)
diff --git a/src/test/java/com/google/devtools/build/lib/syntax/EvaluationTest.java b/src/test/java/com/google/devtools/build/lib/syntax/EvaluationTest.java
index 78f9962d74..d7607dcab4 100644
--- a/src/test/java/com/google/devtools/build/lib/syntax/EvaluationTest.java
+++ b/src/test/java/com/google/devtools/build/lib/syntax/EvaluationTest.java
@@ -578,9 +578,34 @@ public class EvaluationTest extends EvaluationTestCase {
   public void testListComprehensionUpdate() throws Exception {
     new BuildTest()
         .setUp("xs = [1, 2, 3]")
-        .testStatement("[xs.append(4) for x in xs]",
-            MutableList.of(env, Runtime.NONE, Runtime.NONE, Runtime.NONE))
-        .testLookup("xs", MutableList.of(env, 1, 2, 3));
+        .testIfErrorContains("trying to mutate a locked object",
+            "[xs.append(4) for x in xs]");
+  }
+
+  @Test
+  public void testNestedListComprehensionUpdate() throws Exception {
+    new BuildTest()
+        .setUp("xs = [1, 2, 3]")
+        .testIfErrorContains("trying to mutate a locked object",
+            "[xs.append(4) for x in xs for y in xs]");
+  }
+
+  @Test
+  public void testListComprehensionUpdateInClause() throws Exception {
+    new BuildTest()
+        .setUp("xs = [1, 2, 3]")
+        .testIfErrorContains("trying to mutate a locked object",
+            // Use short-circuiting to produce valid output in the event
+            // the exception is not raised.
+            "[y for x in xs for y in (xs.append(4) or xs)]");
+  }
+
+  @Test
+  public void testDictComprehensionUpdate() throws Exception {
+    new BuildTest()
+        .setUp("xs = {1:1, 2:2, 3:3}")
+        .testIfErrorContains("trying to mutate a locked object",
+            "[xs.popitem() for x in xs]");
   }
 
   @Test
diff --git a/src/test/java/com/google/devtools/build/lib/syntax/SkylarkEvaluationTest.java b/src/test/java/com/google/devtools/build/lib/syntax/SkylarkEvaluationTest.java
index b7650e0a4d..81b61732d8 100644
--- a/src/test/java/com/google/devtools/build/lib/syntax/SkylarkEvaluationTest.java
+++ b/src/test/java/com/google/devtools/build/lib/syntax/SkylarkEvaluationTest.java
@@ -357,37 +357,93 @@ public class SkylarkEvaluationTest extends EvaluationTest {
 
   @Test
   public void testForUpdateList() throws Exception {
-    // Check that modifying the list within the loop doesn't affect what gets iterated over.
-    // TODO(brandjon): When we support in-place assignment to a list index, also test that
-    // as a modification here.
     new SkylarkTest().setUp("def foo():",
-        "  xs = ['a', 'b', 'c']",
-        "  s = ''",
+        "  xs = [1, 2, 3]",
         "  for x in xs:",
-        "    s = s + x",
-        "    if x == 'a':",
-        "      xs.pop(0)",
-        "      xs.pop(1)",
-        "    elif x == 'c':",
-        "      xs.append('d')",
-        "  return s",
-        "s = foo()").testLookup("s", "abc");
+        "    if x == 1:",
+        "      xs.append(10)"
+        ).testIfErrorContains("trying to mutate a locked object", "foo()");
   }
 
   @Test
   public void testForUpdateDict() throws Exception {
-    // Check that modifying the dict within the loop doesn't affect what gets iterated over.
     new SkylarkTest().setUp("def foo():",
         "  d = {'a': 1, 'b': 2, 'c': 3}",
-        "  s = ''",
         "  for k in d:",
-        "    s = s + k",
-        "    if k == 'a':",
-        "      d.pop('a')",
-        "      d.pop('b')",
-        "    d.update({'d': 4})",
-        "  return s",
-        "s = foo()").testLookup("s", "abc");
+        "    d[k] *= 2"
+        ).testIfErrorContains("trying to mutate a locked object", "foo()");
+  }
+
+  @Test
+  public void testForUnlockedAfterBreak() throws Exception {
+    new SkylarkTest().setUp("def foo():",
+        "  xs = [1, 2]",
+        "  for x in xs:",
+        "    break",
+        "  xs.append(3)",
+        "  return xs"
+        ).testEval("foo()", "[1, 2, 3]");
+  }
+
+  @Test
+  public void testForNestedOnSameListStillLocked() throws Exception {
+    new SkylarkTest().setUp("def foo():",
+        "  xs = [1, 2]",
+        "  ys = []",
+        "  for x1 in xs:",
+        "    for x2 in xs:",
+        "      ys.append(x1 * x2)",
+        "    xs.append(4)",
+        "  return ys"
+        ).testIfErrorContains("trying to mutate a locked object", "foo()");
+  }
+
+  @Test
+  public void testForNestedOnSameListErrorMessage() throws Exception {
+    new SkylarkTest().setUp("def foo():",
+        "  xs = [1, 2]",
+        "  ys = []",
+        "  for x1 in xs:",
+        "    for x2 in xs:",
+        "      ys.append(x1 * x2)",
+        "      xs.append(4)",
+        "  return ys"
+        // No file name in message, due to how test is set up.
+        ).testIfErrorContains("Object locked at the following locations: 4:3, 5:5", "foo()");
+  }
+
+  @Test
+  public void testForNestedOnSameListUnlockedAtEnd() throws Exception {
+    new SkylarkTest().setUp("def foo():",
+        "  xs = [1, 2]",
+        "  ys = []",
+        "  for x1 in xs:",
+        "    for x2 in xs:",
+        "      ys.append(x1 * x2)",
+        "  xs.append(4)",
+        "  return ys"
+        ).testEval("foo()", "[1, 2, 2, 4]");
+  }
+
+  @Test
+  public void testForNestedWithListCompGood() throws Exception {
+    new SkylarkTest().setUp("def foo():",
+        "  xs = [1, 2]",
+        "  ys = []",
+        "  for x in xs:",
+        "    zs = [None for x in xs for y in (ys.append(x) or ys)]",
+        "  return ys"
+        ).testEval("foo()", "[1, 2, 1, 2]");
+  }
+  @Test
+  public void testForNestedWithListCompBad() throws Exception {
+    new SkylarkTest().setUp("def foo():",
+        "  xs = [1, 2, 3]",
+        "  ys = []",
+        "  for x in xs:",
+        "    zs = [None for x in xs for y in (xs.append(x) or ys)]",
+        "  return ys"
+        ).testIfErrorContains("trying to mutate a locked object", "foo()");
   }
 
   @Test