asyncio: trap on unexpected resume

cl/async.go

@@ -67,7 +67,7 @@ func (p *context) coAwait(b llssa.Builder, args []ssa.Value) llssa.Expr {
 }
 
 func (p *context) coSuspend(b llssa.Builder, final llssa.Expr) {
-	b.CoSuspend(b.AsyncToken(), final)
+	b.CoSuspend(b.AsyncToken(), final, nil)
 }
 
 func (p *context) coDone(b llssa.Builder, args []ssa.Value) llssa.Expr {

cl/compile.go

@@ -232,11 +232,12 @@ func (p *context) compileFuncDecl(pkg llssa.Package, f *ssa.Function) (llssa.Fun
 	nBlkOff := 0
 	if nblk := len(f.Blocks); nblk > 0 {
 		if async {
-			nBlkOff = 4
+			nBlkOff = 5
 			fn.MakeBlock("entry")
 			fn.MakeBlock("alloc")
 			fn.MakeBlock("clean")
 			fn.MakeBlock("suspend")
+			fn.MakeBlock("trap")
 		}
 		fn.MakeBlocks(nblk)   // to set fn.HasBody() = true
 		if f.Recover != nil { // set recover block

ssa/coro.go

@@ -378,8 +378,7 @@ func (b Builder) SetAsyncToken(token Expr) {
 }
 
 func (b Builder) EndAsync() {
-	_, _, cleanBlk := b.onSuspBlk(b.blk)
-	b.Jump(cleanBlk)
+	b.onReturn()
 }
 
 /*
@@ -409,7 +408,8 @@ func (b Builder) BeginAsync(fn Function) {
 	allocBlk := fn.Block(1)
 	cleanBlk := fn.Block(2)
 	suspdBlk := fn.Block(3)
-	beginBlk := fn.Block(4)
+	trapBlk := fn.Block(4)
+	beginBlk := fn.Block(5)
 
 	b.SetBlock(entryBlk)
 	promiseSize := b.Const(constant.MakeUint64(b.Prog.SizeOf(promiseTy)), b.Prog.Int64()).SetName("promise.size")
@@ -448,9 +448,16 @@ func (b Builder) BeginAsync(fn Function) {
 	b.CoEnd(hdl, b.Prog.BoolVal(false), b.Prog.TokenNone())
 	b.Return(promise)
 
+	b.SetBlock(trapBlk)
+	b.LLVMTrap()
+	b.Unreachable()
+
 	b.onSuspBlk = func(nextBlk BasicBlock) (BasicBlock, BasicBlock, BasicBlock) {
 		return suspdBlk, nextBlk, cleanBlk
 	}
+	b.onReturn = func() {
+		b.CoSuspend(b.asyncToken, b.Prog.BoolVal(true), trapBlk)
+	}
 }
 
 // -----------------------------------------------------------------------------
@@ -600,14 +607,15 @@ func (b Builder) coSuspend(save, final Expr) Expr {
 	return b.Call(fn, save, final)
 }
 
-func (b Builder) CoSuspend(save, final Expr) {
+func (b Builder) CoSuspend(save, final Expr, nextBlk BasicBlock) {
 	if !b.async {
 		panic(fmt.Errorf("suspend %v not in async block", b.Func.Name()))
 	}
+	if nextBlk == nil {
+		b.Func.MakeBlock("")
+		nextBlk = b.Func.Block(b.blk.idx + 1)
+	}
 	ret := b.coSuspend(save, final)
-	// add resume block
-	b.Func.MakeBlock("")
-	nextBlk := b.Func.Block(b.blk.idx + 1)
 	susp, next, clean := b.onSuspBlk(nextBlk)
 	swt := b.Switch(ret, susp)
 	swt.Case(b.Const(constant.MakeInt64(0), b.Prog.Byte()), next)
@@ -669,5 +677,5 @@ func (b Builder) CoYield(setValueFn Function, value Expr, final Expr) {
 		panic(fmt.Errorf("yield %v not in async block", b.Func.Name()))
 	}
 	b.Call(setValueFn.Expr, b.promise, value)
-	b.CoSuspend(b.AsyncToken(), b.Prog.BoolVal(false))
+	b.CoSuspend(b.AsyncToken(), final, nil)
 }

ssa/eh.go

@@ -55,6 +55,13 @@ func (p Program) tySiglongjmp() *types.Signature {
 	return p.sigljmpTy
 }
 
+func (p Program) tyLLVMTrap() *types.Signature {
+	if p.llvmTrapTy == nil {
+		p.llvmTrapTy = types.NewSignatureType(nil, nil, nil, nil, nil, false)
+	}
+	return p.llvmTrapTy
+}
+
 func (b Builder) AllocaSigjmpBuf() Expr {
 	prog := b.Prog
 	n := unsafe.Sizeof(sigjmpbuf{})
@@ -77,6 +84,11 @@ func (b Builder) Siglongjmp(jb, retval Expr) {
 	// b.Unreachable()
 }
 
+func (b Builder) LLVMTrap() {
+	fn := b.Pkg.cFunc("llvm.trap", b.Prog.tyLLVMTrap())
+	b.Call(fn)
+}
+
 // -----------------------------------------------------------------------------
 
 const (

ssa/package.go

@@ -193,6 +193,8 @@ type aProgram struct {
 	sigsetjmpTy *types.Signature
 	sigljmpTy   *types.Signature
 
+	llvmTrapTy *types.Signature
+
 	// coroutine manipulation intrinsics (ordered by LLVM coroutine doc)
 	coDestroyTy *types.Signature
 	coResumeTy  *types.Signature

ssa/stmt_builder.go

@@ -68,6 +68,7 @@ type aBuilder struct {
 	asyncToken Expr
 	promise    Expr
 	onSuspBlk  func(blk BasicBlock) (susp BasicBlock, next BasicBlock, clean BasicBlock)
+	onReturn   func()
 	blkOffset  int
 }