Enhanced comprehensions use … syntax

src/comprehension.lsc

@@ -3,7 +3,7 @@ import { transformTails } from './tails'
 import { isa } from './is'
 import { toStatement } from './blocks'
 
-import { getLoc, placeAtLoc as atLoc, placeAtNode as atNode, getSurroundingLoc, span } from 'ast-loc-utils'
+import { getLoc, placeAtLoc as atLoc, placeAtNode as atNode, getSurroundingLoc, span, placeTreeAtLocWhenUnplaced as allAtLoc } from 'ast-loc-utils'
 
 validateComprehensionBody(path) ->
   path.traverse({
@@ -14,7 +14,7 @@ validateComprehensionBody(path) ->
     AwaitExpression(awaitPath) ->
       throw awaitPath.buildCodeFrameError(
         "`await` is not allowed within Comprehensions; " +
-        "instead, await the Comprehension (eg; `y <- [for x of xs: x]`)."
+        "instead, await the Comprehension."
       )
 
     YieldExpression(yieldPath) ->
@@ -37,32 +37,59 @@ iife(body, id, initializer) ->
     []
   )~atLoc(loc)
 
-retailObject(path, id, transformPathName, returnPathName) ->
+isSimpleObject(objExpr) ->
+  objExpr.properties?.length == 1 and
+  (not objExpr~isa("ObjectComprehension")) and
+  objExpr.properties[0].type == "ObjectProperty" and
+  (not objExpr.properties[0].decorators?.length)
+
+retailObject(info, path, id, transformPathName, returnPathName) ->
   transformPath = path.get(transformPathName)
   validateComprehensionBody(transformPath)
   transformTails(
     transformPath
     true
     false
-    (seqExpr, tailPath) ->
-      if (
-        seqExpr.type !== "SequenceExpression" or
-        seqExpr.expressions.length !== 2
-      ):
-        throw tailPath.buildCodeFrameError("Object comprehensions must end" +
-        " with a (key, value) pair.")
-
-      [ keyExpr, valExpr ] = seqExpr.expressions
-
-      t.assignmentExpression("=",
-        t.memberExpression(id, keyExpr, true)~atNode(seqExpr),
-        valExpr
-      )~atNode(seqExpr)
+    (expr, tailPath) ->
+      if info.isLegacy:
+        if (
+          expr.type !== "SequenceExpression" or
+          expr.expressions.length !== 2
+        ):
+          throw tailPath.buildCodeFrameError("Object comprehensions must end" +
+          " with a (key, value) pair.")
+
+        [ keyExpr, valExpr ] = expr.expressions
+
+        return t.assignmentExpression("=",
+          t.memberExpression(id, keyExpr, true)~atNode(expr),
+          valExpr
+        )~atNode(expr)
+
+      if not expr~isa("ObjectExpression"):
+        throw tailPath.buildCodeFrameError("Object comprehensions must end with an object expression.")
+
+      if expr~isSimpleObject():
+        // Simple object case: { [k]: v } --> obj[k] = v
+        { properties: [prop] } = expr
+        t.assignmentExpression("=",
+          t.memberExpression(id, prop.key, prop.computed)~atNode(expr),
+          prop.value
+        )~atNode(expr)
+      else:
+        // Complex object case: { ... } -> Object.assign(obj, { ... })
+        t.callExpression(
+          t.memberExpression(
+            t.identifier("Object")
+            t.identifier("assign")
+          )~allAtLoc(expr~getLoc())
+          [id, expr]
+        )~atNode(expr)
   )
 
   path.get(returnPathName).node
 
-retailArray(path, id, transformPathName, returnPathName) ->
+retailArray(info, path, id, transformPathName, returnPathName) ->
   transformPath = path.get(transformPathName)
   validateComprehensionBody(transformPath)
   transformTails(
@@ -74,23 +101,38 @@ retailArray(path, id, transformPathName, returnPathName) ->
         t.memberExpression(id, t.identifier("push")~atNode(expr))~atNode(expr)
         [expr]
       )~atNode(expr)
+
+      // XXX: below code is for allowing ArrayExpressions in tail position
+      // if not expr~isa("ArrayExpression"):
+      //   throw tailPath.buildCodeFrameError("Array comprehensions must end with an array expression.")
+
+      // t.callExpression(
+      //   t.memberExpression(id, t.identifier("push")~atNode(expr))~atNode(expr)
+      //   if expr.elements?.length == 1 and (not expr~isa("ArrayComprehension")):
+      //     // Shortcut for simple array exprs: just array.push the single entry.
+      //     [expr.elements[0]]
+      //   else:
+      //     // ES6-spread the tail array onto the base array
+      //     [t.spreadElement(expr)~atNode(expr)]
+      // )~atNode(expr)
   )
 
   path.get(returnPathName).node
 
-transformLoop(path, ref, isObject, stmts) ->
-  if isObject:
-    stmts.push(retailObject(path, ref, "loop.body", "loop"))
+transformLoop(info, path, ref, stmts) ->
+  if info.isObject:
+    stmts.push(retailObject(info, path, ref, "loop.body", "loop"))
   else:
-    stmts.push(retailArray(path, ref, "loop.body", "loop"))
+    stmts.push(retailArray(info, path, ref, "loop.body", "loop"))
 
-transformCase(path, ref, isObject, stmts) ->
-  if isObject:
-    stmts.push(retailObject(path, ref, "conditional", "conditional"))
+transformCase(info, path, ref, stmts) ->
+  if info.isObject:
+    stmts.push(retailObject(info, path, ref, "conditional", "conditional"))
   else:
-    stmts.push(retailArray(path, ref, "conditional", "conditional"))
+    stmts.push(retailArray(info, path, ref, "conditional", "conditional"))
 
-pushBundle(stmts, ref, bundle, isObject) ->
+pushBundle(info, stmts, ref, bundle) ->
+  { isObject } = info
   if isObject:
     // _ref.k1 = v1; _ref.k2 = v2; ...
     for elem property in bundle:
@@ -111,36 +153,39 @@ pushBundle(stmts, ref, bundle, isObject) ->
       )~atLoc(loc)~toStatement()
     )
 
-export transformComprehension(path, isObject) ->
+export transformComprehension(info) ->
+  { path, isObject } = info
   { node } = path
   elements = if isObject: node.properties else: node.elements
   nodeKey = if isObject: "properties" else: "elements"
   stmts = []
   id = path.scope.generateUidIdentifier(isObject ? "obj" : "arr")~t.clone()~atLoc(getLoc(node)~span(1))
 
   let i = 0, len = elements.length, bundle = [], first = true, initializer
+
   while i < len:
     element = elements[i]
     if element~isa("Comprehension"):
       if first:
         now initializer = bundle
         now first = false
       else:
-        if bundle.length > 0: stmts~pushBundle(id, bundle, isObject)
+        if bundle.length > 0: pushBundle(info, stmts, id, bundle)
       now bundle = []
 
       match element:
         | ~isa("LoopComprehension"):
-          path.get(`${nodeKey}.${i}`)~transformLoop(id, isObject, stmts)
+          info~transformLoop(path.get(`${nodeKey}.${i}`), id, stmts)
         | ~isa("CaseComprehension"):
-          path.get(`${nodeKey}.${i}`)~transformCase(id, isObject, stmts)
-        | else: throw new Error("Invalid comprehension node (this is an internal error)")
+          info~transformCase(path.get(`${nodeKey}.${i}`), id, stmts)
+        | else:
+          throw new Error("Invalid comprehension node (this is an internal error)")
     else:
       bundle.push(element)
 
     i++
 
-  if bundle.length > 0: stmts~pushBundle(id, bundle, isObject)
+  if bundle.length > 0: pushBundle(info, stmts, id, bundle)
 
   initializerLoc = if initializer.length == 0:
     getLoc(node)~span(1)
@@ -154,36 +199,31 @@ export transformComprehension(path, isObject) ->
 
   path.replaceWith(stmts~iife(id, finalInitializer))
 
+getComprehensionInfo(path, isObject, isLegacy) ->
+  { path, isObject, isLegacy }
+
 export transformArrayComprehension(path): void ->
-  transformComprehension(path, false)
+  getComprehensionInfo(path, false, false)~transformComprehension()
 
 export transformObjectComprehension(path): void ->
-  transformComprehension(path, true)
+  getComprehensionInfo(path, true, false)~transformComprehension()
 
-export transformPlainArrayComprehension(path): void ->
-  // Shim V1 onto V2
+export transformLegacyComprehension(path, isObject): void ->
+  // Shim legacy comprehensions onto new model
   { node } = path
   { loop } = node
   if loop:
     delete node.loop
     // TODO: fix babel patch so there's a builder for this...
-    node.elements = [ {
+    node[if isObject: "properties" else: "elements"] = [ {
       type: "LoopComprehension"
       loop
     }~atNode(node) ]
     path.replaceWith(node)
-    transformArrayComprehension(path)
+    getComprehensionInfo(path, isObject, true)~transformComprehension()
+
+export transformPlainArrayComprehension(path): void ->
+  transformLegacyComprehension(path, false)
 
 export transformPlainObjectComprehension(path): void ->
-  // Shim V1 onto V2
-  { node } = path
-  { loop } = node
-  if loop:
-    delete node.loop
-    // TODO: fix babel patch so there's a builder for this...
-    node.properties = [ {
-      type: "LoopComprehension"
-      loop
-    }~atNode(node) ]
-    path.replaceWith(node)
-    transformObjectComprehension(path)
+  transformLegacyComprehension(path, true)

src/config.lsc

@@ -94,7 +94,7 @@ export getParserOpts(pluginOpts, initialParserOpts) ->
   if pluginOpts?.bangCall != false: plugins.push("bangCall")
   if not pluginOpts?.noEnforcedSubscriptIndentation: plugins.push("enforceSubscriptIndentation")
   if pluginOpts?.flippedImports: plugins.push("flippedImports")
-  if pluginOpts?.enhancedComprehension: plugins.push("enhancedComprehension")
+  if pluginOpts?.enhancedComprehension: plugins.push("splatComprehension")
   if pluginOpts?.whiteblock: plugins.push("whiteblockOnly")
   if pluginOpts?.placeholderArgs: plugins.push("syntacticPlaceholder")
   if pluginOpts?.placeholder:

test/fixtures/comprehensions/await/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "`await` is not allowed within Comprehensions; instead, await the Comprehension (eg; `y <- [for x of xs: x]`)."
+  "throws": "`await` is not allowed within Comprehensions; instead, await the Comprehension."
 }

test/fixtures/enhanced-comprehensions/assignment-expr/actual.js

@@ -1,3 +1,3 @@
-[for idx i in Array(10):
+[...for idx i in Array(10):
   now x = f(i)
 ]

test/fixtures/enhanced-comprehensions/assignment-expr/exec.js

@@ -1,5 +1,5 @@
 let x = 0
-result = [for idx i in Array(10):
+result = [...for idx i in Array(10):
   now x = i
 ]
 assert.deepEqual(result, [0,1,2,3,4,5,6,7,8,9])

test/fixtures/enhanced-comprehensions/await/actual.js

@@ -1,4 +1,4 @@
 f() -/>
-  <- [for now x of arr:
+  <- [...for now x of arr:
     <- x
   ]

test/fixtures/enhanced-comprehensions/await/options.json

@@ -1,3 +1,3 @@
 {
-  "throws": "`await` is not allowed within Comprehensions; instead, await the Comprehension (eg; `y <- [for x of xs: x]`)."
+  "throws": "`await` is not allowed within Comprehensions; instead, await the Comprehension."
 }

test/fixtures/enhanced-comprehensions/basic/actual.js

@@ -1 +1 @@
-[for elem x in y: x]
+[...for elem x in y: x]

test/fixtures/enhanced-comprehensions/closure-nested-deep/actual.js

@@ -1,9 +1,9 @@
-[for idx i in Array(10):
+[...for idx i in Array(10):
   for idx j in a:
     if i < 5:
       f() ->
-        {for idx k in Array(10):
+        {...for idx k in Array(10):
           if k > 7:
-            (k, g() -> function h() { [i,j,k] })
+            {[k]: g() -> function h() { [i,j,k] }}
         }
 ]

test/fixtures/enhanced-comprehensions/closure-nested/actual.js

@@ -1 +1 @@
-[ for idx i in Array(10): f() -> g() -> i ]
+[ ...for idx i in Array(10): f() -> g() -> i ]

test/fixtures/enhanced-comprehensions/closure-nested/exec.js

@@ -0,0 +1,4 @@
+closures = [ ...for idx i in Array(10): f() -> g() -> i ]
+closureResults = [ ...for elem f in closures: f()() ]
+
+assert.deepEqual(closureResults, [0, 1, 2, 3, 4, 5, 6, 7, 8, 9])

test/fixtures/enhanced-comprehensions/closure/actual.js

@@ -1 +1 @@
-[ for idx i in Array(10): f() -> i ]
+[ ...for idx i in Array(10): f() -> i ]

test/fixtures/enhanced-comprehensions/closure/exec.js

@@ -0,0 +1,4 @@
+closures = [ ...for idx i in Array(10): f() -> i ]
+closureResults = [ ...for elem f in closures: f() ]
+
+assert.deepEqual(closureResults, [0, 1, 2, 3, 4, 5, 6, 7, 8, 9])

test/fixtures/enhanced-comprehensions/declared-inner-function/actual.js

@@ -1,4 +1,4 @@
-closures = [for idx i in Array(3):
+closures = [...for idx i in Array(3):
   x = g(i)
   g(x) -> x+1
 ]

test/fixtures/enhanced-comprehensions/declared-inner-function/exec.js

@@ -0,0 +1,7 @@
+closures = [...for idx i in Array(3):
+  x = g(i)
+  g(x) -> x+i
+]
+results = [...for elem f in closures: f(1)]
+
+assert.deepEqual(results, [1,2,3])

test/fixtures/enhanced-comprehensions/from.js

@@ -1,4 +1,4 @@
 arr = [4, 5, 6]
 
-c = [ for idx i, elem x in arr: [i, x] ]
+c = [ ...for idx i, elem x in arr: [i, x] ]
 assert.deepEqual(c, [[0, 4], [1, 5], [2, 6]])

test/fixtures/enhanced-comprehensions/if-elif.js

@@ -1,2 +1,2 @@
-c = [ for let i=0; i<10; i++: if i > 5: i elif i > 3: i * 2 ]
+c = [ ...for let i=0; i<10; i++: if i > 5: i elif i > 3: i * 2 ]
 assert.deepEqual(c, [8, 10, 6, 7, 8, 9])

test/fixtures/enhanced-comprehensions/if-else-if.js

@@ -1,2 +1,2 @@
-c = [ for let i=0; i<10; i++: if i > 5: i else if i > 3: i * 2 ]
+c = [ ...for let i=0; i<10; i++: if i > 5: i else if i > 3: i * 2 ]
 assert.deepEqual(c, [8, 10, 6, 7, 8, 9])

test/fixtures/enhanced-comprehensions/if-else.js

@@ -1,2 +1,2 @@
-c = [ for let i=0; i<10; i++: if i > 5: i else: 0 ]
+c = [ ...for let i=0; i<10; i++: if i > 5: i else: 0 ]
 assert.deepEqual(c, [0, 0, 0, 0, 0, 0, 6, 7, 8, 9])

test/fixtures/enhanced-comprehensions/if.js

@@ -1,2 +1,2 @@
-c = [ for let i=0; i<10; i++: if i > 5: i ]
+c = [ ...for let i=0; i<10; i++: if i > 5: i ]
 assert.deepEqual(c, [6, 7, 8, 9])

test/fixtures/enhanced-comprehensions/in.js

@@ -1,7 +1,7 @@
 arr = [4, 5, 6]
 
-x = [ for const i in arr: parseInt(i) ]
+x = [ ...for const i in arr: parseInt(i) ]
 assert.deepEqual(x, [0, 1, 2])
 
-y = [ for const i in arr: parseInt(i) + 1 ]
+y = [ ...for const i in arr: parseInt(i) + 1 ]
 assert.deepEqual(y, [1, 2, 3])

test/fixtures/enhanced-comprehensions/multi/actual.js

@@ -1,10 +1,10 @@
 [
   1
   2
-  for elem e in ["buckle", "my", "shoe"]: e
-  case true: 3
+  ...for elem e in ["buckle", "my", "shoe"]: e
+  ...if true: 3
   4
-  for elem e in ["shut", "the", "door"]: e
+  ...for elem e in ["shut", "the", "door"]: e
   5
   6
   "pickup sticks"

test/fixtures/enhanced-comprehensions/multiple-for.js

@@ -1,5 +1,5 @@
-c = [ for let i=0;i<3;i++: for let j=5;j<7;j++: [i, j] ]
+c = [ ...for let i=0;i<3;i++: for let j=5;j<7;j++: [i, j] ]
 assert.deepEqual(c, [ [0, 5], [0, 6], [1, 5], [1, 6], [2, 5], [2, 6] ])
 
-d = [ for let i=0;i<3;i++: for let j=5;j<7;j++: if i > 1: [i, j] ]
+d = [ ...for let i=0;i<3;i++: for let j=5;j<7;j++: if i > 1: [i, j] ]
 assert.deepEqual(d, [ [2, 5], [2, 6] ])

test/fixtures/enhanced-comprehensions/nested/actual.js

@@ -0,0 +1 @@
+[...for elem e in x: [...for elem d in y: [d]]]

test/fixtures/enhanced-comprehensions/nested.js → test/fixtures/enhanced-comprehensions/nested/exec.js

@@ -1,2 +1,2 @@
-c = [ for let i=0; i<4; i++: [ for let i=0; i<3; i++: 2 ] ]
+c = [ ...for let i=0; i<4; i++: [ ...for let i=0; i<3; i++: 2 ] ]
 assert.deepEqual(c, [ [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2] ])