llgo/ssa: DeferAlways/DeferInCond/DeferInLoop

_demo/defer/defer.go

@@ -12,7 +12,6 @@ func main() {
 		defer println(s)
 	} else {
 		defer println("world")
-		return
 	}
 	defer println("bye")
 }

cl/compile.go

@@ -148,6 +148,9 @@ 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
+
+	blkInfos []blockInfo
+
 	inits []func()
 	phis  []func()
 }
@@ -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:

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
 }

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()
+	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))
+		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) {
+		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...)
+		}
 	})
 }