From cfca98512af0bb5aeae9a34b3477047fdacde7f9 Mon Sep 17 00:00:00 2001
From: xushiwei <x@goplus.org>
Date: Mon, 3 Jun 2024 02:43:04 +0800
Subject: [PATCH] llgo/ssa: DeferAlways/DeferInCond/DeferInLoop

---
 _demo/defer/defer.go |  1 -
 cl/compile.go        | 37 ++++++++++++++++++++++++++++++++-----
 ssa/expr.go          |  8 +++++---
 ssa/stmt_builder.go  | 39 ++++++++++++++++++++++++++-------------
 4 files changed, 63 insertions(+), 22 deletions(-)

diff --git a/_demo/defer/defer.go b/_demo/defer/defer.go
index 82d34e31..cd95da67 100644
--- a/_demo/defer/defer.go
+++ b/_demo/defer/defer.go
@@ -12,7 +12,6 @@ func main() {
 		defer println(s)
 	} else {
 		defer println("world")
-		return
 	}
 	defer println("bye")
 }
diff --git a/cl/compile.go b/cl/compile.go
index 7ace3f5e..2a8b6fce 100644
--- a/cl/compile.go
+++ b/cl/compile.go
@@ -148,8 +148,11 @@ type context struct {
 	loaded map[*types.Package]*pkgInfo // loaded packages
 	bvals  map[ssa.Value]llssa.Expr    // block values
 	vargs  map[*ssa.Alloc][]llssa.Expr // varargs
-	inits  []func()
-	phis   []func()
+
+	blkInfos []blockInfo
+
+	inits []func()
+	phis  []func()
 }
 
 func (p *context) inMain(instr ssa.Instruction) bool {
@@ -277,10 +280,16 @@ func (p *context) compileFuncDecl(pkg llssa.Package, f *ssa.Function) (llssa.Fun
 			for i, block := range f.Blocks {
 				off[i] = p.compilePhis(b, block)
 			}
-			for i, block := range f.Blocks {
+			p.blkInfos = blockInfos(f.Blocks)
+			i := 0
+			for {
+				block := f.Blocks[i]
 				doMainInit := (i == 0 && name == "main")
 				doModInit := (i == 1 && f.Name() == "init" && sig.Recv() == nil)
 				p.compileBlock(b, block, off[i], doMainInit, doModInit)
+				if i = p.blkInfos[i].next; i < 0 {
+					break
+				}
 			}
 			for _, phi := range p.phis {
 				phi()
@@ -291,6 +300,24 @@ func (p *context) compileFuncDecl(pkg llssa.Package, f *ssa.Function) (llssa.Fun
 	return fn, nil, goFunc
 }
 
+type blockInfo struct {
+	kind llssa.DoAction
+	next int
+}
+
+func blockInfos(blks []*ssa.BasicBlock) []blockInfo {
+	n := len(blks)
+	infos := make([]blockInfo, n)
+	for i := range blks {
+		next := i + 1
+		if next >= n {
+			next = -1
+		}
+		infos[i] = blockInfo{kind: llssa.DeferInCond, next: next}
+	}
+	return infos
+}
+
 // funcOf returns a function by name and set ftype = goFunc, cFunc, etc.
 // or returns nil and set ftype = llgoCstr, llgoAlloca, llgoUnreachable, etc.
 func (p *context) funcOf(fn *ssa.Function) (aFn llssa.Function, pyFn llssa.PyObjRef, ftype int) {
@@ -516,7 +543,7 @@ func (p *context) compilePhis(b llssa.Builder, block *ssa.BasicBlock) int {
 			for n < ninstr && isPhi(block.Instrs[n]) {
 				n++
 			}
-			rets := make([]llssa.Expr, n)
+			rets := make([]llssa.Expr, n) // TODO(xsw): check to remove this
 			for i := 0; i < n; i++ {
 				iv := block.Instrs[i].(*ssa.Phi)
 				rets[i] = p.compilePhi(b, iv)
@@ -798,7 +825,7 @@ func (p *context) compileInstr(b llssa.Builder, instr ssa.Instruction) {
 		val := p.compileValue(b, v.Value)
 		b.MapUpdate(m, key, val)
 	case *ssa.Defer:
-		p.call(b, llssa.Defer, &v.Call)
+		p.call(b, p.blkInfos[v.Block().Index].kind, &v.Call)
 	case *ssa.Go:
 		p.call(b, llssa.Go, &v.Call)
 	case *ssa.RunDefers:
diff --git a/ssa/expr.go b/ssa/expr.go
index 0c81e2a3..c085054d 100644
--- a/ssa/expr.go
+++ b/ssa/expr.go
@@ -803,8 +803,10 @@ type DoAction int
 
 const (
 	Call DoAction = iota
-	Defer
 	Go
+	DeferAlways // defer statement executes always
+	DeferInCond // defer statement executes in a conditional block
+	DeferInLoop // defer statement executes in a loop block
 )
 
 // Do call a function with an action.
@@ -812,10 +814,10 @@ func (b Builder) Do(da DoAction, fn Expr, args ...Expr) (ret Expr) {
 	switch da {
 	case Call:
 		return b.Call(fn, args...)
-	case Defer:
-		b.Defer(fn, args...)
 	case Go:
 		b.Go(fn, args...)
+	default:
+		b.Defer(da, fn, args...)
 	}
 	return
 }
diff --git a/ssa/stmt_builder.go b/ssa/stmt_builder.go
index 36da0a9d..9f59e656 100644
--- a/ssa/stmt_builder.go
+++ b/ssa/stmt_builder.go
@@ -216,24 +216,37 @@ func (b Builder) getDefer() *aDefer {
 }
 
 // Defer emits a defer instruction.
-func (b Builder) Defer(fn Expr, args ...Expr) {
+func (b Builder) Defer(kind DoAction, fn Expr, args ...Expr) {
 	if debugInstr {
 		logCall("Defer", fn, args)
 	}
-	prog := b.Prog
-	self := b.getDefer()
-	next := self.nextBit
-	self.nextBit++
-	bits := b.Load(self.bitsPtr)
-	nextbit := prog.Val(uintptr(1 << next))
-	b.Store(self.bitsPtr, b.BinOp(token.OR, bits, nextbit))
-
+	var prog Program
+	var nextbit Expr
+	var self = b.getDefer()
+	switch kind {
+	case DeferInCond:
+		prog = b.Prog
+		next := self.nextBit
+		self.nextBit++
+		bits := b.Load(self.bitsPtr)
+		nextbit = prog.Val(uintptr(1 << next))
+		b.Store(self.bitsPtr, b.BinOp(token.OR, bits, nextbit))
+	case DeferAlways:
+		// nothing to do
+	default:
+		panic("todo: DeferInLoop is not supported")
+	}
 	self.stmts = append(self.stmts, func(bits Expr) {
-		zero := prog.Val(uintptr(0))
-		has := b.BinOp(token.NEQ, b.BinOp(token.AND, bits, nextbit), zero)
-		b.IfThen(has, func() {
+		switch kind {
+		case DeferInCond:
+			zero := prog.Val(uintptr(0))
+			has := b.BinOp(token.NEQ, b.BinOp(token.AND, bits, nextbit), zero)
+			b.IfThen(has, func() {
+				b.Call(fn, args...)
+			})
+		case DeferAlways:
 			b.Call(fn, args...)
-		})
+		}
 	})
 }