rewrote lambdalifting; fixes deeply nested closures

bin/empty.txt

@@ -1 +1 @@
-This file keeps several tools from deleting this subdirectory.
+This file keeps several tools from deleting this subdirectory.

compiler/ast.nim

@@ -417,6 +417,7 @@ type
                 # efficiency
     nfTransf,   # node has been transformed
     nfSem       # node has been checked for semantics
+    nfLL        # node has gone through lambda lifting
     nfDotField  # the call can use a dot operator
     nfDotSetter # the call can use a setter dot operarator
     nfExplicitCall # x.y() was used instead of x.y
@@ -1504,7 +1505,7 @@ proc isGenericRoutine*(s: PSym): bool =
 proc skipGenericOwner*(s: PSym): PSym =
   internalAssert s.kind in skProcKinds
   ## Generic instantiations are owned by their originating generic
-  ## symbol. This proc skips such owners and goes straigh to the owner
+  ## symbol. This proc skips such owners and goes straight to the owner
   ## of the generic itself (the module or the enclosing proc).
   result = if sfFromGeneric in s.flags: s.owner.owner
            else: s.owner

compiler/lambdalifting.nim

@@ -103,43 +103,50 @@ discard """
 
 """
 
+# Important things to keep in mind:
+# * Don't base the analysis on nkProcDef et al. This doesn't work for
+#   instantiated (formerly generic) procs. The analysis has to look at nkSym.
+#   This also means we need to prevent the same proc is processed multiple
+#   times via the 'processed' set.
+# * Keep in mind that the owner of some temporaries used to be unreliable.
+# * For closure iterators we merge the "real" potential closure with the
+#   local storage requirements for efficiency. This means closure iterators
+#   have slightly different semantics from ordinary closures.
+
+
 const
   upName* = ":up" # field name for the 'up' reference
-  paramName* = ":env"
+  paramName* = ":envP"
   envName* = ":env"
 
 type
-  PInnerContext = ref TInnerContext
   POuterContext = ref TOuterContext
 
+  TIter = object
+    fn, closureParam, state, resultSym: PSym # most are only valid if
+                                             # fn.kind == skClosureIterator
+    obj: PType
+    
   PEnv = ref TEnv
-  TDep = tuple[e: PEnv, field: PSym]
   TEnv {.final.} = object of TObject
-    attachedNode: PNode
+    attachedNode, replacementNode: PNode
     createdVar: PSym        # if != nil it is a used environment
     createdVarComesFromIter: bool
     capturedVars: seq[PSym] # captured variables in this environment
-    deps: seq[TDep]         # dependencies
-    up: PEnv
+    up, next: PEnv          # outer scope and next to keep all in a list
+    upField: PSym        # if != nil the dependency to the outer scope is used
     obj: PType
-  
-  TInnerContext = object
-    fn: PSym
-    closureParam: PSym
-    localsToAccess: TIdNodeTable
+    fn: PSym                # function that belongs to this scope;
+                            # if up.fn != fn then we cross function boundaries.
+                            # This is an important case to consider.
+    vars: TIntSet           # variables belonging to this environment
     
   TOuterContext = object
     fn: PSym # may also be a module!
-    currentEnv: PEnv
-    isIter: bool   # first class iterator?
+    head: PEnv
     capturedVars, processed: TIntSet
-    localsToEnv: TIdTable # PSym->PEnv mapping
     localsToAccess: TIdNodeTable
     lambdasToEnv: TIdTable # PSym->PEnv mapping
-    up: POuterContext
-
-    closureParam, state, resultSym: PSym # only if isIter
-    obj: PType # only if isIter
 
 proc getStateType(iter: PSym): PType =
   var n = newNodeI(nkRange, iter.info)
@@ -164,6 +171,7 @@ proc newIterResult(iter: PSym): PSym =
     iter.ast.add newSymNode(result)
 
 proc addHiddenParam(routine: PSym, param: PSym) =
+  assert param.kind == skParam
   var params = routine.ast.sons[paramsPos]
   # -1 is correct here as param.position is 0 based but we have at position 0
   # some nkEffect node:
@@ -175,7 +183,7 @@ proc addHiddenParam(routine: PSym, param: PSym) =
 proc getHiddenParam(routine: PSym): PSym =
   let params = routine.ast.sons[paramsPos]
   let hidden = lastSon(params)
-  assert hidden.kind == nkSym
+  internalAssert hidden.kind == nkSym and hidden.sym.kind == skParam
   result = hidden.sym
 
 proc getEnvParam(routine: PSym): PSym =
@@ -184,161 +192,182 @@ proc getEnvParam(routine: PSym): PSym =
   if hidden.kind == nkSym and hidden.sym.name.s == paramName:
     result = hidden.sym
 
-proc initIterContext(c: POuterContext, iter: PSym) =
-  c.fn = iter
-  c.capturedVars = initIntSet()
+proc initIter(iter: PSym): TIter =
+  result.fn = iter
+  if iter.kind == skClosureIterator:
+    var cp = getEnvParam(iter)
+    if cp == nil:
+      result.obj = createObj(iter, iter.info)
 
-  var cp = getEnvParam(iter)
-  if cp == nil:
-    c.obj = createObj(iter, iter.info)
+      cp = newSym(skParam, getIdent(paramName), iter, iter.info)
+      incl(cp.flags, sfFromGeneric)
+      cp.typ = newType(tyRef, iter)
+      rawAddSon(cp.typ, result.obj)
+      addHiddenParam(iter, cp)
 
-    cp = newSym(skParam, getIdent(paramName), iter, iter.info)
-    incl(cp.flags, sfFromGeneric)
-    cp.typ = newType(tyRef, iter)
-    rawAddSon(cp.typ, c.obj)
-    addHiddenParam(iter, cp)
-
-    c.state = createStateField(iter)
-    addField(c.obj, c.state)
-  else:
-    c.obj = cp.typ.sons[0]
-    assert c.obj.kind == tyObject
-    if c.obj.n.len > 0:
-      c.state = c.obj.n[0].sym
+      result.state = createStateField(iter)
+      rawAddField(result.obj, result.state)
     else:
-      c.state = createStateField(iter)
-      addField(c.obj, c.state)
+      result.obj = cp.typ.sons[0]
+      assert result.obj.kind == tyObject
+      if result.obj.n.len > 0:
+        result.state = result.obj.n[0].sym
+      else:
+        result.state = createStateField(iter)
+        rawAddField(result.obj, result.state)
+    result.closureParam = cp
+    if iter.typ.sons[0] != nil:
+      result.resultSym = newIterResult(iter)
+      #iter.ast.add(newSymNode(c.resultSym))
 
-  c.closureParam = cp
-  if iter.typ.sons[0] != nil:
-    c.resultSym = newIterResult(iter)
-    #iter.ast.add(newSymNode(c.resultSym))
-
-proc newOuterContext(fn: PSym, up: POuterContext = nil): POuterContext =
+proc newOuterContext(fn: PSym): POuterContext =
   new(result)
   result.fn = fn
   result.capturedVars = initIntSet()
   result.processed = initIntSet()
   initIdNodeTable(result.localsToAccess)
-  initIdTable(result.localsToEnv)
   initIdTable(result.lambdasToEnv)
-  result.isIter = fn.kind == skClosureIterator
-  if result.isIter: initIterContext(result, fn)
 
-proc newInnerContext(fn: PSym): PInnerContext =
+proc newEnv(o: POuterContext; up: PEnv, n: PNode; owner: PSym): PEnv =
   new(result)
-  result.fn = fn
-  initIdNodeTable(result.localsToAccess)
-
-proc newEnv(outerProc: PSym, up: PEnv, n: PNode): PEnv =
-  new(result)
-  result.deps = @[]
   result.capturedVars = @[]
-  result.obj = createObj(outerProc, outerProc.info)
+  result.obj = createObj(owner, owner.info)
   result.up = up
   result.attachedNode = n
+  result.fn = owner
+  result.vars = initIntSet()
+  result.next = o.head
+  o.head = result
 
 proc addCapturedVar(e: PEnv, v: PSym) =
   for x in e.capturedVars:
     if x == v: return
-  # XXX meh, just add the state field for every closure for now, it's too
+  # YYY meh, just add the state field for every closure for now, it's too
   # hard to figure out if it comes from a closure iterator:
   if e.obj.n.len == 0: addField(e.obj, createStateField(v.owner))
   e.capturedVars.add(v)
   addField(e.obj, v)
-  
-proc addDep(e, d: PEnv, owner: PSym): PSym =
-  for x, field in items(e.deps):
-    if x == d: return field
-  var pos = sonsLen(e.obj.n)
-  result = newSym(skField, getIdent(upName & $pos), owner, owner.info)
-  result.typ = newType(tyRef, owner)
-  result.position = pos
-  assert d.obj != nil
-  rawAddSon(result.typ, d.obj)
-  addField(e.obj, result)
-  e.deps.add((d, result))
 
 proc newCall(a, b: PSym): PNode =
   result = newNodeI(nkCall, a.info)
   result.add newSymNode(a)
   result.add newSymNode(b)
 
-proc isInnerProc(s, outerProc: PSym): bool {.inline.} =
-  result = s.kind in {skProc, skMethod, skConverter, skClosureIterator} and
-           s.skipGenericOwner == outerProc
+proc isInnerProc(s, outerProc: PSym): bool =
+  if s.kind in {skProc, skMethod, skConverter, skClosureIterator}:
+    var owner = s.skipGenericOwner
+    while true:
+      if owner.isNil: return false
+      if owner == outerProc: return true
+      owner = owner.owner
   #s.typ.callConv == ccClosure
 
-proc addClosureParam(i: PInnerContext, e: PEnv) =
-  var cp = getEnvParam(i.fn)
+proc addClosureParam(fn: PSym; e: PEnv): PSym =
+  var cp = getEnvParam(fn)
   if cp == nil:
-    cp = newSym(skParam, getIdent(paramName), i.fn, i.fn.info)
+    cp = newSym(skParam, getIdent(paramName), fn, fn.info)
     incl(cp.flags, sfFromGeneric)
-    cp.typ = newType(tyRef, i.fn)
+    cp.typ = newType(tyRef, fn)
     rawAddSon(cp.typ, e.obj)
-    addHiddenParam(i.fn, cp)
+    addHiddenParam(fn, cp)
   else:
+    #assert e.obj == nil or e.obj == cp.typ.sons[0]
     e.obj = cp.typ.sons[0]
     assert e.obj.kind == tyObject
-  i.closureParam = cp
-  #echo "closure param added for ", i.fn.name.s, " ", i.fn.id
-
-proc dummyClosureParam(o: POuterContext, i: PInnerContext) =
-  var e = o.currentEnv
-  if idTableGet(o.lambdasToEnv, i.fn) == nil:
-    idTablePut(o.lambdasToEnv, i.fn, e)
-  if i.closureParam == nil: addClosureParam(i, e)
+  result = cp
 
 proc illegalCapture(s: PSym): bool {.inline.} =
   result = skipTypes(s.typ, abstractInst).kind in 
                    {tyVar, tyOpenArray, tyVarargs} or
       s.kind == skResult
 
-proc captureVar(o: POuterContext, i: PInnerContext, local: PSym, 
-                info: TLineInfo) =
-  # for inlined variables the owner is still wrong, so it can happen that it's
-  # not a captured variable at all ... *sigh* 
-  var it = PEnv(idTableGet(o.localsToEnv, local))
-  if it == nil: return
-  
-  if illegalCapture(local) or o.fn.id != local.owner.id or 
-      i.fn.typ.callConv notin {ccClosure, ccDefault}:
-    # Currently captures are restricted to a single level of nesting:
-    localError(info, errIllegalCaptureX, local.name.s)
-  i.fn.typ.callConv = ccClosure
-  #echo "captureVar ", i.fn.name.s, i.fn.id, " ", local.name.s, local.id
-
-  incl(i.fn.typ.flags, tfCapturesEnv)
-
-  # we need to remember which inner most closure belongs to this lambda:
-  var e = o.currentEnv
-  if idTableGet(o.lambdasToEnv, i.fn) == nil:
-    idTablePut(o.lambdasToEnv, i.fn, e)
-
-  # variable already captured:
-  if idNodeTableGet(i.localsToAccess, local) != nil: return
-  if i.closureParam == nil: addClosureParam(i, e)
-  
-  # check which environment `local` belongs to:
-  var access = newSymNode(i.closureParam)
-  addCapturedVar(it, local)
-  if it == e:
-    # common case: local directly in current environment:
-    discard
-  else:
-    # it's in some upper environment:
-    access = indirectAccess(access, addDep(e, it, i.fn), info)
-  access = indirectAccess(access, local, info)
-  if o.isIter:
-    if not containsOrIncl(o.capturedVars, local.id): addField(o.obj, local)
-  else:
-    incl(o.capturedVars, local.id)
-  idNodeTablePut(i.localsToAccess, local, access)
-
 proc interestingVar(s: PSym): bool {.inline.} =
   result = s.kind in {skVar, skLet, skTemp, skForVar, skParam, skResult} and
     sfGlobal notin s.flags
 
+proc nestedAccess(top: PEnv; local: PSym): PNode =
+  # Parts after the transformation are in []:
+  #
+  #  proc main =
+  #    var [:env.]foo = 23
+  #    proc outer(:paramO) =
+  #      [var :envO; createClosure(:envO); :envO.up = paramO]
+  #      proc inner(:paramI) =
+  #        echo [:paramI.up.]foo
+  #      inner([:envO])
+  #    outer([:env])
+  if not (interestingVar(local) and top.fn != local.owner):
+    return nil
+  # check it's in fact a captured variable:
+  var it = top
+  while it != nil:
+    if it.vars.contains(local.id): break
+    it = it.up
+  if it == nil: return nil
+  let envParam = top.fn.getEnvParam
+  internalAssert(not envParam.isNil)
+  var access = newSymNode(envParam)
+  # we could also simply check the tuple type for the field here, I think.
+  it = top.up
+  while it != nil:
+    if it.vars.contains(local.id):
+      access = indirectAccess(access, local, local.info)
+      return access
+    internalAssert it.upField != nil
+    access = indirectAccess(access, it.upField, local.info)
+    it = it.up
+  return nil
+
+proc createUpField(obj, fieldType: PType): PSym =
+  let pos = obj.n.len
+  result = newSym(skField, getIdent(upName), obj.owner, obj.owner.info)
+  result.typ = newType(tyRef, obj.owner)
+  result.position = pos
+  rawAddSon(result.typ, fieldType)
+  addField(obj, result)
+
+proc captureVar(o: POuterContext; top: PEnv; local: PSym; 
+                info: TLineInfo): bool =
+  # first check if we should be concerned at all:
+  var it = top
+  while it != nil:
+    if it.vars.contains(local.id): break
+    it = it.up
+  if it == nil: return false
+  # yes, so mark every 'up' pointer as taken:
+  if illegalCapture(local) or top.fn.typ.callConv notin {ccClosure, ccDefault}:
+    localError(info, errIllegalCaptureX, local.name.s)
+  it = top
+  while it != nil:
+    if it.vars.contains(local.id): break
+    # keep in mind that the first element of the chain belong to top.fn itself
+    # and these don't need any upFields
+    if it.upField == nil and it.up != nil and it.fn != top.fn:
+      it.upField = createUpField(it.obj, it.up.obj)
+
+    if it.fn != local.owner:
+      it.fn.typ.callConv = ccClosure
+      incl(it.fn.typ.flags, tfCapturesEnv)
+      discard addClosureParam(it.fn, it.up)
+
+      if idTableGet(o.lambdasToEnv, it.fn) == nil:
+        idTablePut(o.lambdasToEnv, it.fn, it.up)
+
+    it = it.up
+  # don't do this: 'top' might not require a closure:
+  #if idTableGet(o.lambdasToEnv, it.fn) == nil:
+  #  idTablePut(o.lambdasToEnv, it.fn, top)
+
+  # mark as captured:
+  #if top.iter != nil:
+  #  if not containsOrIncl(o.capturedVars, local.id):
+  #    #addField(top.iter.obj, local)
+  #    addCapturedVar(it, local)
+  #else:
+  incl(o.capturedVars, local.id)
+  addCapturedVar(it, local)
+  result = true
+
 proc semCaptureSym*(s, owner: PSym) =
   if interestingVar(s) and owner.id != s.owner.id and s.kind != skResult:
     if owner.typ != nil and not isGenericRoutine(owner):
@@ -350,28 +379,20 @@ proc semCaptureSym*(s, owner: PSym) =
     # since the analysis is not entirely correct, we don't set 'tfCapturesEnv'
     # here
 
-proc gatherVars(o: POuterContext, i: PInnerContext, n: PNode) = 
-  # gather used vars for closure generation
-  if n == nil: return
+proc gatherVars(o: POuterContext; e: PEnv; n: PNode): int =
+  # gather used vars for closure generation; returns number of captured vars
+  if n == nil: return 0
   case n.kind
   of nkSym:
     var s = n.sym
-    if interestingVar(s) and i.fn.id != s.owner.id:
-      captureVar(o, i, s, n.info)
-    elif s.kind in {skProc, skMethod, skConverter} and
-            s.skipGenericOwner == o.fn and 
-            tfCapturesEnv in s.typ.flags and s != i.fn:
-      # call to some other inner proc; we need to track the dependencies for
-      # this:
-      let env = PEnv(idTableGet(o.lambdasToEnv, i.fn))
-      if env == nil: internalError(n.info, "no environment computed")
-      if o.currentEnv != env:
-        discard addDep(o.currentEnv, env, i.fn)
-        internalError(n.info, "too complex environment handling required")
-  of nkEmpty..pred(nkSym), succ(nkSym)..nkNilLit, nkClosure: discard
+    if interestingVar(s) and e.fn != s.owner:
+      if captureVar(o, e, s, n.info): result = 1
+  of nkEmpty..pred(nkSym), succ(nkSym)..nkNilLit, nkClosure, nkProcDef, 
+     nkMethodDef, nkConverterDef, nkMacroDef, nkTemplateDef, nkTypeSection:
+    discard
   else:
-    for k in countup(0, sonsLen(n) - 1): 
-      gatherVars(o, i, n.sons[k])
+    for k in countup(0, sonsLen(n) - 1):
+      result += gatherVars(o, e, n.sons[k])
 
 proc generateThunk(prc: PNode, dest: PType): PNode =
   ## Converts 'prc' into '(thunk, nil)' so that it's compatible with
@@ -403,61 +424,112 @@ proc makeClosure(prc, env: PSym, info: TLineInfo): PNode =
   else:
     result.add(newSymNode(env))
 
-proc transformInnerProc(o: POuterContext, i: PInnerContext, n: PNode): PNode =
+proc newClosureCreationVar(e: PEnv): PSym =
+  result = newSym(skVar, getIdent(envName), e.fn, e.attachedNode.info)
+  incl(result.flags, sfShadowed)
+  result.typ = newType(tyRef, e.fn)
+  result.typ.rawAddSon(e.obj)
+
+proc getClosureVar(e: PEnv): PSym =
+  if e.createdVar == nil:
+    result = newClosureCreationVar(e)
+    e.createdVar = result
+  else:
+    result = e.createdVar
+
+proc findEnv(o: POuterContext; s: PSym): PEnv =
+  var env = o.head
+  while env != nil:
+    if env.fn == s: break
+    env = env.next
+  internalAssert env != nil and env.up != nil
+  result = env.up
+  while result.fn == s: result = result.up
+
+proc transformInnerProc(o: POuterContext; e: PEnv, n: PNode): PNode =
   case n.kind
   of nkEmpty..pred(nkSym), succ(nkSym)..nkNilLit: discard
   of nkSym:
     let s = n.sym
-    if s == i.fn: 
+    if s == e.fn:
       # recursive calls go through (lambda, hiddenParam):
-      assert i.closureParam != nil, i.fn.name.s
-      result = makeClosure(s, i.closureParam, n.info)
+      result = makeClosure(s, getEnvParam(s), n.info)
     elif isInnerProc(s, o.fn) and s.typ.callConv == ccClosure:
-      # ugh: call to some other inner proc; 
-      assert i.closureParam != nil
-      # XXX this is not correct in general! may also be some 'closure.upval'
-      result = makeClosure(s, i.closureParam, n.info)
+      # ugh: call to some other inner proc;
+      result = makeClosure(s, findEnv(o, s).getClosureVar, n.info)
     else:
       # captured symbol?
-      result = idNodeTableGet(i.localsToAccess, n.sym)
-  of nkLambdaKinds, nkIteratorDef:
-    if n.typ != nil:
-      result = transformInnerProc(o, i, n.sons[namePos])
+      result = nestedAccess(e, n.sym)
+      #result = idNodeTableGet(i.localsToAccess, n.sym)
+    #of nkLambdaKinds, nkIteratorDef:
+    #  if n.typ != nil:
+    #    result = transformInnerProc(o, e, n.sons[namePos])
+    #of nkClosure:
+    #  let x = transformInnerProc(o, e, n.sons[0])
+    #  if x != nil: n.sons[0] = x
   of nkProcDef, nkMethodDef, nkConverterDef, nkMacroDef, nkTemplateDef,
-      nkClosure:
+     nkLambdaKinds, nkIteratorDef, nkClosure:
     # don't recurse here:
     discard
   else:
     for j in countup(0, sonsLen(n) - 1):
-      let x = transformInnerProc(o, i, n.sons[j])
+      let x = transformInnerProc(o, e, n.sons[j])
       if x != nil: n.sons[j] = x
 
 proc closureCreationPoint(n: PNode): PNode =
-  result = newNodeI(nkStmtList, n.info)
-  result.add(emptyNode)
-  result.add(n)
+  if n.kind == nkStmtList and n.len >= 1 and n[0].kind == nkEmpty:
+    # we already have a free slot
+    result = n
+  else:
+    result = newNodeI(nkStmtList, n.info)
+    result.add(emptyNode)
+    result.add(n)
+  #result.flags.incl nfLL
 
-proc searchForInnerProcs(o: POuterContext, n: PNode) =
+proc addParamsToEnv(fn: PSym; env: PEnv) =
+  let params = fn.typ.n
+  for i in 1.. <params.len: 
+    if params.sons[i].kind != nkSym:
+      internalError(params.info, "liftLambdas: strange params")
+    let param = params.sons[i].sym
+    env.vars.incl(param.id)
+  # put the 'result' into the environment so it can be captured:
+  let ast = fn.ast
+  if resultPos < sonsLen(ast) and ast.sons[resultPos].kind == nkSym:
+    env.vars.incl(ast.sons[resultPos].sym.id)
+
+proc searchForInnerProcs(o: POuterContext, n: PNode, env: PEnv) =
   if n == nil: return
   case n.kind
-  of nkEmpty..pred(nkSym), succ(nkSym)..nkNilLit: 
+  of nkEmpty..pred(nkSym), succ(nkSym)..nkNilLit:
     discard
   of nkSym:
-    if isInnerProc(n.sym, o.fn) and not containsOrIncl(o.processed, n.sym.id):
-      var inner = newInnerContext(n.sym)
-      let body = n.sym.getBody
-      gatherVars(o, inner, body)
+    let fn = n.sym
+    if isInnerProc(fn, o.fn) and not containsOrIncl(o.processed, fn.id):
+      let body = fn.getBody
+
+      # handle deeply nested captures:
+      let ex = closureCreationPoint(body)
+      let envB = newEnv(o, env, ex, fn)
+      addParamsToEnv(fn, envB)
+      searchForInnerProcs(o, body, envB)
+      fn.ast.sons[bodyPos] = ex
+      
+      let capturedCounter = gatherVars(o, envB, body)
       # dummy closure param needed?
-      if inner.closureParam == nil and n.sym.typ.callConv == ccClosure:
+      if capturedCounter == 0 and fn.typ.callConv == ccClosure:
         #assert tfCapturesEnv notin n.sym.typ.flags
-        dummyClosureParam(o, inner)
-      # only transform if it really needs a closure:
-      if inner.closureParam != nil:
-        let ti = transformInnerProc(o, inner, body)
-        if ti != nil: n.sym.ast.sons[bodyPos] = ti
+        if idTableGet(o.lambdasToEnv, fn) == nil:
+          idTablePut(o.lambdasToEnv, fn, env)
+        discard addClosureParam(fn, env)
+
+      elif fn.getEnvParam != nil:
+        # only transform if it really needs a closure:
+        let ti = transformInnerProc(o, envB, body)
+        if ti != nil: fn.ast.sons[bodyPos] = ti
   of nkLambdaKinds, nkIteratorDef:
     if n.typ != nil:
-      searchForInnerProcs(o, n.sons[namePos])
+      searchForInnerProcs(o, n.sons[namePos], env)
   of nkWhileStmt, nkForStmt, nkParForStmt, nkBlockStmt:
     # some nodes open a new scope, so they are candidates for the insertion
     # of closure creation; however for simplicity we merge closures between
@@ -465,14 +537,11 @@ proc searchForInnerProcs(o: POuterContext, n: PNode) =
     # yield observable changes in semantics. For Zahary we also
     # include ``nkBlock``.
     var body = n.len-1
-    for i in countup(0, body - 1): searchForInnerProcs(o, n.sons[i])
+    for i in countup(0, body - 1): searchForInnerProcs(o, n.sons[i], env)
     # special handling for the loop body:
-    let oldEnv = o.currentEnv
     let ex = closureCreationPoint(n.sons[body])
-    o.currentEnv = newEnv(o.fn, oldEnv, ex)
-    searchForInnerProcs(o, n.sons[body])
+    searchForInnerProcs(o, n.sons[body], newEnv(o, env, ex, env.fn))
     n.sons[body] = ex
-    o.currentEnv = oldEnv
   of nkVarSection, nkLetSection:
     # we need to compute a mapping var->declaredBlock. Note: The definition
     # counts, not the block where it is captured!
@@ -481,26 +550,29 @@ proc searchForInnerProcs(o: POuterContext, n: PNode) =
       if it.kind == nkCommentStmt: discard
       elif it.kind == nkIdentDefs:
         var L = sonsLen(it)
-        if it.sons[0].kind != nkSym: internalError(it.info, "transformOuter")
-        #echo "set: ", it.sons[0].sym.name.s, " ", o.currentBlock == nil
-        idTablePut(o.localsToEnv, it.sons[0].sym, o.currentEnv)
-        searchForInnerProcs(o, it.sons[L-1])
+        if it.sons[0].kind == nkSym:
+          # this can be false for recursive invokations that already
+          # transformed it into 'env.varName':
+          env.vars.incl(it.sons[0].sym.id)
+        searchForInnerProcs(o, it.sons[L-1], env)
       elif it.kind == nkVarTuple:
         var L = sonsLen(it)
         for j in countup(0, L-3):
           #echo "set: ", it.sons[j].sym.name.s, " ", o.currentBlock == nil
-          idTablePut(o.localsToEnv, it.sons[j].sym, o.currentEnv)
-        searchForInnerProcs(o, it.sons[L-1])
+          if it.sons[j].kind == nkSym:
+            env.vars.incl(it.sons[j].sym.id)
+        searchForInnerProcs(o, it.sons[L-1], env)
       else:
-        internalError(it.info, "transformOuter")
+        internalError(it.info, "searchForInnerProcs")
+  of nkClosure:
+    searchForInnerProcs(o, n.sons[0], env)
   of nkProcDef, nkMethodDef, nkConverterDef, nkMacroDef, nkTemplateDef, 
-     nkClosure, nkTypeSection:
+     nkTypeSection:
     # don't recurse here:
-    # XXX recurse here and setup 'up' pointers
     discard
   else:
     for i in countup(0, sonsLen(n) - 1):
-      searchForInnerProcs(o, n.sons[i])
+      searchForInnerProcs(o, n.sons[i], env)
 
 proc newAsgnStmt(le, ri: PNode, info: TLineInfo): PNode = 
   # Bugfix: unfortunately we cannot use 'nkFastAsgn' here as that would
@@ -512,19 +584,6 @@ proc newAsgnStmt(le, ri: PNode, info: TLineInfo): PNode =
   result.sons[0] = le
   result.sons[1] = ri
 
-proc newClosureCreationVar(o: POuterContext; e: PEnv): PSym =
-  result = newSym(skVar, getIdent(envName), o.fn, e.attachedNode.info)
-  incl(result.flags, sfShadowed)
-  result.typ = newType(tyRef, o.fn)
-  result.typ.rawAddSon(e.obj)
-
-proc getClosureVar(o: POuterContext; e: PEnv): PSym =
-  if e.createdVar == nil:
-    result = newClosureCreationVar(o, e)
-    e.createdVar = result
-  else:
-    result = e.createdVar
-
 proc rawClosureCreation(o: POuterContext, scope: PEnv; env: PSym): PNode =
   result = newNodeI(nkStmtList, env.info)
   var v = newNodeI(nkVarSection, env.info)
@@ -548,21 +607,25 @@ proc rawClosureCreation(o: POuterContext, scope: PEnv; env: PSym): PNode =
       idNodeTablePut(o.localsToAccess, local, fieldAccess)
     else:
       result.add(newAsgnStmt(fieldAccess, existing, env.info))
-  # add support for 'up' references:
-  for e, field in items(scope.deps):
-    # add ``env.up = env2``
-    result.add(newAsgnStmt(indirectAccess(env, field, env.info),
-               newSymNode(getClosureVar(o, e)), env.info))
+  if scope.upField != nil:
+    # "up" chain has been used:
+    if scope.up.fn != scope.fn:
+      # crosses function boundary:
+      result.add(newAsgnStmt(indirectAccess(env, scope.upField, env.info),
+                 newSymNode(getEnvParam(scope.fn)), env.info))
+    else:
+      result.add(newAsgnStmt(indirectAccess(env, scope.upField, env.info),
+                 newSymNode(getClosureVar(scope.up)), env.info))
   
 proc generateClosureCreation(o: POuterContext, scope: PEnv): PNode =
-  var env = getClosureVar(o, scope)
+  var env = getClosureVar(scope)
   result = rawClosureCreation(o, scope, env)
 
 proc generateIterClosureCreation(o: POuterContext; env: PEnv;
                                  scope: PNode): PSym =
   if env.createdVarComesFromIter or env.createdVar.isNil:
     # we have to create a new closure:
-    result = newClosureCreationVar(o, env)
+    result = newClosureCreationVar(env)
     let cc = rawClosureCreation(o, env, result)
     var insertPoint = scope.sons[0]
     if insertPoint.kind == nkEmpty: scope.sons[0] = cc
@@ -577,21 +640,22 @@ proc generateIterClosureCreation(o: POuterContext; env: PEnv;
 proc interestingIterVar(s: PSym): bool {.inline.} =
   result = s.kind in {skVar, skLet, skTemp, skForVar} and sfGlobal notin s.flags
 
-proc transformOuterProc(o: POuterContext, n: PNode): PNode
+proc transformOuterProc(o: POuterContext, n: PNode, it: TIter): PNode
 
-proc transformYield(c: POuterContext, n: PNode): PNode =
-  inc c.state.typ.n.sons[1].intVal
-  let stateNo = c.state.typ.n.sons[1].intVal
+proc transformYield(c: POuterContext, n: PNode, it: TIter): PNode =
+  inc it.state.typ.n.sons[1].intVal
+  let stateNo = it.state.typ.n.sons[1].intVal
 
   var stateAsgnStmt = newNodeI(nkAsgn, n.info)
-  stateAsgnStmt.add(indirectAccess(newSymNode(c.closureParam),c.state,n.info))
+  stateAsgnStmt.add(rawIndirectAccess(newSymNode(it.closureParam),
+                    it.state, n.info))
   stateAsgnStmt.add(newIntTypeNode(nkIntLit, stateNo, getSysType(tyInt)))
 
   var retStmt = newNodeI(nkReturnStmt, n.info)
   if n.sons[0].kind != nkEmpty:
     var a = newNodeI(nkAsgn, n.sons[0].info)
-    var retVal = transformOuterProc(c, n.sons[0])
-    addSon(a, newSymNode(c.resultSym))
+    var retVal = transformOuterProc(c, n.sons[0], it)
+    addSon(a, newSymNode(it.resultSym))
     addSon(a, if retVal.isNil: n.sons[0] else: retVal)
     retStmt.add(a)
   else:
@@ -605,17 +669,18 @@ proc transformYield(c: POuterContext, n: PNode): PNode =
   result.add(retStmt)
   result.add(stateLabelStmt)
 
-proc transformReturn(c: POuterContext, n: PNode): PNode =
+proc transformReturn(c: POuterContext, n: PNode, it: TIter): PNode =
   result = newNodeI(nkStmtList, n.info)
   var stateAsgnStmt = newNodeI(nkAsgn, n.info)
-  stateAsgnStmt.add(indirectAccess(newSymNode(c.closureParam),c.state,n.info))
+  stateAsgnStmt.add(rawIndirectAccess(newSymNode(it.closureParam), it.state,
+                    n.info))
   stateAsgnStmt.add(newIntTypeNode(nkIntLit, -1, getSysType(tyInt)))
   result.add(stateAsgnStmt)
   result.add(n)
 
-proc outerProcSons(o: POuterContext, n: PNode) =
+proc outerProcSons(o: POuterContext, n: PNode, it: TIter) =
   for i in countup(0, sonsLen(n) - 1):
-    let x = transformOuterProc(o, n.sons[i])
+    let x = transformOuterProc(o, n.sons[i], it)
     if x != nil: n.sons[i] = x
 
 proc liftIterSym*(n: PNode): PNode =
@@ -631,27 +696,106 @@ proc liftIterSym*(n: PNode): PNode =
   addVar(v, newSymNode(env))
   result.add(v)
   # add 'new' statement:
-  result.add(newCall(getSysSym"internalNew", env))
+  result.add newCall(getSysSym"internalNew", env)
   result.add makeClosure(iter, env, n.info)
 
-proc transformOuterProc(o: POuterContext, n: PNode): PNode =
-  if n == nil: return nil
+template envActive(env): expr =
+  (env.capturedVars.len > 0 or env.upField != nil)
+
+# We have to split up environment creation in 2 steps:
+# 1. Generate it and store it in env.replacementNode
+# 2. Insert replacementNode into its forseen slot.
+# This split is necessary so that assignments belonging to closure
+# creation like 'env.param = param' are not transformed
+# into 'env.param = env.param'.
+proc createEnvironments(o: POuterContext) =
+  var env = o.head
+  while env != nil:
+    if envActive(env):
+      var scope = env.attachedNode
+      assert scope.kind == nkStmtList
+      if scope.sons[0].kind == nkEmpty:
+        # prepare for closure construction:
+        env.replacementNode = generateClosureCreation(o, env)
+    env = env.next
+
+proc finishEnvironments(o: POuterContext) =
+  var env = o.head
+  while env != nil:
+    if env.replacementNode != nil:
+      var scope = env.attachedNode
+      assert scope.kind == nkStmtList
+      if scope.sons[0].kind == nkEmpty:
+        # change the empty node to contain the closure construction:
+        scope.sons[0] = env.replacementNode
+    env = env.next
+
+proc transformOuterProcBody(o: POuterContext, n: PNode; it: TIter): PNode =
+  if nfLL in n.flags:
+    result = nil
+  elif it.fn.kind == skClosureIterator:
+    # unfortunately control flow is still convoluted and we can end up
+    # multiple times here for the very same iterator. We shield against this
+    # with some rather primitive check for now:
+    if n.kind == nkStmtList and n.len > 0:
+      if n.sons[0].kind == nkGotoState: return nil
+      if n.len > 1 and n[1].kind == nkStmtList and n[1].len > 0 and 
+          n[1][0].kind == nkGotoState:
+        return nil
+    result = newNodeI(nkStmtList, it.fn.info)
+    var gs = newNodeI(nkGotoState, it.fn.info)
+    assert it.closureParam != nil
+    assert it.state != nil
+    gs.add(rawIndirectAccess(newSymNode(it.closureParam), it.state, it.fn.info))
+    result.add(gs)
+    var state0 = newNodeI(nkState, it.fn.info)
+    state0.add(newIntNode(nkIntLit, 0))
+    result.add(state0)
+    
+    let newBody = transformOuterProc(o, n, it)
+    if newBody != nil:
+      result.add(newBody)
+    else:
+      result.add(n)
+
+    var stateAsgnStmt = newNodeI(nkAsgn, it.fn.info)
+    stateAsgnStmt.add(rawIndirectAccess(newSymNode(it.closureParam),
+                      it.state, it.fn.info))
+    stateAsgnStmt.add(newIntTypeNode(nkIntLit, -1, getSysType(tyInt)))
+    result.add(stateAsgnStmt)
+    result.flags.incl nfLL
+  else:
+    result = transformOuterProc(o, n, it)
+    if result != nil: result.flags.incl nfLL
+
+proc transformOuterProc(o: POuterContext, n: PNode; it: TIter): PNode =
+  if n == nil or nfLL in n.flags: return nil
   case n.kind
   of nkEmpty..pred(nkSym), succ(nkSym)..nkNilLit: discard
   of nkSym:
     var local = n.sym
 
-    if o.isIter and interestingIterVar(local) and o.fn.id == local.owner.id:
-      if not containsOrIncl(o.capturedVars, local.id): addField(o.obj, local)
-      return indirectAccess(newSymNode(o.closureParam), local, n.info)
+    if isInnerProc(local, o.fn) and o.processed.contains(local.id):
+      o.processed.excl(local.id)
+      let body = local.getBody
+      let newBody = transformOuterProcBody(o, body, initIter(local))
+      if newBody != nil:
+        local.ast.sons[bodyPos] = newBody
+
+    if it.fn.kind == skClosureIterator and interestingIterVar(local) and
+        it.fn == local.owner:
+      # every local goes through the closure:
+      if not containsOrIncl(o.capturedVars, local.id):
+        addField(it.obj, local)
+      return indirectAccess(newSymNode(it.closureParam), local, n.info)
 
     var closure = PEnv(idTableGet(o.lambdasToEnv, local))
-
     if local.kind == skClosureIterator:
       # consider: [i1, i2, i1]  Since we merged the iterator's closure
       # with the captured owning variables, we need to generate the
       # closure generation code again:
-      if local == o.fn: message(n.info, errRecursiveDependencyX, local.name.s)
+      if local == o.fn or local == it.fn:
+        message(n.info, errRecursiveDependencyX, local.name.s)
       # XXX why doesn't this work?
       if closure.isNil:
         return liftIterSym(n)
@@ -662,7 +806,6 @@ proc transformOuterProc(o: POuterContext, n: PNode): PNode =
 
     if closure != nil:
       # we need to replace the lambda with '(lambda, env)':
-      
       let a = closure.createdVar
       if a != nil:
         return makeClosure(local, a, n.info)
@@ -678,14 +821,6 @@ proc transformOuterProc(o: POuterContext, n: PNode): PNode =
         return makeClosure(local, x, n.info)
     
     if not contains(o.capturedVars, local.id): return
-    var env = PEnv(idTableGet(o.localsToEnv, local))
-    if env == nil: return
-    var scope = env.attachedNode
-    assert scope.kind == nkStmtList
-    if scope.sons[0].kind == nkEmpty:
-      # change the empty node to contain the closure construction:
-      scope.sons[0] = generateClosureCreation(o, env)
-    
     # change 'local' to 'closure.local', unless it's a 'byCopy' variable:
     # if sfByCopy notin local.flags:
     result = idNodeTableGet(o.localsToAccess, local)
@@ -694,73 +829,46 @@ proc transformOuterProc(o: POuterContext, n: PNode): PNode =
     # to access the local as a local.
   of nkLambdaKinds, nkIteratorDef:
     if n.typ != nil:
-      result = transformOuterProc(o, n.sons[namePos])
+      result = transformOuterProc(o, n.sons[namePos], it)
   of nkProcDef, nkMethodDef, nkConverterDef, nkMacroDef, nkTemplateDef,
       nkClosure:
     # don't recurse here:
     discard
   of nkHiddenStdConv, nkHiddenSubConv, nkConv:
-    let x = transformOuterProc(o, n.sons[1])
+    let x = transformOuterProc(o, n.sons[1], it)
     if x != nil: n.sons[1] = x
     result = transformOuterConv(n)
   of nkYieldStmt:
-    if o.isIter: result = transformYield(o, n)
-    else: outerProcSons(o, n)
+    if it.fn.kind == skClosureIterator: result = transformYield(o, n, it)
+    else: outerProcSons(o, n, it)
   of nkReturnStmt:
-    if o.isIter: result = transformReturn(o, n)
-    else: outerProcSons(o, n)
+    if it.fn.kind == skClosureIterator: result = transformReturn(o, n, it)
+    else: outerProcSons(o, n, it)
   else:
-    outerProcSons(o, n)
-
-proc liftIterator(c: POuterContext, body: PNode): PNode =
-  let iter = c.fn
-  result = newNodeI(nkStmtList, iter.info)
-  var gs = newNodeI(nkGotoState, iter.info)
-  gs.add(indirectAccess(newSymNode(c.closureParam), c.state, iter.info))
-  result.add(gs)
-  var state0 = newNodeI(nkState, iter.info)
-  state0.add(newIntNode(nkIntLit, 0))
-  result.add(state0)
-  
-  let newBody = transformOuterProc(c, body)
-  if newBody != nil:
-    result.add(newBody)
-  else:
-    result.add(body)
-
-  var stateAsgnStmt = newNodeI(nkAsgn, iter.info)
-  stateAsgnStmt.add(indirectAccess(newSymNode(c.closureParam),
-                    c.state,iter.info))
-  stateAsgnStmt.add(newIntTypeNode(nkIntLit, -1, getSysType(tyInt)))
-  result.add(stateAsgnStmt)
+    outerProcSons(o, n, it)
 
 proc liftLambdas*(fn: PSym, body: PNode): PNode =
   # XXX gCmd == cmdCompileToJS does not suffice! The compiletime stuff needs
   # the transformation even when compiling to JS ...
-  if body.kind == nkEmpty or gCmd == cmdCompileToJS:
+  if body.kind == nkEmpty or gCmd == cmdCompileToJS or 
+      fn.skipGenericOwner.kind != skModule:
     # ignore forward declaration:
     result = body
   else:
     var o = newOuterContext(fn)
     let ex = closureCreationPoint(body)
-    o.currentEnv = newEnv(fn, nil, ex)
-    # put all params into the environment so they can be captured:
-    let params = fn.typ.n
-    for i in 1.. <params.len: 
-      if params.sons[i].kind != nkSym:
-        internalError(params.info, "liftLambdas: strange params")
-      let param = params.sons[i].sym
-      idTablePut(o.localsToEnv, param, o.currentEnv)
-    # put the 'result' into the environment so it can be captured:
-    let ast = fn.ast
-    if resultPos < sonsLen(ast) and ast.sons[resultPos].kind == nkSym:
-      idTablePut(o.localsToEnv, ast.sons[resultPos].sym, o.currentEnv)
-    searchForInnerProcs(o, body)
-    if o.isIter:
-      result = liftIterator(o, ex)
+    let env = newEnv(o, nil, ex, fn)
+    addParamsToEnv(fn, env)
+    searchForInnerProcs(o, body, env)
+    createEnvironments(o)
+    if fn.kind == skClosureIterator:
+      result = transformOuterProcBody(o, body, initIter(fn))
     else:
-      discard transformOuterProc(o, body)
+      discard transformOuterProcBody(o, body, initIter(fn))
       result = ex
+    finishEnvironments(o)
+    #if fn.name.s == "factory" or fn.name.s == "factory2":
+    #  echo rendertree(result, {renderIds})
 
 proc liftLambdasForTopLevel*(module: PSym, body: PNode): PNode =
   if body.kind == nkEmpty or gCmd == cmdCompileToJS:
@@ -768,9 +876,11 @@ proc liftLambdasForTopLevel*(module: PSym, body: PNode): PNode =
   else:
     var o = newOuterContext(module)
     let ex = closureCreationPoint(body)
-    o.currentEnv = newEnv(module, nil, ex)
-    searchForInnerProcs(o, body)
-    discard transformOuterProc(o, body)
+    let env = newEnv(o, nil, ex, module)
+    searchForInnerProcs(o, body, env)
+    createEnvironments(o)
+    discard transformOuterProc(o, body, initIter(module))
+    finishEnvironments(o)
     result = ex
 
 # ------------------- iterator transformation --------------------------------
@@ -846,10 +956,6 @@ proc liftForLoop*(body: PNode): PNode =
 
   loopBody.sons[0] = v2
   var bs = newNodeI(nkBreakState, body.info)
-  #if not env.isNil:
-  #  bs.addSon(indirectAccess(env, 
-  #    newSym(skField, getIdent":state", env, env.info), body.info))
-  #else:
   bs.addSon(call.sons[0])
   loopBody.sons[1] = bs
   loopBody.sons[2] = body[L-1]

compiler/lowerings.nim

@@ -64,6 +64,22 @@ proc createObj*(owner: PSym, info: TLineInfo): PType =
   incl result.flags, tfFinal
   result.n = newNodeI(nkRecList, info)
 
+proc rawAddField*(obj: PType; field: PSym) =
+  assert field.kind == skField
+  field.position = sonsLen(obj.n)
+  addSon(obj.n, newSymNode(field))
+
+proc rawIndirectAccess*(a: PNode; field: PSym; info: TLineInfo): PNode = 
+  # returns a[].field as a node
+  assert field.kind == skField
+  var deref = newNodeI(nkHiddenDeref, info)
+  deref.typ = a.typ.skipTypes(abstractInst).sons[0]
+  addSon(deref, a)
+  result = newNodeI(nkDotExpr, info)
+  addSon(result, deref)
+  addSon(result, newSymNode(field))
+  result.typ = field.typ
+
 proc addField*(obj: PType; s: PSym) =
   # because of 'gensym' support, we have to mangle the name with its ID.
   # This is hacky but the clean solution is much more complex than it looks.

lib/system.nim

@@ -2942,9 +2942,10 @@ proc locals*(): TObject {.magic: "Locals", noSideEffect.} =
   ##   # -> B is 1
   discard
 
-proc deepCopy*[T](x: T): T {.magic: "DeepCopy", noSideEffect.}
+proc deepCopy*[T](x: T): T {.magic: "DeepCopy", noSideEffect.} =
   ## performs a deep copy of `x`. This is also used by the code generator
   ## for the implementation of ``spawn``.
+  discard
 
 when not defined(booting):
   type

tests/closure/tclosure.nim

@@ -1,7 +1,6 @@
 discard """
   file: "tclosure.nim"
-  output: "2 4 6 8 10"
-  disabled: true
+  output: "1 3 6 11 20"
 """
 # Test the closure implementation
 
@@ -30,7 +29,8 @@ proc testA() =
 testA()
 
 myData.each do (x: int):
-  write(stout, x)
+  write(stdout, x)
+  write(stdout, " ")
 
 #OUT 2 4 6 8 10
 
@@ -42,6 +42,6 @@ type
 proc getInterf(): ITest =
   var shared: int
   
-  return (setter: proc (x) = shared = x,
+  return (setter: proc (x: int) = shared = x,
           getter: proc (): int = return shared)
 

tests/closure/tnestedclosure.nim

@@ -0,0 +1,36 @@
+discard """
+  output: '''foo88
+23 24foo 88
+foo88
+23 24foo 88'''
+"""
+
+# test nested closure
+proc main(param: int) =
+  var foo = 23
+  proc outer(outerParam: string) =
+    var outerVar = 88
+    echo outerParam, outerVar
+    proc inner() =
+      block Test:
+        echo foo, " ", param, outerParam, " ", outerVar
+    inner()
+  outer("foo")
+
+# test simple closure within dummy 'main':
+proc dummy =
+  proc main2(param: int) =
+    var foo = 23
+    proc outer(outerParam: string) =
+      var outerVar = 88
+      echo outerParam, outerVar
+      proc inner() =
+        block Test:
+          echo foo, " ", param, outerParam, " ", outerVar
+      inner()
+    outer("foo")
+  main2(24)
+
+dummy()
+
+main(24)