Skip to content

Commit

Permalink
Refactor CFG preprocess logic (#249)
Browse files Browse the repository at this point in the history 
This PR refactors the CFG preprocess logic. Specifically, we have
extracted the rich check effect construction out of the preprocessing
logic such that it is now in its standalone file. This offers a lot
better modularity for the already-enormous `assertiontree` package.

Now in `BackpropAcrossFunc` function we are calling preprocess and
rich-check-effect-generating functions separately.
  • Loading branch information
@yuxincs
yuxincs authored May 24, 2024
1 parent a27654d commit 2ae4ebe
Show file tree
Hide file tree
Showing 3 changed files with 261 additions and 255 deletions.
13 changes: 10 additions & 3 deletions assertion/function/assertiontree/backprop.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -814,11 +814,18 @@ func computePostOrder(blocks []*cfg.Block) []int {
// BackpropAcrossFunc is the main driver of the backpropagation, it takes a function declaration
// with accompanying CFG, and back-propagates a tree of assertions across it to generate, at entry
// to the function, the set of assertions that must hold to avoid possible nil flow errors.
func BackpropAcrossFunc(ctx context.Context, pass *analysis.Pass, decl *ast.FuncDecl,
functionContext FunctionContext, graph *cfg.CFG) ([]annotation.FullTrigger, int, int, error) {
func BackpropAcrossFunc(
ctx context.Context,
pass *analysis.Pass,
decl *ast.FuncDecl,
functionContext FunctionContext,
graph *cfg.CFG,
) ([]annotation.FullTrigger, int, int, error) {
// We transform the CFG to have it reflect the implicit control flow that happens
// inside short-circuiting boolean expressions.
graph, richCheckBlocks, exprNonceMap := preprocess(graph, functionContext)
graph = preprocess(graph, functionContext.funcDecl, functionContext.pass)
richCheckBlocks, exprNonceMap := genInitialRichCheckEffects(graph, functionContext)
richCheckBlocks = propagateRichChecks(graph, richCheckBlocks)
blocks, preprocessing := blocksAndPreprocessingFromCFG(pass, graph, richCheckBlocks)

// The assertion nodes for each block and an array of bools to indicate whether each block is
Expand Down
256 changes: 4 additions & 252 deletions assertion/function/assertiontree/preprocess_blocks.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,24 +30,22 @@ import (
//
// The returned RichCheckEffect slices represent the RichCheckEffects present at
// the _end_ of each block
func preprocess(graph *cfg.CFG, fc FunctionContext) (*cfg.CFG, [][]RichCheckEffect, util.ExprNonceMap) {
func preprocess(graph *cfg.CFG, funcDecl *ast.FuncDecl, pass *analysis.Pass) *cfg.CFG {
// The ASTs and CFGs are shared across all analyzers in the nogo framework, so we should never
// modify them directly. Here, we make a copy of the graph (and all blocks in it) and modify
// the copied graph instead.
graph = copyGraph(graph)
restructureBlocks(graph, fc.pass)
richCheckBlocks, exprNonceMap := genInitialRichCheckEffects(graph, fc)
richCheckBlocks = propagateRichChecks(graph, richCheckBlocks)
restructureBlocks(graph, pass)

// Next, we need to re-insert information that is lost during CFG build for *ast.RangeStmt
// and *ast.SwitchStmt by iterating through all blocks. This requires knowing the links between
// the nodes contained within a block to their parents (*ast.RangeStmt or *ast.SwitchStmt nodes).
// So, here establish the link and then do the work.
rangeChildren, switchChildren := collectChildren(fc.funcDecl)
rangeChildren, switchChildren := collectChildren(funcDecl)
markRangeStatements(graph, rangeChildren)
markSwitchStatements(graph, switchChildren)

return graph, richCheckBlocks, exprNonceMap
return graph
}

// copyGraph makes a semi-deep copy of the CFG and returns the copied graph. Note that only the
Expand Down Expand Up @@ -91,61 +89,6 @@ func copyGraph(graph *cfg.CFG) *cfg.CFG {
return newGraph
}

// stripNoops returns a copy of the passed slice `effects`, minus any no-ops
func stripNoops(effects []RichCheckEffect) []RichCheckEffect {
var strippedEffects []RichCheckEffect

for _, effect := range effects {
if !effect.isNoop() {
strippedEffects = append(strippedEffects, effect)
}
}

return strippedEffects
}

// genInitialRichCheckEffects computes an initial array of RichCheckEffect slices for each block,
// not doing any propagation over the CFG except for within each block to track nodes
// that create RichCheckEffects (such as `v, ok := mp[k]`) and make sure it isn't invalidated
// (such as by `ok = true`) before the end of the block.
//
// The returned RichCheckEffect slices represent the RichCheckEffects present at
// the _end_ of each block.
//
// Important: do not duplicate any pointers: each returned RichCheckEffect should be a unique object
func genInitialRichCheckEffects(graph *cfg.CFG, functionContext FunctionContext) (
[][]RichCheckEffect, util.ExprNonceMap) {
richCheckBlocks := make([][]RichCheckEffect, len(graph.Blocks))
nonceGenerator := util.NewGuardNonceGenerator()

// There is no canonical instance of RootAssertionNode until backpropAcrossFunc returns.
// We use a temporary root here as a means to pass contextual information like the function
// declaration and analysis pass.
rootNode := newRootAssertionNode(nonceGenerator.GetExprNonceMap(), functionContext)
for i, block := range graph.Blocks {
var richCheckEffects []RichCheckEffect
for _, node := range block.Nodes {

// invalidate any richCheckEffects that this node invalidates
for j, effect := range richCheckEffects {
if effect.isInvalidatedBy(node) {
richCheckEffects[j] = RichCheckNoop{}
}
}

// check if this node produces a new richCheckEffect
if effects, ok := RichCheckFromNode(rootNode, nonceGenerator, node); ok {
richCheckEffects = append(richCheckEffects, effects...)
}
}
// richCheckEffects is now fully populated

// strip out noops and write into richCheckBlocks
richCheckBlocks[i] = stripNoops(richCheckEffects)
}
return richCheckBlocks, nonceGenerator.GetExprNonceMap()
}

// This function restructures a cfg to reflect short-circuiting and other interesting semantics:
//
// It performs the following short-circuiting:
Expand Down Expand Up @@ -554,194 +497,3 @@ func markSwitchStatements(graph *cfg.CFG, switchChildren map[ast.Node]*ast.Switc
}
}
}

func mergeSlices(useDeepEquality bool, left []RichCheckEffect, rights ...[]RichCheckEffect) []RichCheckEffect {
var eq func(first, second RichCheckEffect) bool
if useDeepEquality {
eq = func(first, second RichCheckEffect) bool {
return first.equals(second)
}
} else {
eq = func(first, second RichCheckEffect) bool {
return first == second
}
}
var out []RichCheckEffect
addToOut := func(effect RichCheckEffect) {
for _, outEffect := range out {
if eq(outEffect, effect) {
return
}
}
out = append(out, effect)
}
for _, l := range left {
addToOut(l)
}
for _, right := range rights {
for _, r := range right {
addToOut(r)
}
}
return out
}

func genPreds(graph *cfg.CFG) [][]int32 {
out := make([][]int32, len(graph.Blocks))
for _, block := range graph.Blocks {
if block.Live {
for _, succ := range block.Succs {
out[succ.Index] = append(out[succ.Index], block.Index)
}
}
}
return out
}

// weakPropagateRichChecks performs a simple form of propagation of rich checks: for each effect, it
// figures out which blocks are reachable from the block it was declared in.
//
// The results are returned as a map from `RichCheckEffect`s to arrays of booleans, representing for
// each block whether it is reached by the block that effect is declared in
func weakPropagateRichChecks(graph *cfg.CFG, richCheckBlocks [][]RichCheckEffect) map[RichCheckEffect][]bool {
reachability := make(map[RichCheckEffect][]bool)
for blockNum := range richCheckBlocks {
for _, check := range richCheckBlocks[blockNum] {
newCheck := make([]bool, len(richCheckBlocks))
newCheck[blockNum] = true // mark each check as reachable in its declaring block
reachability[check] = newCheck
}
}
done := false
for !done {
done = true
for blockNum := range richCheckBlocks {
for _, reachable := range reachability {
if reachable[blockNum] {
for _, nextBlock := range graph.Blocks[blockNum].Succs {
if !reachable[nextBlock.Index] {
reachable[nextBlock.Index] = true
done = false
}
}
}
}
}
}
return reachability
}

// propagateRichChecks takes an initial array richCheckBlocks and flows all of its contained checks
// forwards through the CFG as long as they are not invalidated. A check created by a node in block A
// is determined to flow to block B if every path from A to B does not invalidate the check. We capture
// this criterion by first calling the function weakPropagateRichChecks above to do reachability
// propagation without any knowledge of check invalidation. The real propagation done in this function
// then tempers its computation of checks at a given block via intersection at control flow points by
// including exactly those checks that are present in every predecessor of the block that is reachable
// from the originator block of the check.
func propagateRichChecks(graph *cfg.CFG, richCheckBlocks [][]RichCheckEffect) [][]RichCheckEffect {
n := len(graph.Blocks)
if len(richCheckBlocks) != n {
panic(fmt.Sprintf("richCheckBlocks (len %d) and graph.blocks (len %d) out of "+
"sync - fix generation pass in preprocess_blocks.go", len(richCheckBlocks), n))
}

effectReaches := weakPropagateRichChecks(graph, richCheckBlocks)

currBlocks := richCheckBlocks
nextBlocks := make([][]RichCheckEffect, n)

preds := genPreds(graph)
roundCount := 0

done := false

for !done {

done = true

for i := range preds {

// predRichCheckEffects will be populated with all the rich bool effects that flow
// into this block from one of its 0 or more predecessors
var predRichCheckEffects []RichCheckEffect

if len(preds[i]) >= 1 {
reachingEffects := make(map[RichCheckEffect]bool)

for _, predIndex := range preds[i] {
for _, effect := range currBlocks[predIndex] {
// for each effect in a predecessor, mark it as `true` in `reachingEffects`
// - performing a merge
reachingEffects[effect] = true
}
}

for _, predIndex := range preds[i] {
maskingEffects := make(map[RichCheckEffect]bool)
for effect := range reachingEffects {
if blocksEffectReaches, ok := effectReaches[effect]; ok &&
blocksEffectReaches[predIndex] {
maskingEffects[effect] = true
}
}
for _, effect := range currBlocks[predIndex] {
if maskingEffects[effect] {
maskingEffects[effect] = false
}
}
for effect, present := range maskingEffects {
if present {
reachingEffects[effect] = false
}
}
}

predRichCheckEffects = make([]RichCheckEffect, 0)

for effect := range reachingEffects {
if reachingEffects[effect] {
predRichCheckEffects = append(predRichCheckEffects, effect)
}
}

// This code performs a simple merge instead - but this is very unsound and NOT right
// predRichCheckEffects =
// append(make([]RichCheckEffect, 0, len(currBlocks[preds[i][0]])),
// currBlocks[preds[i][0]]...)
//
// for _, predNum := range preds[i][1:] {
// predRichCheckEffects = mergeSlices(false, predRichCheckEffects, currBlocks[predNum])
// }

for _, node := range graph.Blocks[i].Nodes {
// invalidate any richCheckEffects that this node invalidates
for j, effect := range predRichCheckEffects {
if effect.isInvalidatedBy(node) {
predRichCheckEffects[j] = RichCheckNoop{}
}
}
}
}

nextBlocks[i] = mergeSlices(false, currBlocks[i], stripNoops(predRichCheckEffects))
if len(nextBlocks[i]) > len(currBlocks[i]) {
done = false
}
}

currBlocks = nextBlocks
nextBlocks = make([][]RichCheckEffect, n)

roundCount++

checkCFGFixedPointRuntime("RichCheckEffect Forwards Propagation", roundCount, n)
}

// this strips duplicates from the RichCheckEffect slices
for i := range currBlocks {
currBlocks[i] = mergeSlices(true, currBlocks[i])
}

return currBlocks
}
Loading

0 comments on commit 2ae4ebe

Please sign in to comment.