Fix multiple object comprehension bugs

sjsonnet/src/sjsonnet/Evaluator.scala

@@ -841,42 +841,40 @@ 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 = {
+                lazy val newScope: ValScope = s.extend(newBindings, self, sup)
+                lazy val newBindings = visitBindings(binds, (self, sup) => newScope)
+                visitExpr(e.value)(
+                  s.extend(newBindings, self, null)
+                )
               }
-            )
-            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 =>
+          Error.fail(
+            s"Field name must be string or null, not ${x.prettyName}",
+            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/test/src/sjsonnet/EvaluatorTests.scala

@@ -145,6 +145,58 @@ 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]}"""
+      ) ==> 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))
+      // 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
     }
     test("super") {
       test("implicit") {