From d3df782fca7a1ab3bd8ff6c61382261a36f6c4f3 Mon Sep 17 00:00:00 2001
From: Li Jie <cpunion@gmail.com>
Date: Mon, 5 Aug 2024 17:24:36 +0800
Subject: [PATCH] asyncic: coReturn

---
 cl/async.go | 31 ++++++++++++++-----------------
 cl/instr.go |  2 +-
 ssa/coro.go | 32 ++++++++++++++------------------
 3 files changed, 29 insertions(+), 36 deletions(-)

diff --git a/cl/async.go b/cl/async.go
index df94d60f..bb81e01d 100644
--- a/cl/async.go
+++ b/cl/async.go
@@ -85,33 +85,30 @@ func (p *context) coResume(b llssa.Builder, args []ssa.Value) {
 	}
 }
 
-func (p *context) coReturn(b llssa.Builder, args []ssa.Value) {
-	cargs := make([]llssa.Expr, len(args))
-	for i, arg := range args {
-		cargs[i] = p.compileValue(b, arg)
-	}
-	b.CoReturn(cargs...)
-}
-
-func (p *context) coYield(b llssa.Builder, fn *ssa.Function, args []ssa.Value) {
+func (p *context) getSetValueFunc(fn *ssa.Function) llssa.Function {
 	typ := fn.Signature.Recv().Type()
 	mthds := p.goProg.MethodSets.MethodSet(typ)
-	// TODO(lijie): make llgo instruction callable (e.g. llgo.yield) to avoid extra setValue method
-	var setValue *ssa.Function
 	for i := 0; i < mthds.Len(); i++ {
 		m := mthds.At(i)
 		if ssaMthd := p.goProg.MethodValue(m); ssaMthd != nil {
 			if ssaMthd.Name() == "setValue" || strings.HasPrefix(ssaMthd.Name(), "setValue[") {
-				setValue = ssaMthd
-				break
+				setValueFn, _, _ := p.compileFunction(ssaMthd)
+				return setValueFn
 			}
 		}
 	}
-	if setValue == nil {
-		panic("coYield(): not found method setValue")
-	}
+	panic("method setValue not found on type " + typ.String())
+}
+
+func (p *context) coReturn(b llssa.Builder, fn *ssa.Function, args []ssa.Value) {
+	setValueFn := p.getSetValueFunc(fn)
+	value := p.compileValue(b, args[1])
+	b.CoReturn(setValueFn, value)
+}
+
+func (p *context) coYield(b llssa.Builder, fn *ssa.Function, args []ssa.Value) {
+	setValueFn := p.getSetValueFunc(fn)
 	value := p.compileValue(b, args[1])
-	setValueFn, _, _ := p.compileFunction(setValue)
 	// TODO(lijie): find whether the co.Yield/co.Return is the last instruction
 	final := b.Const(constant.MakeBool(false), b.Prog.Bool())
 	b.CoYield(setValueFn, value, final)
diff --git a/cl/instr.go b/cl/instr.go
index 9a65ba8e..50cfc22f 100644
--- a/cl/instr.go
+++ b/cl/instr.go
@@ -408,7 +408,7 @@ func (p *context) call(b llssa.Builder, act llssa.DoAction, call *ssa.CallCommon
 		case llgoCoResume:
 			p.coResume(b, args)
 		case llgoCoReturn:
-			p.coReturn(b, args)
+			p.coReturn(b, cv, args)
 		case llgoCoYield:
 			p.coYield(b, cv, args)
 		case llgoCoRun:
diff --git a/ssa/coro.go b/ssa/coro.go
index a1e35eaa..9864e2f3 100644
--- a/ssa/coro.go
+++ b/ssa/coro.go
@@ -378,10 +378,6 @@ func (b Builder) AsyncToken() Expr {
 	return b.asyncToken
 }
 
-func (b Builder) SetAsyncToken(token Expr) {
-	b.asyncToken = token
-}
-
 func (b Builder) EndAsync() {
 	b.onReturn()
 }
@@ -452,6 +448,9 @@ func (b Builder) BeginAsync(fn Function, entryBlk, allocBlk, cleanBlk, suspdBlk,
 	b.Unreachable()
 
 	b.onSuspBlk = func(nextBlk BasicBlock) (BasicBlock, BasicBlock, BasicBlock) {
+		if nextBlk == nil {
+			nextBlk = trapBlk
+		}
 		return suspdBlk, nextBlk, cleanBlk
 	}
 	b.onReturn = func() {
@@ -627,16 +626,21 @@ func (b Builder) CoSuspend(save, final Expr, nextBlk BasicBlock) {
 	b.SetBlock(nextBlk)
 }
 
-func (b Builder) CoReturn(args ...Expr) {
+func (b Builder) CoReturn(setValueFn Function, value Expr) {
 	if !b.async {
 		panic(fmt.Errorf("return %v not in async block", b.Func.Name()))
 	}
-
-	b.Func.MakeBlock("")
-	nextBlk := b.Func.Block(b.blk.idx + 1)
-	_, _, cleanBlk := b.onSuspBlk(nextBlk)
+	b.Call(setValueFn.Expr, b.promise, value)
+	_, _, cleanBlk := b.onSuspBlk(nil)
 	b.Jump(cleanBlk)
-	b.SetBlock(nextBlk)
+}
+
+func (b Builder) CoYield(setValueFn Function, value Expr, final Expr) {
+	if !b.async {
+		panic(fmt.Errorf("yield %v not in async block", b.Func.Name()))
+	}
+	b.Call(setValueFn.Expr, b.promise, value)
+	b.CoSuspend(b.AsyncToken(), final, nil)
 }
 
 /*
@@ -676,11 +680,3 @@ func (b Builder) CoAwaitSuspendHandle(awaiter, handle, f Expr) {
 	b.Call(fn, awaiter, handle, f)
 }
 */
-
-func (b Builder) CoYield(setValueFn Function, value Expr, final Expr) {
-	if !b.async {
-		panic(fmt.Errorf("yield %v not in async block", b.Func.Name()))
-	}
-	b.Call(setValueFn.Expr, b.promise, value)
-	b.CoSuspend(b.AsyncToken(), final, nil)
-}