starlark: allow nested def statements

This change removes the restriction that def statements may not be nested.
As in Python, Starlark's nested def statements are lexically scoped,
and are implemented using closures, a technique first described in 1964
(for Landin's SECD machine) and employed in essentially every language
since Scheme in the 1970s.

The resolver computes the free variables of each function, which are
in effect treated as hidden parameters implicitly supplied from the
environment in which the def statement is executed. Local variables
shared between outer and inner functions are indirect and called Cells.

The tests now use the assert_fails(lambda: expr, "expected error")
construct so that they can test failures without aborting the test chunk.
(There is no lambda syntax yet, but its addition is trivial and will
be done in a follow-up; see CL 345746527.)

RELNOTES: Starlark now permits def statements to be nested (closures).
PiperOrigin-RevId: 346365019

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

@@ -175,13 +175,32 @@ private static void execDef(StarlarkThread.Frame fr, DefStatement node)
       defaults = EMPTY;
     }
 
+    // Capture the cells of the function's
+    // free variables from the lexical environment.
+    Object[] freevars = new Object[rfn.getFreeVars().size()];
+    int i = 0;
+    for (Resolver.Binding bind : rfn.getFreeVars()) {
+      // Unlike expr(Identifier), we want the cell itself, not its content.
+      switch (bind.getScope()) {
+        case FREE:
+          freevars[i++] = fn(fr).getFreeVar(bind.getIndex());
+          break;
+        case CELL:
+          freevars[i++] = fr.locals[bind.getIndex()];
+          break;
+        default:
+          throw new IllegalStateException("unexpected: " + bind);
+      }
+    }
+
     // Nested functions use the same globalIndex as their enclosing function,
     // since both were compiled from the same Program.
     StarlarkFunction fn = fn(fr);
     assignIdentifier(
         fr,
         node.getIdentifier(),
-        new StarlarkFunction(rfn, Tuple.wrap(defaults), fn.getModule(), fn.globalIndex));
+        new StarlarkFunction(
+            rfn, fn.getModule(), fn.globalIndex, Tuple.wrap(defaults), Tuple.wrap(freevars)));
   }
 
   private static TokenKind execIf(StarlarkThread.Frame fr, IfStatement node)
@@ -231,8 +250,8 @@ private static void execLoad(StarlarkThread.Frame fr, LoadStatement node) throws
       // loads bind file-locally. Either way, the resolver should designate
       // the proper scope of binding.getLocalName() and this should become
       // simply assign(binding.getLocalName(), value).
-      // Currently, we update the module but not module.exportedGlobals;
-      // changing it to fr.locals.put breaks a test. TODO(adonovan): find out why.
+      // Currently, we update the module but not module.exportedGlobals.
+      // Change it to a local binding now that closures are supported.
       fn(fr).setGlobal(binding.getLocalName().getBinding().getIndex(), value);
     }
   }
@@ -328,6 +347,9 @@ private static void assignIdentifier(StarlarkThread.Frame fr, Identifier id, Obj
       case LOCAL:
         fr.locals[bind.getIndex()] = value;
         break;
+      case CELL:
+        ((StarlarkFunction.Cell) fr.locals[bind.getIndex()]).x = value;
+        break;
       case GLOBAL:
         // Updates a module binding and sets its 'exported' flag.
         // (Only load bindings are not exported.
@@ -637,6 +659,12 @@ private static Object evalIdentifier(StarlarkThread.Frame fr, Identifier id)
       case LOCAL:
         result = fr.locals[bind.getIndex()];
         break;
+      case CELL:
+        result = ((StarlarkFunction.Cell) fr.locals[bind.getIndex()]).x;
+        break;
+      case FREE:
+        result = fn(fr).getFreeVar(bind.getIndex()).x;
+        break;
       case GLOBAL:
         result = fn(fr).getGlobal(bind.getIndex());
         break;

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

@@ -868,8 +868,6 @@ public static Object execFile(
    */
   public static Object execFileProgram(Program prog, Module module, StarlarkThread thread)
       throws EvalException, InterruptedException {
-    Tuple defaultValues = Tuple.empty();
-
     Resolver.Function rfn = prog.getResolvedFunction();
 
     // A given Module may be passed to execFileProgram multiple times in sequence,
@@ -884,7 +882,13 @@ public static Object execFileProgram(Program prog, Module module, StarlarkThread
     // two array lookups.
     int[] globalIndex = module.getIndicesOfGlobals(rfn.getGlobals());
 
-    StarlarkFunction toplevel = new StarlarkFunction(rfn, defaultValues, module, globalIndex);
+    StarlarkFunction toplevel =
+        new StarlarkFunction(
+            rfn,
+            module,
+            globalIndex,
+            /*defaultValues=*/ Tuple.empty(),
+            /*freevars=*/ Tuple.empty());
     return Starlark.fastcall(thread, toplevel, EMPTY, EMPTY);
   }
 
@@ -928,10 +932,10 @@ public static StarlarkFunction newExprFunction(
       ParserInput input, FileOptions options, Module module) throws SyntaxError.Exception {
     Expression expr = Expression.parse(input, options);
     Program prog = Program.compileExpr(expr, module, options);
-    Tuple defaultValues = Tuple.empty();
     Resolver.Function rfn = prog.getResolvedFunction();
     int[] globalIndex = module.getIndicesOfGlobals(rfn.getGlobals()); // see execFileProgram
-    return new StarlarkFunction(rfn, defaultValues, module, globalIndex);
+    return new StarlarkFunction(
+        rfn, module, globalIndex, /*defaultValues=*/ Tuple.empty(), /*freevars=*/ Tuple.empty());
   }
 
   /**

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

@@ -36,15 +36,33 @@
 public final class StarlarkFunction implements StarlarkCallable {
 
   final Resolver.Function rfn;
-  final int[] globalIndex; // index in Module.globals of ith Program global (binding index)
   private final Module module; // a function closes over its defining module
+
+  // Index in Module.globals of ith Program global (Resolver.Binding(GLOBAL).index).
+  // See explanation at Starlark.execFileProgram.
+  final int[] globalIndex;
+
+  // Default values of optional parameters.
+  // Indices correspond to the subsequence of parameters after the initial
+  // required parameters and before *args/**kwargs.
+  // Contain MANDATORY for the required keyword-only parameters.
   private final Tuple defaultValues;
 
-  StarlarkFunction(Resolver.Function rfn, Tuple defaultValues, Module module, int[] globalIndex) {
+  // Cells (shared locals) of enclosing functions.
+  // Indexed by Resolver.Binding(FREE).index values.
+  private final Tuple freevars;
+
+  StarlarkFunction(
+      Resolver.Function rfn,
+      Module module,
+      int[] globalIndex,
+      Tuple defaultValues,
+      Tuple freevars) {
     this.rfn = rfn;
-    this.globalIndex = globalIndex;
     this.module = module;
+    this.globalIndex = globalIndex;
     this.defaultValues = defaultValues;
+    this.freevars = freevars;
   }
 
   // Sets a global variable, given its index in this function's compiled Program.
@@ -153,9 +171,6 @@ public Module getModule() {
   @Override
   public Object fastcall(StarlarkThread thread, Object[] positional, Object[] named)
       throws EvalException, InterruptedException {
-    if (thread.mutability().isFrozen()) {
-      throw Starlark.errorf("Trying to call in frozen environment");
-    }
     if (!thread.isRecursionAllowed() && thread.isRecursiveCall(this)) {
       throw Starlark.errorf("function '%s' called recursively", getName());
     }
@@ -167,9 +182,19 @@ public Object fastcall(StarlarkThread thread, Object[] positional, Object[] name
     StarlarkThread.Frame fr = thread.frame(0);
     fr.locals = new Object[rfn.getLocals().size()];
     System.arraycopy(arguments, 0, fr.locals, 0, rfn.getParameterNames().size());
+
+    // Spill indicated locals to cells.
+    for (int index : rfn.getCellIndices()) {
+      fr.locals[index] = new Cell(fr.locals[index]);
+    }
+
     return Eval.execFunctionBody(fr, rfn.getBody());
   }
 
+  Cell getFreeVar(int index) {
+    return (Cell) freevars.get(index);
+  }
+
   @Override
   public void repr(Printer printer) {
     // TODO(adonovan): use the file name instead. But that's a breaking Bazel change.
@@ -376,9 +401,20 @@ public boolean isImmutable() {
   }
 
   // The MANDATORY sentinel indicates a slot in the defaultValues
-  // tuple corresponding to a required parameter. It is not visible
-  // to Java or Starlark code.
+  // tuple corresponding to a required parameter.
+  // It is not visible to Java or Starlark code.
   static final Object MANDATORY = new Mandatory();
 
   private static class Mandatory implements StarlarkValue {}
+
+  // A Cell is a local variable shared between an inner and an outer function.
+  // It is a StarlarkValue because it is a stack operand and a Tuple element,
+  // but it is not visible to Java or Starlark code.
+  static final class Cell implements StarlarkValue {
+    Object x;
+
+    Cell(Object x) {
+      this.x = x;
+    }
+  }
 }

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

@@ -32,15 +32,17 @@
  * per-thread application state (see {@link #setThreadLocal}) that passes through Starlark functions
  * but does not directly affect them, such as information about the BUILD file being loaded.
  *
- * <p>Every {@code StarlarkThread} has a {@link Mutability} field, and must be used within a
- * function that creates and closes this {@link Mutability} with the try-with-resource pattern. This
- * {@link Mutability} is also used when initializing mutable objects within that {@code
- * StarlarkThread}. When the {@code Mutability} is closed at the end of the computation, it freezes
- * the {@code StarlarkThread} along with all of those objects. This pattern enforces the discipline
- * that there should be no dangling mutable {@code StarlarkThread}, or concurrency between
- * interacting {@code StarlarkThread}s. It is a Starlark-level error to attempt to mutate a frozen
- * {@code StarlarkThread} or its objects, but it is a Java-level error to attempt to mutate an
- * unfrozen {@code StarlarkThread} or its objects from within a different {@code StarlarkThread}.
+ * <p>StarlarkThreads are not thread-safe: they should be confined to a single Java thread.
+ *
+ * <p>Every StarlarkThread has an associated {@link Mutability}, which should be created for that
+ * thread, and closed once the thread's work is done. (A try-with-resources statement is handy for
+ * this purpose.) Starlark values created by the thread are associated with the thread's Mutability,
+ * so that when the Mutability is closed at the end of the computation, all the values created by
+ * the thread become frozen. This pattern ensures that all Starlark values are frozen before they
+ * are published to another thread, and thus that concurrently executing Starlark threads are free
+ * from data races. Once a thread's mutability is frozen, the thread is unlikely to be useful for
+ * further computation because it can no longer create mutable values. (This is occasionally
+ * valuable in tests.)
  */
 public final class StarlarkThread {
 
@@ -136,7 +138,8 @@ static final class Frame implements Debug.Frame {
     private boolean errorLocationSet;
 
     // The locals of this frame, if fn is a StarlarkFunction, otherwise null.
-    // Set by StarlarkFunction.fastcall.
+    // Set by StarlarkFunction.fastcall. Elements may be regular Starlark
+    // values, or wrapped in StarlarkFunction.Cells if shared with a nested function.
     @Nullable Object[] locals;
 
     @Nullable private Object profileSpan; // current span of walltime call profiler
@@ -181,8 +184,12 @@ public ImmutableMap<String, Object> getLocals() {
       ImmutableMap.Builder<String, Object> env = ImmutableMap.builder();
       if (fn instanceof StarlarkFunction) {
         for (int i = 0; i < locals.length; i++) {
-          if (locals[i] != null) {
-            env.put(((StarlarkFunction) fn).rfn.getLocals().get(i).getName(), locals[i]);
+          Object local = locals[i];
+          if (local != null) {
+            if (local instanceof StarlarkFunction.Cell) {
+              local = ((StarlarkFunction.Cell) local).x;
+            }
+            env.put(((StarlarkFunction) fn).rfn.getLocals().get(i).getName(), local);
           }
         }
       }
@@ -332,9 +339,9 @@ boolean isRecursiveCall(StarlarkFunction fn) {
     // Find fn buried within stack. (The top of the stack is assumed to be fn.)
     for (int i = callstack.size() - 2; i >= 0; --i) {
       Frame fr = callstack.get(i);
-      // TODO(adonovan): compare code, not closure values, otherwise
-      // one can defeat this check by writing the Y combinator.
-      if (fr.fn.equals(fn)) {
+      // We compare code, not closure values, otherwise one can defeat the
+      // check by writing the Y combinator.
+      if (fr.fn instanceof StarlarkFunction && ((StarlarkFunction) fr.fn).rfn.equals(fn.rfn)) {
         return true;
       }
     }