further steps for closure support

7 files changed, 101 insertions, 57 deletions

diff --git a/compiler/cgen.nim b/compiler/cgen.nim
index 9784b21bb..e9d7673bb 100755
--- a/compiler/cgen.nim
+++ b/compiler/cgen.nim
@@ -578,6 +578,7 @@ proc closureSetup(p: BProc, prc: PSym) =
   if prc.typ.callConv != ccClosure: return
   # prc.ast[paramsPos].last contains the type we're after:
   var env = lastSon(prc.ast[paramsPos]).sym
+  #echo "created environment: ", env.id, " for ", prc.name.s
   assignLocalVar(p, env)
   # generate cast assignment:
   appcg(p, cpsStmts, "$1 = ($2) ClEnv;$n", rdLoc(env.loc), 
diff --git a/compiler/lambdalifting.nim b/compiler/lambdalifting.nim
index 70db9334e..83885029a 100644
--- a/compiler/lambdalifting.nim
+++ b/compiler/lambdalifting.nim
@@ -10,8 +10,9 @@
 # This include file implements lambda lifting for the transformator.
 
 const
-  procDefs = {nkLambda, nkProcDef, nkMethodDef, nkIteratorDef, nkMacroDef,
+  declarativeDefs = {nkProcDef, nkMethodDef, nkIteratorDef, nkMacroDef,
      nkConverterDef}
+  procDefs = {nkLambda} + declarativeDefs
 
 proc indirectAccess(a, b: PSym, info: TLineInfo): PNode = 
   # returns a[].b as a node
@@ -47,17 +48,16 @@ proc captureToTuple(cap: TCapture, owner: PSym): PType =
     addSon(result.n, newSymNode(field))
     addSon(result, typ)
 
+proc interestingVar(s: PSym): bool {.inline.} =
+  result = s.kind in {skVar, skLet, skTemp, skForVar, skParam, skResult} and
+    sfGlobal notin s.flags
+
 proc gatherVars(c: PTransf, n: PNode, outerProc: PSym, cap: var TCapture) = 
   # gather used vars for closure generation into 'cap'
   case n.kind
   of nkSym:
     var s = n.sym
-    var found = false
-    case s.kind
-    of skVar, skLet: found = sfGlobal notin s.flags
-    of skTemp, skForVar, skParam, skResult: found = true
-    else: nil
-    if found and outerProc.id == s.owner.id:
+    if interestingVar(s) and outerProc.id == s.owner.id:
       #echo "captured: ", s.name.s
       Capture(cap, s)
   of nkEmpty..pred(nkSym), succ(nkSym)..nkNilLit: nil
@@ -70,37 +70,34 @@ proc replaceVars(c: PTransf, n: PNode, outerProc, env: PSym) =
     let a = n.sons[i]
     if a.kind == nkSym:
       let s = a.sym
-      var found = false
-      case s.kind
-      of skVar, skLet: found = sfGlobal notin s.flags
-      of skTemp, skForVar, skParam, skResult: found = true
-      else: nil
-      if found and outerProc.id == s.owner.id:
+      if interestingVar(s) and outerProc == s.owner:
         # access through the closure param:
         n.sons[i] = indirectAccess(env, s, n.info)
     else:
       replaceVars(c, a, outerProc, env)
 
-proc addFormalParam(routine: PType, param: PSym) =
-  addSon(routine, param.typ)
-  addSon(routine.n, newSymNode(param))
-
-proc addFormalParam(routine: PSym, param: PSym) = 
-  #addFormalParam(routine.typ, param)
-  addSon(routine.ast.sons[paramsPos], newSymNode(param))
+proc addHiddenParam(routine: PSym, param: PSym) =
+  var params = routine.ast.sons[paramsPos]
+  let L = params.len-1
+  if L >= 0:
+    # update if we already added a hidden parameter:
+    if params.sons[L].kind == nkSym and params.sons[L].sym.kind == skTemp: 
+      params.sons[L].sym = param
+      return
+  addSon(params, newSymNode(param))
+  #echo "produced environment: ", param.id, " for ", routine.name.s
 
 proc isInnerProc(s, outerProc: PSym): bool {.inline.} =
   result = s.kind in {skProc, skMacro, skIterator, skMethod, skConverter} and
-    s.owner.id == outerProc.id and not isGenericRoutine(s) and
-    s.typ.callConv == ccClosure
+    s.owner == outerProc and not isGenericRoutine(s)
+  #s.typ.callConv == ccClosure
 
 proc searchForInnerProcs(c: PTransf, n: PNode, outerProc: PSym,
                          cap: var TCapture) =
   case n.kind
   of nkSym:
-    let s = n.sym
-    if isInnerProc(s, outerProc):
-      gatherVars(c, s.getBody, outerProc, cap)
+    if isInnerProc(n.sym, outerProc):
+      gatherVars(c, n.sym.getBody, outerProc, cap)
   of nkEmpty..pred(nkSym), succ(nkSym)..nkNilLit: nil
   else:
     for i in 0.. <len(n):
@@ -109,29 +106,47 @@ proc searchForInnerProcs(c: PTransf, n: PNode, outerProc: PSym,
 proc makeClosure(c: PTransf, prc, env: PSym, info: TLineInfo): PNode =
   result = newNodeIT(nkClosure, info, prc.typ)
   result.add(newSymNode(prc))
-  result.add(newSymNode(env))
+  if env == nil:
+    result.add(newNodeIT(nkNilLit, info, getSysType(tyNil)))
+  else:
+    result.add(newSymNode(env))
   
 proc transformInnerProcs(c: PTransf, n: PNode, outerProc, env: PSym) =
   case n.kind
   of nkSym:
     let innerProc = n.sym
-    if isInnerProc(innerProc, outerProc):
-      # inner proc could capture outer vars:
-      var param = newTemp(c, env.typ, n.info)
-      param.kind = skParam
-      addFormalParam(innerProc, param)
-      # 'anon' should be replaced by '(anon, env)':
-      IdNodeTablePut(c.transCon.mapping, innerProc, 
-                     makeClosure(c, innerProc, env, n.info))
-      # access all non-local vars through the 'env' param:
-      var body = innerProc.getBody
-      # XXX does not work with recursion!
-      replaceVars(c, body, outerProc, param)
-      innerProc.ast.sons[bodyPos] = body
+    if isInnerProc(innerProc, outerProc) and not 
+        containsOrIncl(c.transformedInnerProcs, innerProc.id):
+      if env == nil:
+        innerProc.ast.sons[bodyPos] = transform(c, innerProc.getBody).pnode
+      else:
+        # inner proc could capture outer vars:
+        var param = newTemp(c, env.typ, n.info)
+        
+        # recursive calls go through (f, hiddenParam):
+        IdNodeTablePut(c.transCon.mapping, innerProc, 
+                       makeClosure(c, innerProc, param, n.info))
+        # access all non-local vars through the 'env' param:
+        replaceVars(c, innerProc.getBody, outerProc, param)
+
+        innerProc.ast.sons[bodyPos] = transform(c, innerProc.getBody).pnode
+        addHiddenParam(innerProc, param)
+        
+        # 'anon' should be replaced by '(anon, env)' in the outer proc:
+        IdNodeTablePut(c.transCon.mapping, innerProc, 
+                       makeClosure(c, innerProc, env, n.info))
   of nkEmpty..pred(nkSym), succ(nkSym)..nkNilLit: nil
   else:
     for i in 0.. <len(n):
       transformInnerProcs(c, n.sons[i], outerProc, env)
+  
+template checkInvariant(n: PNode, s: PSym) =
+  when false:
+    if s.ast != n:
+      echo renderTree(s.ast)
+      echo " -------------- "
+      echo n.renderTree
+    assert s.ast == n
 
 proc newCall(a, b: PSym): PNode =
   result = newNodeI(nkCall, a.info)
@@ -158,10 +173,13 @@ proc createEnvStmt(c: PTransf, varList: TCapture, env: PSym): PTransNode =
     IdNodeTablePut(c.transCon.mapping, v, fieldAccess)
   
 proc transformProcFin(c: PTransf, n: PNode, s: PSym): PTransNode =
-  # to be safe: XXX this a mystery how it could ever happen that: s.ast != n.
-  s.ast.sons[bodyPos] = n.sons[bodyPos]
-  if n.kind == nkMethodDef: methodDef(s, false)
+  if n.kind == nkLambda:
+    # for lambdas we transformed 'n.sons[bodyPos]', but not 'ast.n[bodyPos]'!
+    s.ast.sons[bodyPos] = n.sons[bodyPos]
+  else:
+    assert s.ast == n
   
+  if n.kind == nkMethodDef: methodDef(s, false)
   # should 's' be replaced by a tuple ('s', env)?
   var tc = c.transCon
   var repl: PNode = nil
@@ -181,26 +199,35 @@ proc transformProc(c: PTransf, n: PNode): PTransNode =
   
   var s = n.sons[namePos].sym
   var body = s.getBody
-  if body.kind == nkEmpty:
-    return PTransNode(n)    
+  if body.kind == nkEmpty or n.sons[bodyPos].kind == nkEmpty or
+     containsOrIncl(c.transformedInnerProcs, s.id):
+    return PTransNode(n)
+    
+  checkInvariant(n, s)
   
-  if not containsNode(body, procDefs):
+  if not containsNode(body, procDefs) and s.typ.callConv != ccClosure:
     # fast path: no inner procs, so no closure needed:
     n.sons[bodyPos] = PNode(transform(c, body))
+    checkInvariant(n, s)
     return transformProcFin(c, n, s)
 
   # create environment:
   var cap: TCapture = @[]
   searchForInnerProcs(c, body, s, cap)
+
+  var envType = newType(tyRef, s)
+  addSon(envType, captureToTuple(cap, s))
+  if s.typ.callConv == ccClosure:
+    addHiddenParam(s, newTemp(c, envType, n.info))
+    IdNodeTablePut(c.transCon.mapping, s, 
+                   makeClosure(c, s, nil, n.info))
   
   if cap.len == 0:
     # fast path: no captured variables, so no closure needed:
+    transformInnerProcs(c, body, s, nil)
     n.sons[bodyPos] = PNode(transform(c, body))
     return transformProcFin(c, n, s)
   
-  var envType = newType(tyRef, s)
-  addSon(envType, captureToTuple(cap, s))
-  
   # Currently we always do a heap allocation. A simple escape analysis
   # could turn the closure into a stack allocation. Later versions might 
   # implement that. This would require backend changes too though.
@@ -211,11 +238,11 @@ proc transformProc(c: PTransf, n: PNode): PTransNode =
   # mapping entries that turn (localProc) into (localProc, env):
   transformInnerProcs(c, body, s, envSym)
 
-  # now we can transform 'body' as all rewriting entries have been created.
-  # Careful this transforms the inner procs too!
+  # now we can transform 'body' as all rewriting entries have been created:
   newBody.add(transform(c, body))
   n.sons[bodyPos] = newBody.pnode
   result = transformProcFin(c, n, s)
+  checkInvariant(n, s)
 
 proc generateThunk(c: PTransf, prc: PNode, dest: PType): PNode =
   ## Converts 'prc' into '(thunk, nil)' so that it's compatible with
diff --git a/compiler/msgs.nim b/compiler/msgs.nim
index 43ab7a192..3862c5751 100755
--- a/compiler/msgs.nim
+++ b/compiler/msgs.nim
@@ -94,6 +94,7 @@ type
     errUnhandledExceptionX, errCyclicTree, errXisNoMacroOrTemplate, 
     errXhasSideEffects, errIteratorExpected, errLetNeedsInit,
     errThreadvarCannotInit, errWrongSymbolX, errIllegalCaptureX,
+    errXCannotBeClosure,
     errUser,
     warnCannotOpenFile, 
     warnOctalEscape, warnXIsNeverRead, warnXmightNotBeenInit, 
@@ -326,6 +327,7 @@ const
     errThreadvarCannotInit: "a thread var cannot be initialized explicitly",
     errWrongSymbolX: "usage of \'$1\' is a user-defined error", 
     errIllegalCaptureX: "illegal capture '$1'",
+    errXCannotBeClosure: "'$1' cannot have 'closure' calling convention",
     errUser: "$1", 
     warnCannotOpenFile: "cannot open \'$1\' [CannotOpenFile]",
     warnOctalEscape: "octal escape sequences do not exist; leading zero is ignored [OctalEscape]", 
diff --git a/compiler/semexprs.nim b/compiler/semexprs.nim
index 24161e85e..a90c475e2 100755
--- a/compiler/semexprs.nim
+++ b/compiler/semexprs.nim
@@ -92,7 +92,8 @@ proc semSym(c: PContext, n: PNode, s: PSym, flags: TExprFlags): PNode =
     # if a proc accesses a global variable, it is not side effect free:
     if sfGlobal in s.flags:
       incl(c.p.owner.flags, sfSideEffect)
-    elif s.owner != c.p.owner and s.owner.kind != skModule:
+    elif s.owner != c.p.owner and s.owner.kind != skModule and 
+        c.p.owner.typ != nil and not IsGenericRoutine(s.owner):
       c.p.owner.typ.callConv = ccClosure
       if illegalCapture(s) or c.p.next.owner != s.owner:
         # Currently captures are restricted to a single level of nesting:
diff --git a/compiler/semstmts.nim b/compiler/semstmts.nim
index 424950056..f761fd454 100755
--- a/compiler/semstmts.nim
+++ b/compiler/semstmts.nim
@@ -658,6 +658,8 @@ proc semLambda(c: PContext, n: PNode): PNode =
   else: 
     LocalError(n.info, errImplOfXexpected, s.name.s)
   sideEffectsCheck(c, s)
+  if s.typ.callConv == ccClosure and s.owner.kind == skModule:
+    localError(s.info, errXCannotBeClosure, s.name.s)
   closeScope(c.tab)           # close scope for parameters
   popOwner()
   result.typ = s.typ
@@ -754,6 +756,8 @@ proc semProcAux(c: PContext, n: PNode, kind: TSymKind,
       incl(s.flags, sfForward)
     elif sfBorrow in s.flags: semBorrow(c, n, s)
   sideEffectsCheck(c, s)
+  if s.typ.callConv == ccClosure and s.owner.kind == skModule:
+    localError(s.info, errXCannotBeClosure, s.name.s)
   closeScope(c.tab)           # close scope for parameters
   popOwner()
   
diff --git a/compiler/transf.nim b/compiler/transf.nim
index b5e875c53..509aa9320 100755
--- a/compiler/transf.nim
+++ b/compiler/transf.nim
@@ -45,9 +45,9 @@ type
     module: PSym
     transCon: PTransCon      # top of a TransCon stack
     inlining: int            # > 0 if we are in inlining context (copy vars)
+    nestedProcs: int         # > 0 if we are in a nested proc
     blocksyms: seq[PSym]
-    procToEnv: TIdTable      # mapping from a proc to its generated explicit
-                             # 'env' var (for closure generation)
+    transformedInnerProcs: TIntSet
   PTransf = ref TTransfContext
 
 proc newTransNode(a: PNode): PTransNode {.inline.} = 
@@ -616,7 +616,16 @@ proc transform(c: PTransf, n: PNode): PTransNode =
     # nothing to be done for leaves:
     result = PTransNode(n)
   of nkBracketExpr: result = transformArrayAccess(c, n)
-  of procDefs: result = transformProc(c, n)
+  of procDefs: 
+    if c.nestedProcs == 0:
+      inc c.nestedProcs
+      result = transformProc(c, n)
+      dec c.nestedProcs
+    else:
+      result = PTransNode(n)
+      if n.sons[namePos].kind == nkSym:
+        let x = transformSym(c, n.sons[namePos])
+        if x.pnode.kind == nkClosure: result = x
   of nkForStmt: result = transformFor(c, n)
   of nkCaseStmt: result = transformCase(c, n)
   of nkContinueStmt:
@@ -680,7 +689,7 @@ proc openTransf(module: PSym, filename: string): PPassContext =
   new(n)
   n.blocksyms = @[]
   n.module = module
-  initIdTable(n.procToEnv)
+  n.transformedInnerProcs = initIntSet()
   result = n
 
 proc openTransfCached(module: PSym, filename: string, 
diff --git a/todo.txt b/todo.txt
index 7a5ec6054..1bf562db7 100755
--- a/todo.txt
+++ b/todo.txt
@@ -2,9 +2,9 @@ version 0.8.14
 ==============
 
 - implement closures
-  - test evals.nim with closures
+  - fix evals.nim with closures
   - deactivate lambda lifting for JS backend
-  - Test capture of for loop vars; test generics; test recursion
+  - Test capture of for loop vars; test generics;
   - test constant closures
   - 'closureEnv' magic for easy interfacing with C