'return' for first class iterators

6 files changed, 80 insertions, 21 deletions

diff --git a/compiler/lambdalifting.nim b/compiler/lambdalifting.nim
index c0098fb17..309e01c6f 100644
--- a/compiler/lambdalifting.nim
+++ b/compiler/lambdalifting.nim
@@ -594,9 +594,11 @@ type
     tup: PType
 
 proc newIterResult(iter: PSym): PSym =
-  result = newSym(skResult, getIdent":result", iter, iter.info)
-  result.typ = iter.typ.sons[0]
-  incl(result.flags, sfUsed)
+  result = iter.ast.sons[resultPos].sym
+  when false:
+    result = newSym(skResult, getIdent":result", iter, iter.info)
+    result.typ = iter.typ.sons[0]
+    incl(result.flags, sfUsed)
 
 proc interestingIterVar(s: PSym): bool {.inline.} =
   result = s.kind in {skVar, skLet, skTemp, skForVar} and sfGlobal notin s.flags
@@ -636,6 +638,13 @@ proc transfIterBody(c: var TIterContext, n: PNode): PNode =
     result.add(stateAsgnStmt)
     result.add(retStmt)
     result.add(stateLabelStmt)
+  of nkReturnStmt:
+    result = newNodeI(nkStmtList, n.info)
+    var stateAsgnStmt = newNodeI(nkAsgn, n.info)
+    stateAsgnStmt.add(indirectAccess(newSymNode(c.closureParam),c.state,n.info))
+    stateAsgnStmt.add(newIntTypeNode(nkIntLit, -1, getSysType(tyInt)))
+    result.add(stateAsgnStmt)
+    result.add(n)
   else:
     for i in countup(0, sonsLen(n)-1):
       let x = transfIterBody(c, n.sons[i])
diff --git a/compiler/semexprs.nim b/compiler/semexprs.nim
index 910100e92..e1d69c0bc 100755
--- a/compiler/semexprs.nim
+++ b/compiler/semexprs.nim
@@ -1098,23 +1098,25 @@ proc semAsgn(c: PContext, n: PNode): PNode =
     asgnToResultVar(c, n, n.sons[0], n.sons[1])
   result = n
 
-proc SemReturn(c: PContext, n: PNode): PNode = 
+proc SemReturn(c: PContext, n: PNode): PNode =
   result = n
   checkSonsLen(n, 1)
-  if c.p.owner.kind notin {skConverter, skMethod, skProc, skMacro}:
+  if c.p.owner.kind in {skConverter, skMethod, skProc, skMacro} or
+     (c.p.owner.kind == skIterator and c.p.owner.typ.callConv == ccClosure):
+    if n.sons[0].kind != nkEmpty:
+      # transform ``return expr`` to ``result = expr; return``
+      if c.p.resultSym != nil: 
+        var a = newNodeI(nkAsgn, n.sons[0].info)
+        addSon(a, newSymNode(c.p.resultSym))
+        addSon(a, n.sons[0])
+        n.sons[0] = semAsgn(c, a)
+        # optimize away ``result = result``:
+        if n[0][1].kind == nkSym and n[0][1].sym == c.p.resultSym: 
+          n.sons[0] = ast.emptyNode
+      else:
+        LocalError(n.info, errNoReturnTypeDeclared)
+  else:
     LocalError(n.info, errXNotAllowedHere, "\'return\'")
-  elif n.sons[0].kind != nkEmpty:
-    # transform ``return expr`` to ``result = expr; return``
-    if c.p.resultSym != nil: 
-      var a = newNodeI(nkAsgn, n.sons[0].info)
-      addSon(a, newSymNode(c.p.resultSym))
-      addSon(a, n.sons[0])
-      n.sons[0] = semAsgn(c, a)
-      # optimize away ``result = result``:
-      if n[0][1].kind == nkSym and n[0][1].sym == c.p.resultSym: 
-        n.sons[0] = ast.emptyNode
-    else:
-      LocalError(n.info, errNoReturnTypeDeclared)
 
 proc semProcBody(c: PContext, n: PNode): PNode =
   openScope(c.tab)
diff --git a/compiler/semstmts.nim b/compiler/semstmts.nim
index 37ffe2889..d68d0da1a 100755
--- a/compiler/semstmts.nim
+++ b/compiler/semstmts.nim
@@ -792,7 +792,8 @@ proc semProcAux(c: PContext, n: PNode, kind: TSymKind,
     if n.sons[genericParamsPos].kind == nkEmpty: 
       ParamsTypeCheck(c, s.typ)
       pushProcCon(c, s)
-      if s.typ.sons[0] != nil and kind != skIterator: 
+      if s.typ.sons[0] != nil and
+          (kind != skIterator or s.typ.callConv == ccClosure):
         addResult(c, s.typ.sons[0], n.info, kind)
         addResultNode(c, n)
       if sfImportc notin s.flags:
diff --git a/doc/manual.txt b/doc/manual.txt
index c97e4ef5c..f6ff576bf 100755
--- a/doc/manual.txt
+++ b/doc/manual.txt
@@ -1199,6 +1199,31 @@ details like this when mixing garbage collected data with unmanaged memory.
 .. XXX finalizers for traced objects
 
 
+Not nil annotation
+~~~~~~~~~~~~~~~~~~
+
+All types for that ``nil`` is a valid value can be annotated to 
+exclude ``nil`` as a valid value with the `not nil`:idx: annotation:
+
+.. code-block:: nimrod
+  type
+    PObject = ref TObj not nil
+    TProc = (proc (x, y: int)) not nil
+    
+  proc p(x: PObject) =
+    echo "not nil"
+  
+  # compiler catches this:
+  p(nil)
+  
+  # but not this:
+  var x: PObject
+  p(x)
+
+As shown in the example this is merely an annotation for documentation purposes;
+for now the compiler can only catch the most trivial type violations.
+
+
 Procedural type
 ~~~~~~~~~~~~~~~
 A `procedural type`:idx: is internally a pointer to a procedure. ``nil`` is
diff --git a/tests/run/titer9.nim b/tests/run/titer9.nim
new file mode 100644
index 000000000..0d6c466c1
--- /dev/null
+++ b/tests/run/titer9.nim
@@ -0,0 +1,20 @@
+discard """
+  output: '''5
+14
+0'''
+"""
+
+iterator count(x: int, skip: bool): int {.closure.} =
+  if skip: return x+10
+  else: yield x+1
+
+  if skip: return x+10
+  else: yield x+2
+
+proc takeProc(x: iterator (x: int, skip: bool): int) =
+  echo x(4, false)
+  echo x(4, true)
+  echo x(4, false)
+
+takeProc(count)
+
diff --git a/todo.txt b/todo.txt
index 7d8a3809d..7e940c3e4 100755
--- a/todo.txt
+++ b/todo.txt
@@ -2,7 +2,6 @@ version 0.9.2
 =============
   
 - test&finish first class iterators:
-  * allow return in first class iterators
   * nested iterators
   * arglist as a type?
 
@@ -17,7 +16,7 @@ version 0.9.2
 
 - improve the compiler as a service
 - ``=`` should be overloadable; requires specialization for ``=``
-- implement constructors and non-nil types
+- implement constructors
 - make 'bind' default for templates and introduce 'mixin'; 
   special rule for ``[]=``
 - implicit deref for parameter matching; overloading based on 'var T'
@@ -38,6 +37,7 @@ Bugs
 version 0.9.XX
 ==============
 
+- improve not-nil types
 - make:
   p(a, b): 
     echo a
@@ -144,7 +144,9 @@ Version 2 and beyond
   (or implicit) syncGC() calls in loops. Automatic loop injection seems
   troublesome, but maybe we can come up with a simple heuristic. (All procs
   that `new` shared memory are syncGC() candidates... But then 'new' itself
-  calls syncGC() so that's pointless.)
+  calls syncGC() so that's pointless.) Hm instead of an heuristic simply
+  provide a ``syncgc`` pragma to trigger compiler injection --> more general:
+  an ``injectLoop`` pragma
 
 - const ptr/ref  --> pointless because of aliasing; 
   much better: 'writes: []' effect