compiler: use Tarjan's SCC algorithm to detect loops for defer

compiler/asserts.go

@@ -245,7 +245,7 @@ func (b *builder) createRuntimeAssert(assert llvm.Value, blockPrefix, assertFunc
 	// current insert position.
 	faultBlock := b.ctx.AddBasicBlock(b.llvmFn, blockPrefix+".throw")
 	nextBlock := b.insertBasicBlock(blockPrefix + ".next")
-	b.blockExits[b.currentBlock] = nextBlock // adjust outgoing block for phi nodes
+	b.currentBlockInfo.exit = nextBlock // adjust outgoing block for phi nodes
 
 	// Now branch to the out-of-bounds or the regular block.
 	b.CreateCondBr(assert, faultBlock, nextBlock)

compiler/compiler.go

@@ -152,10 +152,12 @@ type builder struct {
 	llvmFnType        llvm.Type
 	llvmFn            llvm.Value
 	info              functionInfo
-	locals            map[ssa.Value]llvm.Value            // local variables
-	blockEntries      map[*ssa.BasicBlock]llvm.BasicBlock // a *ssa.BasicBlock may be split up
-	blockExits        map[*ssa.BasicBlock]llvm.BasicBlock // these are the exit blocks
+	locals            map[ssa.Value]llvm.Value // local variables
+	blockInfo         []blockInfo
 	currentBlock      *ssa.BasicBlock
+	currentBlockInfo  *blockInfo
+	tarjanStack       []uint
+	tarjanIndex       uint
 	phis              []phiNode
 	deferPtr          llvm.Value
 	deferFrame        llvm.Value
@@ -187,11 +189,22 @@ func newBuilder(c *compilerContext, irbuilder llvm.Builder, f *ssa.Function) *bu
 		info:            c.getFunctionInfo(f),
 		locals:          make(map[ssa.Value]llvm.Value),
 		dilocals:        make(map[*types.Var]llvm.Metadata),
-		blockEntries:    make(map[*ssa.BasicBlock]llvm.BasicBlock),
-		blockExits:      make(map[*ssa.BasicBlock]llvm.BasicBlock),
 	}
 }
 
+type blockInfo struct {
+	// entry is the LLVM basic block corresponding to the start of this *ssa.Block.
+	entry llvm.BasicBlock
+
+	// exit is the LLVM basic block corresponding to the end of this *ssa.Block.
+	// It will be different than entry if any of the block's instructions contain internal branches.
+	exit llvm.BasicBlock
+
+	// tarjan holds state for applying Tarjan's strongly connected components algorithm to the CFG.
+	// This is used by defer.go to determine whether to stack- or heap-allocate defer data.
+	tarjan tarjanNode
+}
+
 type deferBuiltin struct {
 	callName string
 	pos      token.Pos
@@ -1220,14 +1233,29 @@ func (b *builder) createFunctionStart(intrinsic bool) {
 		// intrinsic (like an atomic operation). Create the entry block
 		// manually.
 		entryBlock = b.ctx.AddBasicBlock(b.llvmFn, "entry")
+		// Intrinsics may create internal branches (e.g. nil checks).
+		// They will attempt to access b.currentBlockInfo to update the exit block.
+		// Create some fake block info for them to access.
+		blockInfo := []blockInfo{
+			{
+				entry: entryBlock,
+				exit:  entryBlock,
+			},
+		}
+		b.blockInfo = blockInfo
+		b.currentBlockInfo = &blockInfo[0]
 	} else {
+		blocks := b.fn.Blocks
+		blockInfo := make([]blockInfo, len(blocks))
 		for _, block := range b.fn.DomPreorder() {
+			info := &blockInfo[block.Index]
 			llvmBlock := b.ctx.AddBasicBlock(b.llvmFn, block.Comment)
-			b.blockEntries[block] = llvmBlock
-			b.blockExits[block] = llvmBlock
+			info.entry = llvmBlock
+			info.exit = llvmBlock
 		}
+		b.blockInfo = blockInfo
 		// Normal functions have an entry block.
-		entryBlock = b.blockEntries[b.fn.Blocks[0]]
+		entryBlock = blockInfo[0].entry
 	}
 	b.SetInsertPointAtEnd(entryBlock)
 
@@ -1323,8 +1351,9 @@ func (b *builder) createFunction() {
 		if b.DumpSSA {
 			fmt.Printf("%d: %s:\n", block.Index, block.Comment)
 		}
-		b.SetInsertPointAtEnd(b.blockEntries[block])
 		b.currentBlock = block
+		b.currentBlockInfo = &b.blockInfo[block.Index]
+		b.SetInsertPointAtEnd(b.currentBlockInfo.entry)
 		for _, instr := range block.Instrs {
 			if instr, ok := instr.(*ssa.DebugRef); ok {
 				if !b.Debug {
@@ -1384,7 +1413,7 @@ func (b *builder) createFunction() {
 		block := phi.ssa.Block()
 		for i, edge := range phi.ssa.Edges {
 			llvmVal := b.getValue(edge, getPos(phi.ssa))
-			llvmBlock := b.blockExits[block.Preds[i]]
+			llvmBlock := b.blockInfo[block.Preds[i].Index].exit
 			phi.llvm.AddIncoming([]llvm.Value{llvmVal}, []llvm.BasicBlock{llvmBlock})
 		}
 	}
@@ -1498,11 +1527,11 @@ func (b *builder) createInstruction(instr ssa.Instruction) {
 	case *ssa.If:
 		cond := b.getValue(instr.Cond, getPos(instr))
 		block := instr.Block()
-		blockThen := b.blockEntries[block.Succs[0]]
-		blockElse := b.blockEntries[block.Succs[1]]
+		blockThen := b.blockInfo[block.Succs[0].Index].entry
+		blockElse := b.blockInfo[block.Succs[1].Index].entry
 		b.CreateCondBr(cond, blockThen, blockElse)
 	case *ssa.Jump:
-		blockJump := b.blockEntries[instr.Block().Succs[0]]
+		blockJump := b.blockInfo[instr.Block().Succs[0].Index].entry
 		b.CreateBr(blockJump)
 	case *ssa.MapUpdate:
 		m := b.getValue(instr.Map, getPos(instr))

compiler/defer.go

@@ -100,7 +100,7 @@ func (b *builder) createLandingPad() {
 
 	// Continue at the 'recover' block, which returns to the parent in an
 	// appropriate way.
-	b.CreateBr(b.blockEntries[b.fn.Recover])
+	b.CreateBr(b.blockInfo[b.fn.Recover.Index].entry)
 }
 
 // Create a checkpoint (similar to setjmp). This emits inline assembly that
@@ -234,41 +234,108 @@ func (b *builder) createInvokeCheckpoint() {
 	continueBB := b.insertBasicBlock("")
 	b.CreateCondBr(isZero, continueBB, b.landingpad)
 	b.SetInsertPointAtEnd(continueBB)
-	b.blockExits[b.currentBlock] = continueBB
+	b.currentBlockInfo.exit = continueBB
 }
 
-// isInLoop checks if there is a path from a basic block to itself.
-func isInLoop(start *ssa.BasicBlock) bool {
-	// Use a breadth-first search to scan backwards through the block graph.
-	queue := []*ssa.BasicBlock{start}
-	checked := map[*ssa.BasicBlock]struct{}{}
-
-	for len(queue) > 0 {
-		// pop a block off of the queue
-		block := queue[len(queue)-1]
-		queue = queue[:len(queue)-1]
-
-		// Search through predecessors.
-		// Searching backwards means that this is pretty fast when the block is close to the start of the function.
-		// Defers are often placed near the start of the function.
-		for _, pred := range block.Preds {
-			if pred == start {
-				// cycle found
-				return true
-			}
+// isInLoop checks if there is a path from the current block to itself.
+// Use Tarjan's strongly connected components algorithm to search for cycles.
+// A one-node SCC is a cycle iff there is an edge from the node to itself.
+// A multi-node SCC is always a cycle.
+// https://en.wikipedia.org/wiki/Tarjan%27s_strongly_connected_components_algorithm
+func (b *builder) isInLoop() bool {
+	if b.currentBlockInfo.tarjan.lowLink == 0 {
+		b.strongConnect(b.currentBlock)
+	}
+	return b.currentBlockInfo.tarjan.cyclic
+}
 
-			if _, ok := checked[pred]; ok {
-				// block already checked
-				continue
-			}
+func (b *builder) strongConnect(block *ssa.BasicBlock) {
+	// Assign a new index.
+	// Indices start from 1 so that 0 can be used as a sentinel.
+	assignedIndex := b.tarjanIndex + 1
+	b.tarjanIndex = assignedIndex
+
+	// Apply the new index.
+	blockIndex := block.Index
+	node := &b.blockInfo[blockIndex].tarjan
+	node.lowLink = assignedIndex
+
+	// Push the node onto the stack.
+	node.onStack = true
+	b.tarjanStack = append(b.tarjanStack, uint(blockIndex))
+
+	// Process the successors.
+	for _, successor := range block.Succs {
+		// Look up the successor's state.
+		successorIndex := successor.Index
+		if successorIndex == blockIndex {
+			// Handle a self-cycle specially.
+			node.cyclic = true
+			continue
+		}
+		successorNode := &b.blockInfo[successorIndex].tarjan
 
-			// add to queue and checked map
-			queue = append(queue, pred)
-			checked[pred] = struct{}{}
+		switch {
+		case successorNode.lowLink == 0:
+			// This node has not yet been visisted.
+			b.strongConnect(successor)
+
+		case !successorNode.onStack:
+			// This node has been visited, but is in a different SCC.
+			// Ignore it, and do not update lowLink.
+			continue
+		}
+
+		// Update the lowLink index.
+		// This always uses the min-of-lowlink instead of using index in the on-stack case.
+		// This is done for two reasons:
+		// 1. The lowLink update can be shared between the new-node and on-stack cases.
+		// 2. The assigned index does not need to be saved - it is only needed for root node detection.
+		if successorNode.lowLink < node.lowLink {
+			node.lowLink = successorNode.lowLink
 		}
 	}
 
-	return false
+	if node.lowLink == assignedIndex {
+		// This is a root node.
+		// Pop the SCC off the stack.
+		stack := b.tarjanStack
+		top := stack[len(stack)-1]
+		stack = stack[:len(stack)-1]
+		blocks := b.blockInfo
+		topNode := &blocks[top].tarjan
+		topNode.onStack = false
+
+		if top != uint(blockIndex) {
+			// The root node is not the only node in the SCC.
+			// Mark all nodes in this SCC as cyclic.
+			topNode.cyclic = true
+			for top != uint(blockIndex) {
+				top = stack[len(stack)-1]
+				stack = stack[:len(stack)-1]
+				topNode = &blocks[top].tarjan
+				topNode.onStack = false
+				topNode.cyclic = true
+			}
+		}
+
+		b.tarjanStack = stack
+	}
+}
+
+// tarjanNode holds per-block state for isInLoop and strongConnect.
+type tarjanNode struct {
+	// lowLink is the index of the first visited node that is reachable from this block.
+	// The lowLink indices are assigned by the SCC search, and do not correspond to b.Index.
+	// A lowLink of 0 is used as a sentinel to mark a node which has not yet been visited.
+	lowLink uint
+
+	// onStack tracks whether this node is currently on the SCC search stack.
+	onStack bool
+
+	// cyclic indicates whether this block is in a loop.
+	// If lowLink is 0, strongConnect must be called before reading this field.
+	cyclic bool
 }
 
 // createDefer emits a single defer instruction, to be run when this function
@@ -410,7 +477,10 @@ func (b *builder) createDefer(instr *ssa.Defer) {
 
 	// Put this struct in an allocation.
 	var alloca llvm.Value
-	if !isInLoop(instr.Block()) {
+	if instr.Block() != b.currentBlock {
+		panic("block mismatch")
+	}
+	if !b.isInLoop() {
 		// This can safely use a stack allocation.
 		alloca = llvmutil.CreateEntryBlockAlloca(b.Builder, deferredCallType, "defer.alloca")
 	} else {

compiler/interface.go

@@ -737,7 +737,7 @@ func (b *builder) createTypeAssert(expr *ssa.TypeAssert) llvm.Value {
 	prevBlock := b.GetInsertBlock()
 	okBlock := b.insertBasicBlock("typeassert.ok")
 	nextBlock := b.insertBasicBlock("typeassert.next")
-	b.blockExits[b.currentBlock] = nextBlock // adjust outgoing block for phi nodes
+	b.currentBlockInfo.exit = nextBlock // adjust outgoing block for phi nodes
 	b.CreateCondBr(commaOk, okBlock, nextBlock)
 
 	// Retrieve the value from the interface if the type assert was