Fix multiple object comprehension bugs (#358)

This PR fixes multiple bugs in object comprehensions (fixes #357, fixes
#331, plus another bug).

## Mis-handling of `self` and `super` in object comprehension locals
when inherited

Object comprehensions may define local variables and those variables may
reference `self` or `super`.

Consider the reproduction reported in
#357:

```jsonnet
local lib = {
  foo()::
    {
      local global = self,

      [iterParam]: global.base {
        foo: iterParam
      }
      for iterParam in ["foo"]
    },
};

{
 base:: {}
}
+ lib.foo()
```

This is supposed to output

```json
{
  "foo": {
    "foo": "foo"
  }
}
```

but sjsonnet outputs

```
sjsonnet.Error: Field does not exist: base
      at [Select base].(:7:26)
```

This bug occurs because of how bindings were managed: there's some
tricky circular reference and variable rebinding logic (which I believe
was added to handle references _between_ locals) and that was rebinding
the `self` reference to the comprehension result itself: that is wrong
because the `self` reference may change if an object is inherited or
extended.

While digging into this, I also discovered a related bug where `super`
wasn't being properly re-bound in object comprehension locals or values:

```jsonnet
local lib = {
  foo():: {
    local sx = super.x,
    [k]: sx + 1
    for k in ["x"]
  },
};

{ x: 2 }
+ lib.foo()
```

was failing with

```
java.lang.Exception: sjsonnet.Error: Attempt to use `super` when there is no super class
      at [SelectSuper x].(:4:21)
      at [ValidId sx].(:5:10)
      at [BinaryOp +].(:5:13)
```

## Silent dropping of fields if key is not a string

The object comprehension

```jsonnet
{[k]: k for k in [1]}
```

fails in regular jsonnet with an error

```
RUNTIME ERROR: field must be string, got: number
        <cmdline>:1:1-22
```

but in sjsonnet it was silently returning `{}`. This regression was
[introduced](https://github.com/databricks/sjsonnet/pull/226/files#diff-7e786f55f42d493f85bfb37e5181c1c942ea039b21a307680e4e6913a8084aa2L628-R631)
in #226 by adding a default `case _ =>` (presumably to fix a compiler or
IDE warning).

## Support function definition via object comprehensions

This fixes #331 by adding
parser support for defining functions using method syntax, e.g.

```jsonnet
{
  [a](x): x * 2
  for a in ["f1", "f2", "f3"]
}
```

## Fixes

This PR fixes the above bugs by changing logic within
`Evaluator.visitObjComp`. It is best reviewed commit-by-commit or using
a whitespace-free diff.

Author: JoshRosen

sjsonnet/src/sjsonnet/Evaluator.scala

@@ -680,15 +680,15 @@ class Evaluator(
         visitExpr(k) match {
           case Val.Str(_, k1) => k1
           case Val.Null(_)    => null
-          case x =>
-            Error.fail(
-              s"Field name must be string or null, not ${x.prettyName}",
-              pos
-            )
+          case x              => fieldNameTypeError(x, pos)
         }
     }
   }
 
+  private def fieldNameTypeError(fieldName: Val, pos: Position): Nothing = {
+    Error.fail(s"Field name must be string or null, not ${fieldName.prettyName}", pos)
+  }
+
   def visitMethod(rhs: Expr, params: Params, outerPos: Position)(implicit
       scope: ValScope): Val.Func =
     new Val.Func(outerPos, scope, params) {
@@ -697,20 +697,16 @@ class Evaluator(
       override def evalDefault(expr: Expr, vs: ValScope, es: EvalScope) = visitExpr(expr)(vs)
     }
 
-  def visitBindings(
-      bindings: Array[Bind],
-      scope: (Val.Obj, Val.Obj) => ValScope): Array[(Val.Obj, Val.Obj) => Lazy] = {
-    val arrF = new Array[(Val.Obj, Val.Obj) => Lazy](bindings.length)
+  def visitBindings(bindings: Array[Bind], scope: => ValScope): Array[Lazy] = {
+    val arrF = new Array[Lazy](bindings.length)
     var i = 0
     while (i < bindings.length) {
       val b = bindings(i)
       arrF(i) = b.args match {
         case null =>
-          (self: Val.Obj, sup: Val.Obj) =>
-            new LazyWithComputeFunc(() => visitExpr(b.rhs)(scope(self, sup)))
+          new LazyWithComputeFunc(() => visitExpr(b.rhs)(scope))
         case argSpec =>
-          (self: Val.Obj, sup: Val.Obj) =>
-            new LazyWithComputeFunc(() => visitMethod(b.rhs, argSpec, b.pos)(scope(self, sup)))
+          new LazyWithComputeFunc(() => visitMethod(b.rhs, argSpec, b.pos)(scope))
       }
       i += 1
     }
@@ -841,42 +837,37 @@ class Evaluator(
   def visitObjComp(e: ObjBody.ObjComp, sup: Val.Obj)(implicit scope: ValScope): Val.Obj = {
     val binds = e.preLocals ++ e.postLocals
     val compScope: ValScope = scope // .clearSuper
-
-    lazy val newSelf: Val.Obj = {
-      val builder = new java.util.LinkedHashMap[String, Val.Obj.Member]
-      for (s <- visitComp(e.first :: e.rest, Array(compScope))) {
-        lazy val newScope: ValScope = s.extend(newBindings, newSelf, null)
-
-        lazy val newBindings = visitBindings(binds, (self, sup) => newScope)
-
-        visitExpr(e.key)(s) match {
-          case Val.Str(_, k) =>
-            val prev_length = builder.size()
-            builder.put(
-              k,
-              new Val.Obj.Member(e.plus, Visibility.Normal) {
-                def invoke(self: Val.Obj, sup: Val.Obj, fs: FileScope, ev: EvalScope): Val =
-                  visitExpr(e.value)(
-                    s.extend(newBindings, self, null)
-                  )
+    val builder = new java.util.LinkedHashMap[String, Val.Obj.Member]
+    for (s <- visitComp(e.first :: e.rest, Array(compScope))) {
+      visitExpr(e.key)(s) match {
+        case Val.Str(_, k) =>
+          val prev_length = builder.size()
+          builder.put(
+            k,
+            new Val.Obj.Member(e.plus, Visibility.Normal) {
+              def invoke(self: Val.Obj, sup: Val.Obj, fs: FileScope, ev: EvalScope): Val = {
+                // There is a circular dependency between `newScope` and `newBindings` because
+                // bindings may refer to other bindings (e.g. chains of locals that build on
+                // each other):
+                lazy val newScope: ValScope = s.extend(newBindings, self, sup)
+                lazy val newBindings = visitBindings(binds, newScope)
+                visitExpr(e.value)(newScope)
               }
-            )
-            if (prev_length == builder.size() && settings.noDuplicateKeysInComprehension) {
-              Error.fail(s"Duplicate key ${k} in evaluated object comprehension.", e.pos);
             }
-          case Val.Null(_) => // do nothing
-          case _           =>
-        }
-      }
-      val valueCache = if (sup == null) {
-        Val.Obj.getEmptyValueCacheForObjWithoutSuper(builder.size())
-      } else {
-        new java.util.HashMap[Any, Val]()
+          )
+          if (prev_length == builder.size() && settings.noDuplicateKeysInComprehension) {
+            Error.fail(s"Duplicate key ${k} in evaluated object comprehension.", e.pos);
+          }
+        case Val.Null(_) => // do nothing
+        case x           => fieldNameTypeError(x, e.pos)
       }
-      new Val.Obj(e.pos, builder, false, null, sup, valueCache)
     }
-
-    newSelf
+    val valueCache = if (sup == null) {
+      Val.Obj.getEmptyValueCacheForObjWithoutSuper(builder.size())
+    } else {
+      new java.util.HashMap[Any, Val]()
+    }
+    new Val.Obj(e.pos, builder, false, null, sup, valueCache)
   }
 
   @tailrec

sjsonnet/src/sjsonnet/Parser.scala

@@ -430,8 +430,20 @@ class Parser(
       val preLocals = exprs
         .takeWhile(_.isInstanceOf[Expr.Bind])
         .map(_.asInstanceOf[Expr.Bind])
-      val Expr.Member.Field(offset, Expr.FieldName.Dyn(lhs), plus, args, Visibility.Normal, rhs) =
+      val Expr.Member.Field(
+        offset,
+        Expr.FieldName.Dyn(lhs),
+        plus,
+        args,
+        Visibility.Normal,
+        rhsBody
+      ) =
         exprs(preLocals.length): @unchecked
+      val rhs = if (args == null) {
+        rhsBody
+      } else {
+        Expr.Function(offset, args, rhsBody)
+      }
       val postLocals = exprs
         .drop(preLocals.length + 1)
         .takeWhile(_.isInstanceOf[Expr.Bind])

sjsonnet/src/sjsonnet/ValScope.scala

@@ -18,23 +18,11 @@ final class ValScope private (val bindings: Array[Lazy]) extends AnyVal {
 
   def length: Int = bindings.length
 
-  def extend(
-      newBindingsF: Array[(Val.Obj, Val.Obj) => Lazy] = null,
-      newSelf: Val.Obj,
-      newSuper: Val.Obj): ValScope = {
-    val by = if (newBindingsF == null) 2 else newBindingsF.length + 2
-    val b = Arrays.copyOf(bindings, bindings.length + by)
+  def extend(newBindings: Array[Lazy], newSelf: Val.Obj, newSuper: Val.Obj): ValScope = {
+    val b = Arrays.copyOf(bindings, bindings.length + newBindings.length + 2)
     b(bindings.length) = newSelf
     b(bindings.length + 1) = newSuper
-    if (newBindingsF != null) {
-      var i = 0
-      var j = bindings.length + 2
-      while (i < newBindingsF.length) {
-        b(j) = newBindingsF(i).apply(newSelf, newSuper)
-        i += 1
-        j += 1
-      }
-    }
+    System.arraycopy(newBindings, 0, b, bindings.length + 2, newBindings.length)
     new ValScope(b)
   }
 

sjsonnet/test/src-jvm/sjsonnet/ErrorTestsJvmOnly.scala

@@ -22,6 +22,17 @@ object ErrorTestsJvmOnly extends TestSuite {
   }
 
   val tests: Tests = Tests {
+    // Hack: this test suite may flakily fail if this suite runs prior to other error tests:
+    // if classloading or linking happens inside the StackOverflowError handling then that
+    // may will trigger a secondary StackOverflowError and cause the test to fail.
+    // As a temporary solution, we redundantly run one of the other error tests first.
+    // A better long term solution would be to change how we handle StackOverflowError
+    // to avoid this failure mode, but for now we add this hack to avoid CI flakiness:
+    test("02") - check(
+      """sjsonnet.Error: Foo.
+        |    at [Error].(sjsonnet/test/resources/test_suite/error.02.jsonnet:17:1)
+        |""".stripMargin
+    )
     test("array_recursive_manifest") - check(
       """sjsonnet.Error: Stackoverflow while materializing, possibly due to recursive value
         |    at .(sjsonnet/test/resources/test_suite/error.array_recursive_manifest.jsonnet:17:12)

sjsonnet/test/src/sjsonnet/EvaluatorTests.scala

@@ -145,6 +145,97 @@ object EvaluatorTests extends TestSuite {
         """{local y = $["2"], [x]: if x == "1" then y else 0, for x in ["1", "2"]}["1"]""",
         useNewEvaluator = useNewEvaluator
       ) ==> ujson.Num(0)
+      // References between locals in an object comprehension:
+      eval(
+        """{local a = 1, local b = a + 1, [k]: b + 1 for k in ["x"]}""",
+        useNewEvaluator = useNewEvaluator
+      ) ==> ujson.Obj("x" -> ujson.Num(3))
+      // Locals which reference variables from the comprehension:
+      eval(
+        """{local x2 = k*2, [std.toString(k)]: x2 for k in [1]}""",
+        useNewEvaluator = useNewEvaluator
+      ) ==> ujson.Obj("1" -> ujson.Num(2))
+      // Regression test for https://github.com/databricks/sjsonnet/issues/357
+      // self references in object comprehension locals are properly rebound during inheritance:
+      eval(
+        """
+          |local lib = {
+          |  foo()::
+          |    {
+          |      local global = self,
+          |
+          |      [iterParam]: global.base {
+          |        foo: iterParam
+          |      }
+          |      for iterParam in ["foo"]
+          |    },
+          |};
+          |
+          |{
+          | base:: {}
+          |}
+          |+ lib.foo()
+          |""".stripMargin,
+        useNewEvaluator = useNewEvaluator
+      ) ==> ujson.Obj("foo" -> ujson.Obj("foo" -> "foo"))
+      // Regression test for a related bug involving local references to `super`:
+      eval(
+        """
+          |local lib = {
+          |  foo():: {
+          |    local sx = super.x,
+          |    [k]: sx + 1
+          |    for k in ["x"]
+          |  },
+          |};
+          |
+          |{ x: 2 }
+          |+ lib.foo()
+          |""".stripMargin,
+        useNewEvaluator = useNewEvaluator
+      ) ==> ujson.Obj("x" -> ujson.Num(3))
+      // Yet another related bug involving super references _not_ in locals:
+      eval(
+        """
+          |local lib = {
+          |  foo():: {
+          |    [k]: super.x + 1
+          |    for k in ["x"]
+          |  },
+          |};
+          |
+          |{ x: 2 }
+          |+ lib.foo()
+          |""".stripMargin,
+        useNewEvaluator = useNewEvaluator
+      ) ==> ujson.Obj("x" -> ujson.Num(3))
+      // Regression test for a bug in handling of non-string field names:
+      evalErr("{[k]: k for k in [1]}", useNewEvaluator = useNewEvaluator) ==>
+        """sjsonnet.Error: Field name must be string or null, not number
+          |at .(:1:2)""".stripMargin
+      // Basic function support:
+      eval(
+        """
+          |local funcs = {
+          |  [a](x): x * 2
+          |  for a in ["f1", "f2", "f3"]
+          |};
+          |funcs.f1(10)
+          |""".stripMargin,
+        useNewEvaluator = useNewEvaluator
+      ) ==> ujson.Num(20)
+      // Functions which use locals from the comprehension:
+      eval(
+        """
+          |local funcs = {
+          |  local y = b,
+          |  [a](x): x * y
+          |  for a in ["f1", "f2", "f3"] for b in [2]
+          |};
+          |funcs.f1(10)
+          |""".stripMargin,
+        useNewEvaluator = useNewEvaluator
+      ) ==> ujson.Num(20)
     }
     test("super") {
       test("implicit") {