From bb03df70590ec3d0e8e796cd1ff3ffccfbe7f68b Mon Sep 17 00:00:00 2001
From: Li Jie <cpunion@gmail.com>
Date: Mon, 5 Aug 2024 16:42:51 +0800
Subject: [PATCH] ssa: refactor Builder.BeginAsync

---
 cl/compile.go | 16 +++++++++-------
 ssa/coro.go   | 11 ++---------
 2 files changed, 11 insertions(+), 16 deletions(-)

diff --git a/cl/compile.go b/cl/compile.go
index 1ebef907..81f716c6 100644
--- a/cl/compile.go
+++ b/cl/compile.go
@@ -231,15 +231,17 @@ func (p *context) compileFuncDecl(pkg llssa.Package, f *ssa.Function) (llssa.Fun
 	}
 	nBlkOff := 0
 	if nblk := len(f.Blocks); nblk > 0 {
+		var entryBlk, allocBlk, cleanBlk, suspdBlk, trapBlk, beginBlk llssa.BasicBlock
 		if async {
 			nBlkOff = 5
-			fn.MakeBlock("entry")
-			fn.MakeBlock("alloc")
-			fn.MakeBlock("clean")
-			fn.MakeBlock("suspend")
-			fn.MakeBlock("trap")
+			entryBlk = fn.MakeBlock("entry")
+			allocBlk = fn.MakeBlock("alloc")
+			cleanBlk = fn.MakeBlock("clean")
+			suspdBlk = fn.MakeBlock("suspend")
+			trapBlk = fn.MakeBlock("trap")
 		}
-		fn.MakeBlocks(nblk)   // to set fn.HasBody() = true
+		fn.MakeBlocks(nblk) // to set fn.HasBody() = true
+		beginBlk = fn.Block(nBlkOff)
 		if f.Recover != nil { // set recover block
 			// TODO(lijie): fix this for async function because of the block offset increase
 			fn.SetRecover(fn.Block(f.Recover.Index + nBlkOff))
@@ -260,7 +262,7 @@ func (p *context) compileFuncDecl(pkg llssa.Package, f *ssa.Function) (llssa.Fun
 			b := fn.NewBuilder()
 			b.SetBlockOffset(nBlkOff)
 			if async {
-				b.BeginAsync(fn)
+				b.BeginAsync(fn, entryBlk, allocBlk, cleanBlk, suspdBlk, trapBlk, beginBlk)
 			}
 			p.bvals = make(map[ssa.Value]llssa.Expr)
 			off := make([]int, len(f.Blocks))
diff --git a/ssa/coro.go b/ssa/coro.go
index b059e0dd..07f150d8 100644
--- a/ssa/coro.go
+++ b/ssa/coro.go
@@ -21,7 +21,6 @@ import (
 	"go/constant"
 	"go/token"
 	"go/types"
-	"log"
 )
 
 // declare void @llvm.coro.destroy(ptr <handle>)
@@ -382,6 +381,7 @@ func (b Builder) EndAsync() {
 }
 
 /*
+retPtr := malloc(sizeof(Promise))
 id := @llvm.coro.id(0, null, null, null)
 frameSize := @llvm.coro.size.i64()
 needAlloc := @llvm.coro.alloc(id)
@@ -395,7 +395,7 @@ frame := null
 hdl := @llvm.coro.begin(id, frame)
 *retPtr = hdl
 */
-func (b Builder) BeginAsync(fn Function) {
+func (b Builder) BeginAsync(fn Function, entryBlk, allocBlk, cleanBlk, suspdBlk, trapBlk, beginBlk BasicBlock) {
 	ty := fn.Type.RawType().(*types.Signature).Results().At(0).Type()
 	ptrTy, ok := ty.(*types.Pointer)
 	if !ok {
@@ -404,19 +404,12 @@ func (b Builder) BeginAsync(fn Function) {
 	promiseTy := b.Prog.Type(ptrTy.Elem(), InGo)
 
 	b.async = true
-	entryBlk := fn.Block(0)
-	allocBlk := fn.Block(1)
-	cleanBlk := fn.Block(2)
-	suspdBlk := fn.Block(3)
-	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")
 	promise := b.AllocZ(promiseSize).SetName("promise")
 	promise.Type = b.Prog.Pointer(promiseTy)
 	b.promise = promise
-	log.Printf("promise ptr: %v", promise.RawType())
 	align := b.Const(constant.MakeInt64(0), b.Prog.CInt()).SetName("align")
 	null := b.Const(nil, b.Prog.CIntPtr())
 	id := b.CoID(align, null, null, null).SetName("id")