1 files changed, 59 insertions, 33 deletions

diff --git a/lib/pure/sugar.nim b/lib/pure/sugar.nim
index 062341ccb..dbc60b8a9 100644
--- a/lib/pure/sugar.nim
+++ b/lib/pure/sugar.nim
@@ -11,16 +11,25 @@
 ## macro system.
 
 import std/private/since
-import macros
-import typetraits
+import macros, typetraits
+
+proc checkPragma(ex, prag: var NimNode) =
+  since (1, 3):
+    if ex.kind == nnkPragmaExpr:
+      prag = ex[1]
+      if ex[0].kind == nnkPar and ex[0].len == 1:
+        ex = ex[0][0]
+      else:
+        ex = ex[0]
 
 proc createProcType(p, b: NimNode): NimNode {.compileTime.} =
-  #echo treeRepr(p)
-  #echo treeRepr(b)
   result = newNimNode(nnkProcTy)
-  var formalParams = newNimNode(nnkFormalParams)
+  var
+    formalParams = newNimNode(nnkFormalParams).add(b)
+    p = p
+    prag = newEmptyNode()
 
-  formalParams.add b
+  checkPragma(p, prag)
 
   case p.kind
   of nnkPar, nnkTupleConstr:
@@ -44,9 +53,7 @@ proc createProcType(p, b: NimNode): NimNode {.compileTime.} =
     formalParams.add identDefs
 
   result.add formalParams
-  result.add newEmptyNode()
-  #echo(treeRepr(result))
-  #echo(result.toStrLit())
+  result.add prag
 
 macro `=>`*(p, b: untyped): untyped =
   ## Syntax sugar for anonymous procedures.
@@ -58,56 +65,75 @@ macro `=>`*(p, b: untyped): untyped =
   ##
   ##   passTwoAndTwo((x, y) => x + y) # 4
 
-  #echo treeRepr(p)
-  #echo(treeRepr(b))
-  var params: seq[NimNode] = @[newIdentNode("auto")]
+  var
+    params = @[ident"auto"]
+    name = newEmptyNode()
+    kind = nnkLambda
+    pragma = newEmptyNode()
+    p = p
+
+  checkPragma(p, pragma)
+
+  if p.kind == nnkInfix and p[0].kind == nnkIdent and p[0].eqIdent"->":
+    params[0] = p[2]
+    p = p[1]
+
+  checkPragma(p, pragma) # check again after -> transform
+
+  since (1, 3):
+    if p.kind == nnkCall:
+      # foo(x, y) => x + y
+      kind = nnkProcDef
+      name = p[0]
+      let newP = newNimNode(nnkPar)
+      for i in 1..<p.len:
+        newP.add(p[i])
+      p = newP
 
   case p.kind
   of nnkPar, nnkTupleConstr:
-    for c in children(p):
+    var untypedBeforeColon = 0
+    for i, c in p:
       var identDefs = newNimNode(nnkIdentDefs)
       case c.kind
       of nnkExprColonExpr:
+        let t = c[1]
+        since (1, 3):
+          # + 1 here because of return type in params
+          for j in (i - untypedBeforeColon + 1) .. i:
+            params[j][1] = t
+        untypedBeforeColon = 0
         identDefs.add(c[0])
-        identDefs.add(c[1])
+        identDefs.add(t)
         identDefs.add(newEmptyNode())
       of nnkIdent:
         identDefs.add(c)
         identDefs.add(newIdentNode("auto"))
         identDefs.add(newEmptyNode())
+        inc untypedBeforeColon
       of nnkInfix:
-        if c[0].kind == nnkIdent and c[0].ident == !"->":
+        if c[0].kind == nnkIdent and c[0].eqIdent"->":
           var procTy = createProcType(c[1], c[2])
           params[0] = procTy[0][0]
           for i in 1 ..< procTy[0].len:
             params.add(procTy[0][i])
         else:
-          error("Expected proc type (->) got (" & $c[0].ident & ").")
+          error("Expected proc type (->) got (" & c[0].strVal & ").", c)
         break
       else:
-        echo treeRepr c
-        error("Incorrect procedure parameter list.")
+        error("Incorrect procedure parameter list.", c)
       params.add(identDefs)
   of nnkIdent:
     var identDefs = newNimNode(nnkIdentDefs)
     identDefs.add(p)
-    identDefs.add(newIdentNode("auto"))
+    identDefs.add(ident"auto")
     identDefs.add(newEmptyNode())
     params.add(identDefs)
-  of nnkInfix:
-    if p[0].kind == nnkIdent and p[0].ident == !"->":
-      var procTy = createProcType(p[1], p[2])
-      params[0] = procTy[0][0]
-      for i in 1 ..< procTy[0].len:
-        params.add(procTy[0][i])
-    else:
-      error("Expected proc type (->) got (" & $p[0].ident & ").")
   else:
-    error("Incorrect procedure parameter list.")
-  result = newProc(params = params, body = b, procType = nnkLambda)
-  #echo(result.treeRepr)
-  #echo(result.toStrLit())
-  #return result # TODO: Bug?
+    error("Incorrect procedure parameter list.", p)
+  result = newProc(body = b, params = params,
+                   pragmas = pragma, name = name,
+                   procType = kind)
 
 macro `->`*(p, b: untyped): untyped =
   ## Syntax sugar for procedure types.
@@ -190,7 +216,7 @@ macro capture*(locals: varargs[typed], body: untyped): untyped {.since: (1, 1).}
   result.add(newProc(newEmptyNode(), params, body, nnkProcDef))
   for arg in locals: result.add(arg)
 
-when (NimMajor, NimMinor) >= (1, 1):
+since (1, 1):
   import std / private / underscored_calls
 
   macro dup*[T](arg: T, calls: varargs[untyped]): T =