From 570101d43c204078e0e90fb0e48303de551ea876 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stanis=C5=82aw=20Barzowski?= <stanislaw.barzowski@gmail.com>
Date: Sun, 18 Oct 2020 13:26:37 +0200
Subject: [PATCH] Desugar locals in object comprehension.

Desugar the locals in object comprehensions
"traditionally" instead of handling them manually.

Object comprehensions allow the locals to depend
on the index variable which means that they are separate
for each field. It doesn't make sense to treat them as
a property of the whole object.

Fixes #358.
---
 builtins.go                         | 32 ++++++++---------------------
 internal/program/desugarer.go       | 16 ++++++++++++++-
 testdata/object_comp4.golden        |  4 ++++
 testdata/object_comp4.jsonnet       | 12 +++++++++++
 testdata/object_comp4.linter.golden |  0
 5 files changed, 39 insertions(+), 25 deletions(-)
 create mode 100644 testdata/object_comp4.golden
 create mode 100644 testdata/object_comp4.jsonnet
 create mode 100644 testdata/object_comp4.linter.golden

diff --git a/builtins.go b/builtins.go
index 0dba2dc..0e70455 100644
--- a/builtins.go
+++ b/builtins.go
@@ -1111,30 +1111,25 @@ func builtinUglyObjectFlatMerge(i *interpreter, trace traceElement, x value) (va
 		return nil, err
 	}
 	newFields := make(simpleObjectFieldMap)
-	var anyObj *simpleObject
 	for _, elem := range objarr.elements {
 		obj, err := i.evaluateObject(elem, trace)
 		if err != nil {
 			return nil, err
 		}
+
 		// starts getting ugly - we mess with object internals
 		simpleObj := obj.uncached.(*simpleObject)
+
+		if len(simpleObj.locals) > 0 {
+			panic("Locals should have been desugared in object comprehension.")
+		}
+
 		// there is only one field, really
 		for fieldName, fieldVal := range simpleObj.fields {
 			if _, alreadyExists := newFields[fieldName]; alreadyExists {
 				return nil, i.Error(duplicateFieldNameErrMsg(fieldName), trace)
 			}
 
-			// Here is the tricky part. Each field in a comprehension has different
-			// upValues, because for example in {[v]: v for v in ["x", "y", "z"] },
-			// the v is different for each field.
-			// Yet, even though upValues are field-specific, they are shadowed by object locals,
-			// so we need to make holes to let them pass through
-			upValues := simpleObj.upValues
-			for _, l := range simpleObj.locals {
-				delete(upValues, l.name)
-			}
-
 			newFields[fieldName] = simpleObjectField{
 				hide: fieldVal.hide,
 				field: &bindingsUnboundField{
@@ -1143,24 +1138,13 @@ func builtinUglyObjectFlatMerge(i *interpreter, trace traceElement, x value) (va
 				},
 			}
 		}
-		anyObj = simpleObj
-	}
-
-	var locals []objectLocal
-	var localUpValues bindingFrame
-	if len(objarr.elements) > 0 {
-		// another ugliness - we just take the locals of our last object,
-		// we assume that the locals are the same for each of merged objects
-		locals = anyObj.locals
-		// note that there are already holes for object locals
-		localUpValues = anyObj.upValues
 	}
 
 	return makeValueSimpleObject(
-		localUpValues,
+		nil,
 		newFields,
 		[]unboundField{}, // No asserts allowed
-		locals,
+		nil,
 	), nil
 }
 
diff --git a/internal/program/desugarer.go b/internal/program/desugarer.go
index e9c8dce..881b53b 100644
--- a/internal/program/desugarer.go
+++ b/internal/program/desugarer.go
@@ -204,8 +204,22 @@ func desugarObjectComp(comp *ast.ObjectComp, objLevel int) (ast.Node, error) {
 		return nil, err
 	}
 
+	// Magic merging which follows doesn't support object locals, so we need
+	// to desugar them completely, i.e. put them inside the fields. The locals
+	// can be different for each field in a comprehension (unlike locals in
+	// "normal" objects which have a fixed value), so it's not even too wasteful.
+	if len(obj.Locals) > 0 {
+		field := &obj.Fields[0]
+		field.Body = &ast.Local{
+			Body:  field.Body,
+			Binds: obj.Locals,
+			// TODO(sbarzowski) should I set some NodeBase stuff here?
+		}
+		obj.Locals = nil
+	}
+
 	if len(obj.Fields) != 1 {
-		panic("Too many fields in object comprehension, it should have been caught during parsing")
+		panic("Wrong number of fields in object comprehension, it should have been caught during parsing")
 	}
 
 	desugaredArrayComp, err := desugarForSpec(wrapInArray(obj), &comp.Spec, objLevel)
diff --git a/testdata/object_comp4.golden b/testdata/object_comp4.golden
new file mode 100644
index 0000000..5bc11c7
--- /dev/null
+++ b/testdata/object_comp4.golden
@@ -0,0 +1,4 @@
+{
+   "a": "A",
+   "b": "B"
+}
diff --git a/testdata/object_comp4.jsonnet b/testdata/object_comp4.jsonnet
new file mode 100644
index 0000000..ee17c68
--- /dev/null
+++ b/testdata/object_comp4.jsonnet
@@ -0,0 +1,12 @@
+local data = {
+  a: 'A',
+  b: 'B',
+};
+
+local process(input) = {
+  local v = input[k],
+  [k]: v
+  for k in std.objectFields(input)
+};
+
+process(data)
diff --git a/testdata/object_comp4.linter.golden b/testdata/object_comp4.linter.golden
new file mode 100644
index 0000000..e69de29