Change cfg to be more efficient by not calculating all paths

Author: Amit Davidson

Diff for: domain/BlockState.go

@@ -34,11 +34,9 @@ func (existingBlock *BlockState) AddResult(newBlock *BlockState, shouldMergeLock
 // Merge child B unto A:
 // A -> B
 // Will Merge B unto A
-func (existingBlock *BlockState) MergeChildBlock(newBlock *BlockState, shouldUpdateGALockset bool) {
-	if shouldUpdateGALockset {
-		for _, guardedAccess := range newBlock.GuardedAccesses {
-			guardedAccess.Lockset.UpdateLockSet(existingBlock.Lockset.Locks, existingBlock.Lockset.Unlocks)
-		}
+func (existingBlock *BlockState) MergeChildBlock(newBlock *BlockState) {
+	for _, guardedAccess := range newBlock.GuardedAccesses {
+		guardedAccess.Lockset.UpdateLockSet(existingBlock.Lockset.Locks, existingBlock.Lockset.Unlocks)
 	}
 	existingBlock.GuardedAccesses = append(existingBlock.GuardedAccesses, newBlock.GuardedAccesses...)
 	existingBlock.DeferredFunctions.MergeStacks(newBlock.DeferredFunctions)

Diff for: ssaUtils/CFG.go

@@ -7,81 +7,51 @@ import (
 )
 
 type CFG struct {
-	visitedBlocksStack *stacks.BasicBlockStack
+	visitedBlocksStack *stacks.BlockMap
 
 	ComputedBlocks      map[int]*domain.BlockState
 	ComputedDeferBlocks map[int]*domain.BlockState
-	calculatedState     *domain.BlockState
 }
 
 func newCFG() *CFG {
 	return &CFG{
-		visitedBlocksStack:  stacks.NewBasicBlockStack(),
+		visitedBlocksStack:  stacks.NewBlockMap(),
 		ComputedBlocks:      make(map[int]*domain.BlockState),
 		ComputedDeferBlocks: make(map[int]*domain.BlockState),
 	}
 }
 
-func (cfg *CFG) calculateFunctionStatePathSensitive(context *domain.Context, block *ssa.BasicBlock) {
-	firstBlock := &ssa.BasicBlock{Index: -1, Succs: []*ssa.BasicBlock{block}}
-	cfg.traverseGraph(context, firstBlock)
-}
+func (cfg *CFG) CalculateFunctionStatePathSensitive(context *domain.Context, block *ssa.BasicBlock) *domain.BlockState {
+	cfg.visitedBlocksStack.Add(block)
+	defer cfg.visitedBlocksStack.Remove(block)
+	cfg.calculateBlockState(context, block)
+
+	// Regular flow
+	blockState := cfg.ComputedBlocks[block.Index]
 
-func (cfg *CFG) traverseGraph(context *domain.Context, block *ssa.BasicBlock) {
+	// recursion
+	var branchState *domain.BlockState
 	for _, nextBlock := range block.Succs {
-		cfg.calculateBlockState(context, block)
-		if len(nextBlock.Succs) == 0 || cfg.visitedBlocksStack.Contains(nextBlock) {
-			// If a return is reached or if a cycle of a loop is completed.
-			cfg.calculateBlockState(context, nextBlock)
-			cfg.visitedBlocksStack.Push(nextBlock)
-			cfg.CalculatePath()
-			cfg.visitedBlocksStack.Pop()
-		} else {
-			cfg.visitedBlocksStack.Push(nextBlock)
-			cfg.traverseGraph(context, nextBlock)
-			cfg.visitedBlocksStack.Pop()
+		// if it's a cycle we skip it
+		if cfg.visitedBlocksStack.Contains(nextBlock.Index) {
+			continue
 		}
-	}
-}
-
-func (cfg *CFG) CalculatePath() {
-	path := cfg.visitedBlocksStack.GetItems()
-	block := path[0]
-	state := cfg.ComputedBlocks[block.Index].Copy()
-	for _, nextBlock := range path[1:] {
-		nextState := cfg.ComputedBlocks[nextBlock.Index].Copy()
-		state.MergeChildBlock(nextState, true)
-	}
-
-	var firstDeferState *domain.BlockState
-	var firstDeferIndex int
-	for i := len(path) - 1; i >= 0; i-- {
-		deferIndex := path[i].Index
-		deferState, ok := cfg.ComputedDeferBlocks[deferIndex]
-		if ok {
-			firstDeferState = deferState
-			firstDeferIndex = i
-			break
+		retBlockState := cfg.CalculateFunctionStatePathSensitive(context, nextBlock)
+		if branchState == nil {
+			branchState = retBlockState.Copy()
+		} else {
+			branchState.MergeSiblingBlock(retBlockState)
 		}
 	}
-	if firstDeferState != nil {
-		deferStateCopy := firstDeferState.Copy()
-		for i := firstDeferIndex - 1; i >= 0; i-- {
-			nextState, ok := cfg.ComputedDeferBlocks[path[i].Index]
-			if !ok {
-				continue
-			}
-			nextStateCopy := nextState.Copy()
-			deferStateCopy.MergeChildBlock(nextStateCopy, true)
-		}
-		state.AddResult(deferStateCopy, true)
+	if branchState != nil {
+		blockState.MergeChildBlock(branchState)
 	}
 
-	if cfg.calculatedState == nil {
-		cfg.calculatedState = state
-	} else {
-		cfg.calculatedState.MergeSiblingBlock(state)
+	// Defer
+	if deferState, ok := cfg.ComputedDeferBlocks[block.Index]; ok {
+		blockState.MergeChildBlock(deferState)
 	}
+	return blockState
 }
 
 func (cfg *CFG) calculateBlockState(context *domain.Context, block *ssa.BasicBlock) {

Diff for: ssaUtils/Functions.go

@@ -171,35 +171,19 @@ func GetBlockSummary(context *domain.Context, block *ssa.BasicBlock) *domain.Blo
 
 func (cfg *CFG) runDefers(context *domain.Context, defers *stacks.CallCommonStack) *domain.BlockState {
 	calculatedState := domain.GetEmptyBlockState()
-	for {
-		deferFunction := defers.Pop()
+	defersItems := defers.GetItems()
+	for i := len(defersItems) - 1; i >= 0; i-- {
+		deferFunction := defersItems[i]
 		if deferFunction == nil {
 			break
 		}
 		retState := HandleCallCommon(context, deferFunction, deferFunction.Pos())
-		calculatedState.MergeChildBlock(retState, true)
+		calculatedState.MergeChildBlock(retState)
 	}
 	return calculatedState
 
 }
 
-func (cfg *CFG) calculateFunctionStatePathInsensitive(context *domain.Context, blocks []*ssa.BasicBlock) {
-	for _, block := range blocks {
-		cfg.calculateBlockState(context, block)
-		state := cfg.ComputedBlocks[block.Index]
-		if cfg.calculatedState == nil {
-			cfg.calculatedState = state.Copy()
-		} else {
-			cfg.calculatedState.MergeChildBlock(state, false)
-		}
-
-		deferBlock := cfg.ComputedDeferBlocks[block.Index]
-		if deferBlock != nil {
-			cfg.calculatedState.MergeChildBlock(deferBlock, false)
-		}
-	}
-}
-
 func HandleFunction(context *domain.Context, fn *ssa.Function) *domain.BlockState {
 	funcState := domain.GetEmptyBlockState()
 	if fn.Pkg == nil {
@@ -219,12 +203,6 @@ func HandleFunction(context *domain.Context, fn *ssa.Function) *domain.BlockStat
 		panic("Function is being iterated but wasn't found when iterating on program functions in preprocess")
 	}
 	cfg := newCFG()
-	cfg.calculateFunctionStatePathSensitive(context, fn.Blocks[0])
-	//if isContainingLocks {
-	//	cfg.calculateFunctionStatePathSensitive(context, fn.Blocks[0])
-	//} else {
-	//	cfg.calculateFunctionStatePathInsensitive(context, fn.Blocks)
-	//}
-
-	return cfg.calculatedState
+	calculatedState := cfg.CalculateFunctionStatePathSensitive(context, fn.Blocks[0])
+	return calculatedState
 }

Diff for: utils/stacks/basicBlockStack.go

@@ -4,78 +4,23 @@ import "golang.org/x/tools/go/ssa"
 
 // Used by CFG to traverse the graph. It uses both as a stack for traversal by order, and as a map to count occurrences
 // and fast retrieval of items.
-type BasicBlockStack struct {
-	stack     blocksStack
-	blocksMap blocksMap
-}
-
-func NewBasicBlockStack() *BasicBlockStack {
-	stack := make([]*ssa.BasicBlock, 0)
-	blocksMap := make(blocksMap)
-	basicBlockStack := &BasicBlockStack{stack: stack, blocksMap: blocksMap}
-	return basicBlockStack
-}
 
-func (s *BasicBlockStack) Push(v *ssa.BasicBlock) {
-	s.stack.Push(v)
-	if _, ok := s.blocksMap[v.Index]; !ok {
-		s.blocksMap[v.Index] = &basicBlockWithCount{block: v, count: 1}
-	} else {
-		s.blocksMap[v.Index].count += 1
-	}
-}
+type BlockMap map[int]struct{}
 
-func (s *BasicBlockStack) Pop() *ssa.BasicBlock {
-	v := s.stack.Pop()
-	if v == nil {
-		return nil
-	}
-	s.blocksMap[v.Index].count -= 1
-	if s.blocksMap[v.Index].count == 0 {
-		delete(s.blocksMap, v.Index)
-	}
-	return v
+func NewBlockMap() *BlockMap {
+	blocksMap := make(BlockMap)
+	return &blocksMap
 }
 
-func (s *BasicBlockStack) GetItems() []*ssa.BasicBlock {
-	return s.stack
+func (s *BlockMap) Add(v *ssa.BasicBlock) {
+	(*s)[v.Index] = struct{}{}
 }
 
-type basicBlockWithCount struct {
-	block *ssa.BasicBlock
-	count int
+func (s *BlockMap) Remove(v *ssa.BasicBlock) {
+	delete(*s, v.Index)
 }
 
-type blocksMap map[int]*basicBlockWithCount
-
-type blocksStack []*ssa.BasicBlock
-
-func (s *blocksStack) Push(v *ssa.BasicBlock) {
-	*s = append(*s, v)
-}
-
-func (s *blocksStack) Pop() *ssa.BasicBlock {
-	if len(*s) == 0 {
-		return nil
-	}
-	v := (*s)[len(*s)-1]
-	*s = (*s)[:len(*s)-1]
-	return v
-}
-
-
-// The function check if the block already exists in the stack and if it does, return true only if it appear twice or more.
-// This is used to handle cycles where a path might visit a node more then once. This way a cycle is allowed, but only once
-// to avoid infinite recursion. Finding all valid flows is hard so this is chosen as an approximation. There are edge
-//cases such as: A->B->D->C->D->E(end) where the graph is not circular, but the traversal won't reach to E since D
-//appear a second time and it will exit at that point. This flow is rare and can be achieved using gotos
-func (s *BasicBlockStack) Contains(v *ssa.BasicBlock) bool {
-	block, ok := s.blocksMap[v.Index]
-	if !ok {
-		return false
-	}
-	if block.count >= 1 {
-		return true
-	}
-	return false
+func (s *BlockMap) Contains(v int) bool {
+	_, ok := (*s)[v]
+	return ok
 }

Diff for: utils/stacks/callCommonStack.go

@@ -10,9 +10,7 @@ func NewCallCommonStack() *CallCommonStack {
 }
 
 func (s *CallCommonStack) GetItems() []*ssa.CallCommon {
-	tmp := make([]*ssa.CallCommon, len(*s))
-	copy(tmp, *s)
-	return tmp
+	return *s
 }
 
 func (s *CallCommonStack) Push(v *ssa.CallCommon) {