Iterator fixes (#11234)

8 files changed, 142 insertions, 17 deletions

diff --git a/compiler/ast.nim b/compiler/ast.nim
index ad2928521..604e9d728 100644
--- a/compiler/ast.nim
+++ b/compiler/ast.nim
@@ -1813,6 +1813,9 @@ template detailedInfo*(sym: PSym): string =
 proc isInlineIterator*(s: PSym): bool {.inline.} =
   s.kind == skIterator and s.typ.callConv != ccClosure
 
+proc isClosureIterator*(s: PSym): bool {.inline.} =
+  s.kind == skIterator and s.typ.callConv == ccClosure
+
 proc isSinkParam*(s: PSym): bool {.inline.} =
   s.kind == skParam and (s.typ.kind == tySink or tfHasOwned in s.typ.flags)
 
diff --git a/compiler/semexprs.nim b/compiler/semexprs.nim
index 44d7cfffd..91da4a1ec 100644
--- a/compiler/semexprs.nim
+++ b/compiler/semexprs.nim
@@ -1696,8 +1696,8 @@ proc semAsgn(c: PContext, n: PNode; mode=asgnNormal): PNode =
 proc semReturn(c: PContext, n: PNode): PNode =
   result = n
   checkSonsLen(n, 1, c.config)
-  if c.p.owner.kind in {skConverter, skMethod, skProc, skFunc, skMacro} or (
-     c.p.owner.kind == skIterator and c.p.owner.typ.callConv == ccClosure):
+  if c.p.owner.kind in {skConverter, skMethod, skProc, skFunc, skMacro} or
+      isClosureIterator(c.p.owner):
     if n.sons[0].kind != nkEmpty:
       # transform ``return expr`` to ``result = expr; return``
       if c.p.resultSym != nil:
@@ -1742,8 +1742,12 @@ proc semProcBody(c: PContext, n: PNode): PNode =
       c.p.resultSym.typ = errorType(c)
       c.p.owner.typ.sons[0] = nil
     else:
-      localError(c.config, c.p.resultSym.info, errCannotInferReturnType)
-
+      localError(c.config, c.p.resultSym.info, errCannotInferReturnType %
+        c.p.owner.name.s)
+  if isInlineIterator(c.p.owner) and c.p.owner.typ.sons[0] != nil and
+      c.p.owner.typ.sons[0].kind == tyUntyped:
+    localError(c.config, c.p.owner.info, errCannotInferReturnType %
+      c.p.owner.name.s)
   closeScope(c)
 
 proc semYieldVarResult(c: PContext, n: PNode, restype: PType) =
@@ -1786,6 +1790,8 @@ proc semYield(c: PContext, n: PNode): PNode =
       if resultTypeIsInferrable(restype):
         let inferred = n.sons[0].typ
         iterType.sons[0] = inferred
+        if c.p.resultSym != nil:
+          c.p.resultSym.typ = inferred
 
       semYieldVarResult(c, n, restype)
     else:
diff --git a/compiler/semstmts.nim b/compiler/semstmts.nim
index c84fbd14e..6423b8fb3 100644
--- a/compiler/semstmts.nim
+++ b/compiler/semstmts.nim
@@ -27,7 +27,7 @@ const
     "it is used as an operand to another routine and the types " &
     "of the generic paramers can be inferred from the expected signature."
   errCannotInferTypeOfTheLiteral = "cannot infer the type of the $1"
-  errCannotInferReturnType = "cannot infer the return type of the proc"
+  errCannotInferReturnType = "cannot infer the return type of '$1'"
   errCannotInferStaticParam = "cannot infer the value of the static param '$1'"
   errProcHasNoConcreteType = "'$1' doesn't have a concrete type, due to unspecified generic parameters."
   errLetNeedsInit = "'let' symbol requires an initialization"
@@ -35,6 +35,9 @@ const
   errImplOfXexpected = "implementation of '$1' expected"
   errRecursiveDependencyX = "recursive dependency: '$1'"
   errPragmaOnlyInHeaderOfProcX = "pragmas are only allowed in the header of a proc; redefinition of $1"
+  errCannotAssignMacroSymbol = "cannot assign macro symbol to $1 here. Forgot to invoke the macro with '()'?"
+  errInvalidTypeDescAssign = "'typedesc' metatype is not valid here; typed '=' instead of ':'?"
+  errInlineIteratorNotFirstClass = "inline iterators are not first-class / cannot be assigned to variables"
 
 proc semDiscard(c: PContext, n: PNode): PNode =
   result = n
@@ -445,12 +448,16 @@ proc semVarOrLet(c: PContext, n: PNode, symkind: TSymKind): PNode =
     var def: PNode = c.graph.emptyNode
     if a.sons[length-1].kind != nkEmpty:
       def = semExprWithType(c, a.sons[length-1], {efAllowDestructor})
-      if def.typ.kind == tyProc and def.kind == nkSym and def.sym.kind == skMacro:
-        localError(c.config, def.info, "cannot assign macro symbol to variable here. Forgot to invoke the macro with '()'?")
-        def.typ = errorType(c)
+      if def.typ.kind == tyProc and def.kind == nkSym:
+        if def.sym.kind == skMacro:
+          localError(c.config, def.info, errCannotAssignMacroSymbol % "variable")
+          def.typ = errorType(c)
+        elif isInlineIterator(def.sym):
+          localError(c.config, def.info, errInlineIteratorNotFirstClass)
+          def.typ = errorType(c)
       elif def.typ.kind == tyTypeDesc and c.p.owner.kind != skMacro:
         # prevent the all too common 'var x = int' bug:
-        localError(c.config, def.info, "'typedesc' metatype is not valid here; typed '=' instead of ':'?")
+        localError(c.config, def.info, errInvalidTypeDescAssign)
         def.typ = errorType(c)
 
       if typ != nil:
@@ -589,12 +596,16 @@ proc semConst(c: PContext, n: PNode): PNode =
       localError(c.config, a.sons[length-1].info, errConstExprExpected)
       continue
 
-    if def.typ.kind == tyProc and def.kind == nkSym and def.sym.kind == skMacro:
-        localError(c.config, def.info, "cannot assign macro symbol to constant here. Forgot to invoke the macro with '()'?")
+    if def.typ.kind == tyProc and def.kind == nkSym:
+      if def.sym.kind == skMacro:
+        localError(c.config, def.info, errCannotAssignMacroSymbol % "constant")
+        def.typ = errorType(c)
+      elif isInlineIterator(def.sym):
+        localError(c.config, def.info, errInlineIteratorNotFirstClass)
         def.typ = errorType(c)
     elif def.typ.kind == tyTypeDesc and c.p.owner.kind != skMacro:
       # prevent the all too common 'const x = int' bug:
-      localError(c.config, def.info, "'typedesc' metatype is not valid here; typed '=' instead of ':'?")
+      localError(c.config, def.info, errInvalidTypeDescAssign)
       def.typ = errorType(c)
 
     # check type compatibility between def.typ and typ:
@@ -1534,8 +1545,7 @@ proc activate(c: PContext, n: PNode) =
       discard
 
 proc maybeAddResult(c: PContext, s: PSym, n: PNode) =
-  if s.typ.sons[0] != nil and not
-      (s.kind == skIterator and s.typ.callConv != ccClosure):
+  if s.typ.sons[0] != nil and not isInlineIterator(s):
     addResult(c, s.typ.sons[0], n.info, s.kind)
     addResultNode(c, n)
 
@@ -1768,8 +1778,9 @@ proc semProcAux(c: PContext, n: PNode, kind: TSymKind,
   elif s.kind == skFunc:
     incl(s.flags, sfNoSideEffect)
     incl(s.typ.flags, tfNoSideEffect)
-  var proto = searchForProc(c, oldScope, s)
-  if proto == nil or isAnon:
+  var proto: PSym = if isAnon: nil
+                    else: searchForProc(c, oldScope, s)
+  if proto == nil:
     if s.kind == skIterator:
       if s.typ.callConv != ccClosure:
         s.typ.callConv = if isAnon: ccClosure else: ccInline
@@ -1921,7 +1932,7 @@ proc semIterator(c: PContext, n: PNode): PNode =
   if t.sons[0] == nil and s.typ.callConv != ccClosure:
     localError(c.config, n.info, "iterator needs a return type")
   if isAnon and s.typ.callConv == ccInline:
-    localError(c.config, n.info, "inline iterators are not first-class / cannot be assigned to variables")
+    localError(c.config, n.info, errInlineIteratorNotFirstClass)
   # iterators are either 'inline' or 'closure'; for backwards compatibility,
   # we require first class iterators to be marked with 'closure' explicitly
   # -- at least for 0.9.2.
diff --git a/tests/iter/titer13.nim b/tests/iter/titer13.nim
new file mode 100644
index 000000000..716f59900
--- /dev/null
+++ b/tests/iter/titer13.nim
@@ -0,0 +1,72 @@
+discard """
+  output: '''b yields
+c yields
+a returns
+b yields
+b returns
+c yields
+
+
+1
+2
+3
+4
+'''
+"""
+
+block:
+  template tloop(iter: untyped) =
+    for i in iter():
+      echo i
+
+  template twhile(iter: untyped) =
+    let it = iter
+    while not finished(it):
+      echo it()
+
+  iterator a(): auto {.closure.} =
+    if true: return "a returns"
+    yield "a yields"
+
+  iterator b(): auto {.closure.} =
+    yield "b yields"
+    if true: return "b returns"
+
+  iterator c(): auto {.closure.} =
+    yield "c yields"
+    if true: return
+
+  iterator d(): auto {.closure.} =
+    if true: return
+    yield "d yields"
+
+  tloop(a)
+  tloop(b)
+  tloop(c)
+  tloop(d)
+  twhile(a)
+  twhile(b)
+  twhile(c)
+  twhile(d)
+
+block:
+  iterator a: auto =
+    yield 1
+  for x in a():
+    echo x
+
+  let b = iterator: int =
+    yield 2
+  for x in b():
+    echo x
+
+  let c = iterator: auto =
+    yield 3
+  for x in c():
+    echo x
+
+block:
+  iterator myIter2(): auto {.closure.} =
+    yield 4
+  for a in myIter2():
+    echo a
diff --git a/tests/iter/titerassignerr.nim b/tests/iter/titerassignerr.nim
new file mode 100644
index 000000000..caa56c4ad
--- /dev/null
+++ b/tests/iter/titerassignerr.nim
@@ -0,0 +1,8 @@
+discard """
+  errormsg: "inline iterators are not first-class / cannot be assigned to variables"
+  line: 8
+"""
+
+iterator foo: int =
+  yield 2
+let x = foo
diff --git a/tests/iter/titerautoerr1.nim b/tests/iter/titerautoerr1.nim
new file mode 100644
index 000000000..c7e5642c8
--- /dev/null
+++ b/tests/iter/titerautoerr1.nim
@@ -0,0 +1,8 @@
+discard """
+  errormsg: "type mismatch: got <int literal(1)> but expected 'string'"
+  line: 8
+"""
+
+iterator a(): auto {.closure.} =
+  if true: return "str"
+  yield 1
diff --git a/tests/iter/titerautoerr2.nim b/tests/iter/titerautoerr2.nim
new file mode 100644
index 000000000..eb48a36f9
--- /dev/null
+++ b/tests/iter/titerautoerr2.nim
@@ -0,0 +1,8 @@
+discard """
+  errormsg: "type mismatch: got <string> but expected 'int literal(1)'"
+  line: 8
+"""
+
+iterator b(): auto {.closure.} =
+  yield 1
+  if true: return "str"
diff --git a/tests/iter/titerautoerr3.nim b/tests/iter/titerautoerr3.nim
new file mode 100644
index 000000000..e47b65540
--- /dev/null
+++ b/tests/iter/titerautoerr3.nim
@@ -0,0 +1,9 @@
+discard """
+  errormsg: "cannot infer the return type of 'bar'"
+  line: 6
+"""
+
+iterator bar(): auto =
+  discard
+for t in bar():
+  discard