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

The compiler needs to know whether a defer is in a loop to determine whether to allocate stack or heap memory.
Previously, this performed a DFS of the CFG every time a defer was found.
This resulted in time complexity jointly proportional to the number of defers and the number of blocks in the function.

Now, the compiler will instead use Tarjan's strongly connected components algorithm to find cycles in linear time.
The search is performed lazily, so this has minimal performance impact on functions without defers.

In order to implement Tarjan's SCC algorithm, additional state needed to be attached to the blocks.
I chose to merge all of the per-block state into a single slice to simplify memory management.

Author: niaow

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
@@ -1221,13 +1234,17 @@ func (b *builder) createFunctionStart(intrinsic bool) {
 		// manually.
 		entryBlock = b.ctx.AddBasicBlock(b.llvmFn, "entry")
 	} 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 +1340,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 +1402,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 +1516,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,88 @@ 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.
+func (b *builder) isInLoop() bool {
+	if b.currentBlockInfo.tarjan.index == 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.
+	assignedIndex := b.tarjanIndex + 1
+	b.tarjanIndex = assignedIndex
+
+	// Apply the new index.
+	blockIndex := block.Index
+	node := &b.blockInfo[blockIndex].tarjan
+	node.index = assignedIndex
+	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
+
+		if successorNode.index == 0 {
+			// This node has not yet been visisted.
+			b.strongConnect(successor)
+		} else if !successorNode.onStack {
+			// This node is has been visited, but is in a different SCC.
+			// Ignore it.
+			continue
+		}
+		if successorNode.lowLink < node.lowLink {
+			node.lowLink = successorNode.lowLink
+		}
+	}
 
-			// add to queue and checked map
-			queue = append(queue, pred)
-			checked[pred] = struct{}{}
+	if node.lowLink == node.index {
+		// 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) {
+			// 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
 	}
+}
 
-	return false
+type tarjanNode struct {
+	index   uint
+	lowLink uint
+	onStack bool
+	cyclic  bool
 }
 
 // createDefer emits a single defer instruction, to be run when this function
@@ -410,7 +457,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