introduce coroutines to the language

compiler/ast/ast_query.nim

@@ -625,6 +625,14 @@ proc isSinkParam*(s: PSym): bool {.inline.} =
 proc isSinkType*(t: PType): bool {.inline.} =
   t.kind == tySink
 
+proc isCoroutineConstr*(prc: PSym): bool {.inline.} =
+  ## Returns whether `prc` is the symbol of a coroutine *constructor*.
+  sfCoroutine in prc.flags and tfCoroutine in prc.typ.flags
+
+proc isCoroutine*(prc: PSym): bool {.inline.} =
+  ## Returns whether `prc` is the symbol of a coroutine.
+  sfCoroutine in prc.flags and tfCoroutine notin prc.typ.flags
+
 const magicsThatCanRaise* = {
   mNone, mParseExprToAst, mParseStmtToAst, mEcho, mChckRange }
 

compiler/ast/ast_types.nim

@@ -413,6 +413,7 @@ type
     sfUsedInFinallyOrExcept  ## symbol is used inside an 'except' or 'finally'
     sfNoalias         ## 'noalias' annotation, means C's 'restrict'
     sfEffectsDelayed  ## an 'effectsDelayed' parameter
+    sfCoroutine       ## the routine is a coroutine
 
   TSymFlags* = set[TSymFlag]
 
@@ -640,6 +641,7 @@ type
     tfByCopy,         ## pass object/tuple by copy (C backend)
     tfByRef,          ## pass object/tuple by reference (C backend)
     tfIterator,       ## type is really an iterator, not a tyProc
+    tfCoroutine,      ## type is that of a coroutine
     tfNotNil,         ## type cannot be 'nil'
     tfRequiresInit,   ## type constains a "not nil" constraint somewhere or
                       ## a `requiresInit` field, so the default zero init
@@ -799,6 +801,7 @@ type
     mIsMainModule, mCompileDate, mCompileTime, mProcCall,
     mCpuEndian, mHostOS, mHostCPU, mBuildOS, mBuildCPU, mAppType,
     mCompileOption, mCompileOptionArg,
+    mSuspend,
     mNLen, mNChild, mNSetChild, mNAdd, mNAddMultiple, mNDel,
     mNKind, mNSymKind,
 

compiler/ast/trees.nim

@@ -183,6 +183,16 @@ proc effectSpec*(n: PNode, effectType: TSpecialWord): PNode =
         result.add(it[1])
       return
 
+proc coroutineSpec*(n: PNode): PType =
+  ## Returns the instance type specified by the ``.coroutine`` pragma in
+  ## pragma list `n`, or nil, if no instance type is specified.
+  let p = findPragma(n, wCoroutine)
+  assert p != nil, "has no coroutine specification"
+  if p.kind == nkExprColonExpr:
+    p[1].typ
+  else:
+    nil
+
 proc unnestStmts(n, result: PNode) =
   case n.kind
   of nkStmtList:

compiler/ast/types.nim

@@ -1227,6 +1227,11 @@ proc getProcConvMismatch*(
     result[0].incl pcmNotIterator
     result[1] = isNone
 
+  if (tfCoroutine in f.flags) != (tfCoroutine in a.flags):
+    # TODO: use the a dedicated enum field
+    result[0].incl pcmNotIterator
+    result[1] = isNone
+
   if f.callConv != a.callConv:
     # valid to pass a 'nimcall' thingie to 'closure':
     if f.callConv == ccClosure and a.callConv == ccNimCall:
@@ -1544,3 +1549,14 @@ proc classifyBackendView*(t: PType): BackendViewKind =
      tyGenericParam, tyForward, tyBuiltInTypeClass, tyCompositeTypeClass,
      tyAnd, tyOr, tyNot, tyAnything, tyFromExpr:
     unreachable()
+
+proc lookupInType*(typ: PType, field: PIdent): PSym =
+  ## Searches for a field with the given identifier (`field`) in the object
+  ## type hierarchy of `typ`.
+  var typ = typ
+  while typ != nil:
+    typ = typ.skipTypes(skipPtrs)
+    result = lookupInRecord(typ.n, field)
+    if result != nil:
+      return
+    typ = typ.base

compiler/ast/typesrenderer.nim

@@ -406,6 +406,9 @@ proc typeToString*(typ: PType, prefer: TPreferedDesc = preferName): string =
       result.add(')')
       if t.len > 0 and t[0] != nil: result.add(": " & typeToString(t[0]))
       var prag = if t.callConv == ccNimCall and tfExplicitCallConv notin t.flags: "" else: $t.callConv
+      if tfCoroutine in t.flags:
+        addSep(prag)
+        prag.add "coroutine"
       if tfNoSideEffect in t.flags:
         addSep(prag)
         prag.add("noSideEffect")

compiler/ast/wordrecg.nim

@@ -101,7 +101,7 @@ type
     wStdIn = "stdin", wStdOut = "stdout", wStdErr = "stderr",
 
     wInOut = "inout", wByCopy = "bycopy", wByRef = "byref", wOneWay = "oneway",
-    wBitsize = "bitsize", wImportHidden = "all",
+    wBitsize = "bitsize", wImportHidden = "all", wCoroutine = "coroutine"
 
   TSpecialWords* = set[TSpecialWord]
 

compiler/front/condsyms.nim

@@ -77,3 +77,4 @@ proc initDefines*(symbols: StringTableRef) =
   defineSymbol("nimskullNewExceptionRt")
   defineSymbol("nimskullNoNkStmtListTypeAndNkBlockType")
   defineSymbol("nimskullNoNkNone")
+  defineSymbol("nimskullHasCoroutines")

compiler/sem/coroutines.nim

@@ -0,0 +1,195 @@
+## Implements the coroutine-related transformations. Despite being a separate
+## module, the coroutine transformation(s) is heavily intertwined with
+## `transf <#transf>`_.
+
+import
+  compiler/ast/[
+    ast,
+    idents,
+    types
+  ],
+  compiler/modules/[
+    magicsys,
+    modulegraphs
+  ],
+  compiler/sem/[
+    closureiters,
+    lambdalifting,
+    lowerings
+  ],
+  compiler/utils/[
+    idioms,
+  ]
+
+proc computeFieldStart(t: PType, pos: var int) =
+  ## Recursively traverses `t` and updates `pos` to the field position
+  ## of the first field a type inheriting from `t` would have.
+  proc aux(n: PNode, p: var int) =
+    case n.kind
+    of nkRecCase:
+      inc p # discriminator
+      for i in 1..<n.len:
+        aux(n[i], p)
+    of nkRecList:
+      for it in n.items:
+        aux(it, p)
+    of nkSym:
+      inc p
+    else:
+      unreachable()
+
+  let t = t.skipTypes(skipPtrs)
+  assert t.kind == tyObject
+  if t.len > 0 and t.base != nil:
+    computeFieldStart(t.base, pos)
+
+  aux(t.n, pos)
+
+proc preTransformConstr*(g: ModuleGraph, idgen: IdGenerator, prc: PSym, body: PNode): PNode =
+  ## Transforms the `body` of coroutine constructor `prc` into:
+  ##
+  ##   .. code-block:: nim
+  ##
+  ##     discard (proc inner() {.closure.} =
+  ##       body
+  ##     )
+  ##
+  ## Since the owner of the locals within `body` is not changed, the lambda-
+  ## lifting pass will lift them into the environment type.
+  let cache = g.cache
+  let base = g.getCompilerProc("CoroutineBase").typ
+
+  # setup the actual coroutine:
+  let inner = newSym(skProc, prc.name, nextSymId idgen, prc, prc.info,
+                     prc.options)
+  inner.flags.incl sfCoroutine
+  inner.ast = newProcNode(nkLambda, prc.info, body,
+                          newTree(nkFormalParams, newNodeIT(nkType, prc.info, base)),
+                          newSymNode(inner),
+                          newNode(nkEmpty),
+                          newNode(nkEmpty),
+                          newNode(nkEmpty),
+                          newNode(nkEmpty))
+  inner.typ = newProcType(prc.info, nextTypeId idgen, inner)
+  inner.typ[0] = base # return type
+  # temporarily mark the procedure as a closure procedure, so that the lambda-
+  # lifting pass visits it
+  inner.typ.callConv = ccClosure
+
+  # temporarily stash the ``self`` symbol node in the inner procedure's
+  # dispatcher slot
+  inner.ast.sons.setLen(dispatcherPos + 1)
+  inner.ast[dispatcherPos] = move prc.ast[dispatcherPos]
+
+  # result symbol for the coroutine:
+  inner.ast[resultPos] = newSymNode:
+    newSym(skResult, cache.getIdent("result"), nextSymId idgen, inner,
+           inner.info, base)
+
+  # place the instance base type in the constructor's dispatcher
+  # slot, for the lambda-lifting pass to later fetch it
+  prc.ast[dispatcherPos] = newNodeIT(nkType, prc.info, prc.typ[0])
+
+  # fix the result variable for the constructor:
+  prc.ast[resultPos] = newSymNode:
+    newSym(skResult, cache.getIdent("result"), nextSymId idgen, prc,
+           prc.info, prc.typ[0])
+
+  let body = copyNodeWithKids(inner.ast)
+  body.typ = inner.typ
+  result = nkDiscardStmt.newTree(body)
+
+proc transformCoroutineConstr*(g: ModuleGraph, idgen: IdGenerator, prc: PSym,
+                               body: PNode): PNode =
+  ## Post-processes the transformed coroutine instance constructor procedure,
+  ## completing the body of the constructor.
+  ##
+  ## The `body` is expected to have undergone the transf pass, including
+  ## lambda-lifting.
+  let
+    cache = g.cache
+    base = g.getCompilerProc("CoroutineBase").typ
+    # this is a bit brittle. We rely on the exact positions in the AST
+    inner = body.lastSon[0][0].sym
+    selfSym = move inner.ast[dispatcherPos]
+    res = prc.ast[resultPos].sym
+
+  result = body
+
+  # remove the discard statement injected earlier:
+  result.delSon(result.len - 1)
+
+  let
+    envLocal = body[0][0][0] # the local injected by lambda-lifting
+    constr = body[0][0][2] # the env construction expression
+
+  # patch the environment construction:
+  constr.add nkExprColonExpr.newTree(
+    newSymNode lookupInType(base, cache.getIdent("fn")),
+    newSymNode inner
+    )
+  # the state needs to be initialized to -4, to signal that the instance is
+  # suspended:
+  constr.add nkExprColonExpr.newTree(
+    newSymNode lookupInType(base, cache.getIdent("state")),
+    newIntTypeNode(-4, g.getSysType(prc.info, tyInt32))
+  )
+
+  # add the result assignment:
+  result.add newAsgnStmt(newSymNode(res),
+                        newTreeIT(nkObjDownConv, prc.info, res.typ, envLocal))
+  # init the lifted ``self`` symbol:
+  if getFieldFromObj(envLocal.typ.base, selfSym.sym) != nil:
+    result.add newAsgnStmt(indirectAccess(envLocal, selfSym.sym, selfSym.info),
+                           newSymNode res)
+
+proc transformCoroutine*(g: ModuleGraph, idgen: IdGenerator, prc: PSym,
+                         body: PNode): PNode =
+  ## Given the ``trans``formed and lambda-lifted `body`, applies the
+  ## transformation for turning the coroutine `prc` into a resumable procedure
+  ## (uses `closureiters <#closureiters>`_ underneath).
+  ##
+  ## Also takes care of fixing the signature of `prc`.
+  let
+    base = g.getCompilerProc("CoroutineBase").typ
+    body = transformClosureIterator(g, idgen, prc, body)
+    param = prc.getEnvParam()
+
+  # replace the hidden parameter with one that has the correct type:
+  let newParam = copySym(param, nextSymId idgen)
+  newParam.typ = base
+
+  # the original parameter is turned into a cursor local and is injected
+  # into the body. Reusing the symbol means that body doesn't need
+  # to be patched
+  param.kind = skLet
+  param.flags.incl sfCursor
+
+  # fix the signature. It needs to be ``CoroutineBase -> CoroutineBase``
+  prc.typ.callConv = ccNimCall
+  prc.typ.rawAddSon(base) # first parameter type
+  prc.typ.n.add newSymNode(newParam)
+  # replace the hidden parameter (it's no longer hidden):
+  prc.ast[paramsPos][^1] = newSymNode(newParam)
+
+  # inject a definition for the local and emit ``result`` initialization
+  result = nkStmtList.newTree(
+    nkLetSection.newTree(
+      newIdentDefs(newSymNode(param),
+                   newTreeIT(nkObjDownConv, body.info, param.typ,
+                             newSymNode(newParam)))),
+    # XXX: always assigning the continuation to the result is inefficient; it
+    #      should happen on return
+    newAsgnStmt(prc.ast[resultPos],
+                indirectAccess(newSymNode(newParam), "next", newParam.info,
+                               g.cache)),
+    body
+  )
+
+  # the fields in the constructed environment are wrong, they need to be
+  # patched
+  let obj = param.typ.base
+  var start = 0
+  computeFieldStart(obj.base, start)
+  for it in obj.n.items:
+    it.sym.position += start

compiler/sem/lambdalifting.nim

@@ -167,7 +167,10 @@ proc createStateField(g: ModuleGraph; iter: PSym; idgen: IdGenerator): PSym =
   result.typ = createClosureIterStateType(g, iter, idgen)
 
 proc createEnvObj(g: ModuleGraph; idgen: IdGenerator; owner: PSym; info: TLineInfo): PType =
-  result = createObj(g, idgen, owner, info, final=false)
+  if sfCoroutine in owner.flags:
+    result = createObj(g, idgen, owner, info, owner.ast[dispatcherPos].typ)
+  else:
+    result = createObj(g, idgen, owner, info, g.getCompilerProc("RootObj").typ)
 
 proc getClosureIterResult*(g: ModuleGraph; iter: PSym; idgen: IdGenerator): PSym =
   if resultPos < iter.ast.len:
@@ -591,7 +594,12 @@ proc accessEnv(available: PSym, wanted: PType, info: TLineInfo,
     result = rawIndirectAccess(result, upField, info)
 
 proc getStateField*(g: ModuleGraph; owner: PSym): PSym =
-  getHiddenParam(g, owner).typ.base.n[0].sym
+  if sfCoroutine in owner.flags:
+    # for coroutines, the ``state`` field symbol comes from the
+    # ``CoroutineBase`` type
+    g.getCompilerProc("CoroutineBase").typ.base.n[1].sym
+  else:
+    getHiddenParam(g, owner).typ.base.n[0].sym
 
 proc symToClosure(n: PNode; graph: ModuleGraph; idgen: IdGenerator;
                   c: LiftingPass): PNode =
@@ -682,9 +690,10 @@ proc liftCapturedVars(n: PNode, graph: ModuleGraph, idgen: IdGenerator,
       n[1] = liftCapturedVars(n[1], graph, idgen, c)
       if n[1].kind == nkClosure: result = n[1]
   of nkReturnStmt:
-    if n[0].kind in {nkAsgn, nkFastAsgn}:
+    if n[0].kind in {nkAsgn, nkFastAsgn} and sfCoroutine notin c.owner.flags:
       # let's not touch the LHS in order to make the lifting pass
-      # correct when `result` is lifted
+      # correct when `result` is lifted. For coroutines, the LHS needs to
+      # be rewritten too!
       n[0][1] = liftCapturedVars(n[0][1], graph, idgen, c)
     else:
       n[0] = liftCapturedVars(n[0], graph, idgen, c)
@@ -775,6 +784,19 @@ proc liftLambdas*(g: ModuleGraph; fn: PSym, body: PNode;
     # callsite
     result.body = liftCapturedVars(body, g, idgen, initLiftingPass(d, param))
     result.env = param
+  elif fn.isCoroutine and fn.typ.callConv == ccClosure:
+    # coroutines are somewhat special. In the context of lambda-lifting,
+    # they're similar to closure iterators, with the difference that the
+    # environment parameter type is already correct
+    var d = initDetectionPass(g, fn, idgen)
+    detectCapturedVars(body, fn, d)
+
+    let param = getEnvParam(fn)
+    assert param != nil
+
+    prepareInnerRoutines(d, idgen, param.typ, fn.info)
+    result.body = liftCapturedVars(body, g, idgen, initLiftingPass(d, param))
+    result.env = param
   elif body.kind != nkEmpty:
     assert fn.typ.callConv != ccClosure or getEnvParam(fn) != nil,
       "missing environment parameter"
@@ -895,7 +917,7 @@ proc liftForLoop*(g: ModuleGraph; body: PNode; idgen: IdGenerator;
       callee = newSym(skLet, op[0].sym.name, nextSymId(idgen), owner, op.info)
     callee.typ = op.typ
 
-    if owner.isIterator:
+    if owner.isIterator or owner.isCoroutine:
       # meh, we have to add the local to the environment; it might be used
       # across yields
       op = freshVarForClosureIter(g, callee, idgen, owner)