--- /dev/null
+package irutil
+
+import "honnef.co/go/tools/go/ir"
+
+type Loop struct{ *ir.BlockSet }
+
+func FindLoops(fn *ir.Function) []Loop {
+ if fn.Blocks == nil {
+ return nil
+ }
+ tree := fn.DomPreorder()
+ var sets []Loop
+ for _, h := range tree {
+ for _, n := range h.Preds {
+ if !h.Dominates(n) {
+ continue
+ }
+ // n is a back-edge to h
+ // h is the loop header
+ if n == h {
+ set := Loop{ir.NewBlockSet(len(fn.Blocks))}
+ set.Add(n)
+ sets = append(sets, set)
+ continue
+ }
+ set := Loop{ir.NewBlockSet(len(fn.Blocks))}
+ set.Add(h)
+ set.Add(n)
+ for _, b := range allPredsBut(n, h, nil) {
+ set.Add(b)
+ }
+ sets = append(sets, set)
+ }
+ }
+ return sets
+}
+
+func allPredsBut(b, but *ir.BasicBlock, list []*ir.BasicBlock) []*ir.BasicBlock {
+outer:
+ for _, pred := range b.Preds {
+ if pred == but {
+ continue
+ }
+ for _, p := range list {
+ // TODO improve big-o complexity of this function
+ if pred == p {
+ continue outer
+ }
+ }
+ list = append(list, pred)
+ list = allPredsBut(pred, but, list)
+ }
+ return list
+}