github.com/amarpal/go-tools@v0.0.0-20240422043104-40142f59f616/go/ir/irutil/loops.go (about)

     1  package irutil
     2  
     3  import "github.com/amarpal/go-tools/go/ir"
     4  
     5  type Loop struct{ *ir.BlockSet }
     6  
     7  func FindLoops(fn *ir.Function) []Loop {
     8  	if fn.Blocks == nil {
     9  		return nil
    10  	}
    11  	tree := fn.DomPreorder()
    12  	var sets []Loop
    13  	for _, h := range tree {
    14  		for _, n := range h.Preds {
    15  			if !h.Dominates(n) {
    16  				continue
    17  			}
    18  			// n is a back-edge to h
    19  			// h is the loop header
    20  			if n == h {
    21  				set := Loop{ir.NewBlockSet(len(fn.Blocks))}
    22  				set.Add(n)
    23  				sets = append(sets, set)
    24  				continue
    25  			}
    26  			set := Loop{ir.NewBlockSet(len(fn.Blocks))}
    27  			set.Add(h)
    28  			set.Add(n)
    29  			for _, b := range allPredsBut(n, h, nil) {
    30  				set.Add(b)
    31  			}
    32  			sets = append(sets, set)
    33  		}
    34  	}
    35  	return sets
    36  }
    37  
    38  func allPredsBut(b, but *ir.BasicBlock, list []*ir.BasicBlock) []*ir.BasicBlock {
    39  outer:
    40  	for _, pred := range b.Preds {
    41  		if pred == but {
    42  			continue
    43  		}
    44  		for _, p := range list {
    45  			// TODO improve big-o complexity of this function
    46  			if pred == p {
    47  				continue outer
    48  			}
    49  		}
    50  		list = append(list, pred)
    51  		list = allPredsBut(pred, but, list)
    52  	}
    53  	return list
    54  }