A new algorithm for identifying loops in decompilation

havlak.cc

// Havlak-Tarjan loop-nesting-forest algorithm, self-contained.
//
// Two passes over the CFG: (1) a DFS spanning tree recording preorder and finish
// times, then (2) candidate headers in reverse preorder, each flood-filling its
// loop body backward through predecessors, with a UNION-FIND collapse of the
// already-formed inner loops. Each loop tracks its ENTRY set, so absorbing an
// irreducible multi-entry subloop re-examines the predecessors of ALL its
// entries -- pulling in the outer nodes that reach it through a non-header
// entry. It builds the same loop-nesting forest as the single-pass wei.cpp and
// as llvm/include/llvm/ADT/GenericCycleImpl.h.
//
// Output matches wei.cpp: each node's innermost loop header, then the forest
// with each loop's non-header entries.
//
// Input (same format as the natural-loops post):
//   n m
//   u v          (m directed edges u->v; node 0 is the entry)
//
// Successors are visited in input order. Phase-1 recursion depth equals the
// longest DFS path, so extremely deep graphs need `ulimit -s unlimited`.
#include <algorithm>
#include <cstdio>
#include <vector>
using namespace std;

vector<vector<int>> succ, pred; // successor / predecessor lists, input order
vector<int> pre, fin;           // DFS preorder index (-1 = unvisited); finish rank
vector<int> preorder;           // nodes in DFS preorder
int preCnt = 0, finCnt = 0;

// Phase 1: DFS spanning tree. u is a DFS-descendant of h iff
//   pre[h] <= pre[u] && fin[u] <= fin[h].
void dfs(int u) {
  pre[u] = preCnt++;
  preorder.push_back(u);
  for (int w : succ[u])
    if (pre[w] < 0)
      dfs(w);
  fin[u] = finCnt++;
}

// Recursively print loop h after all of its subloops (children before parents),
// mirroring wei.cpp. entries[h] lists h's non-header entries (empty if reducible).
void emitLoop(int h, const vector<vector<int>> &members,
              const vector<vector<int>> &children,
              const vector<vector<int>> &entries) {
  for (int c : children[h])
    emitLoop(c, members, children, entries);
  printf("loop %d", h);
  if (!entries[h].empty()) {
    printf(" (non-header entries");
    for (int e : entries[h])
      printf(" %d", e);
    printf(")");
  }
  printf(":");
  printf(" %d", h); // the header is a member of its own loop
  for (int v : members[h])
    printf(" %d", v);
  for (int c : children[h])
    printf(" (loop %d)", c);
  puts("");
}

int main() {
  int n, m;
  if (scanf("%d %d", &n, &m) != 2)
    return 0;
  succ.assign(n, {});
  pred.assign(n, {});
  for (int i = 0; i < m; i++) {
    int u, v;
    if (scanf("%d %d", &u, &v) != 2)
      return 1;
    succ[u].push_back(v);
    pred[v].push_back(u);
  }
  pre.assign(n, -1);
  fin.assign(n, 0);
  if (n > 0)
    dfs(0);

  // --- Phase 2: Havlak-Tarjan flood-fill with per-loop entry sets. ---
  vector<int> loopHeader;          // loopHeader[L] = header node of loop L
  vector<int> loopParent;          // immediate enclosing loop, -1 if top-level
  vector<int> uf;                  // path-halved union-find over loopParent roots
  vector<vector<int>> loopEntries; // loop L's entry nodes (header included)
  vector<int> toLoop(n, -1);       // innermost loop directly owning a node
  vector<int> entryOf(n, -1);      // node is an entry of loop entryOf[node]
  vector<int> worklist;
  int curH = -1, curL = -1; // header/index of the loop being built

  auto inSubtree = [&](int p) {
    return pre[curH] <= pre[p] && fin[p] <= fin[curH];
  };
  // Enqueue predecessors of `block`: those inside curH's subtree feed the flood;
  // a reachable predecessor outside the subtree marks `block` an entry of curL.
  auto processPreds = [&](int block) {
    for (int p : pred[block]) {
      if (pre[p] < 0)
        continue; // unreachable predecessor
      if (inSubtree(p))
        worklist.push_back(p);
      else if (entryOf[block] != curL) {
        entryOf[block] = curL;
        loopEntries[curL].push_back(block);
      }
    }
  };

  for (int idx = (int)preorder.size() - 1; idx >= 0; --idx) {
    int h = preorder[idx];
    curH = h;

    // A loop exists only if some predecessor of h is a back edge (in h's
    // subtree, which also captures a self-loop where the predecessor is h).
    bool hasBack = false;
    for (int p : pred[h])
      if (pre[p] >= 0 && inSubtree(p)) {
        hasBack = true;
        break;
      }
    if (!hasBack)
      continue;

    int L = (int)loopHeader.size();
    curL = L;
    loopHeader.push_back(h);
    loopParent.push_back(-1);
    uf.push_back(L);
    loopEntries.push_back({});

    // The header is always an entry and a member of its own loop.
    entryOf[h] = L;
    loopEntries[L].push_back(h);
    toLoop[h] = L;

    // Seed the flood with the back-edge tails (subtree predecessors of h).
    worklist.clear();
    for (int p : pred[h])
      if (pre[p] >= 0 && inSubtree(p))
        worklist.push_back(p);

    while (!worklist.empty()) {
      int v = worklist.back();
      worklist.pop_back();
      if (v == h)
        continue;
      if (toLoop[v] >= 0) {
        // v already belongs to a loop: absorb that loop's outermost ancestor as
        // a subloop of L, then continue the flood from the predecessors of ALL
        // of the subloop's entries (irreducible re-entry handling).
        int sub = toLoop[v];
        while (uf[sub] != sub) { // find root with path halving
          uf[sub] = uf[uf[sub]];
          sub = uf[sub];
        }
        if (sub != L) {
          loopParent[sub] = L;
          uf[sub] = L;
          for (int e : loopEntries[sub])
            processPreds(e);
        }
      } else {
        toLoop[v] = L;
        processPreds(v);
      }
    }
  }

  // --- Phase 3: canonical innermost-loop-header per node, and per-header data. ---
  vector<int> ilh(n, -1);
  vector<char> header(n, 0);
  for (int L = 0; L < (int)loopHeader.size(); ++L) {
    int h = loopHeader[L];
    header[h] = 1;
    ilh[h] = loopParent[L] >= 0 ? loopHeader[loopParent[L]] : -1;
  }
  for (int v = 0; v < n; ++v)
    if (!header[v] && toLoop[v] >= 0)
      ilh[v] = loopHeader[toLoop[v]];

  // Non-header entries per header node (ascending), from the loop entry sets.
  vector<vector<int>> entries(n);
  for (int L = 0; L < (int)loopHeader.size(); ++L) {
    int h = loopHeader[L];
    for (int e : loopEntries[L])
      if (e != h)
        entries[h].push_back(e);
    sort(entries[h].begin(), entries[h].end());
  }

  // Each node's innermost loop header (-1 = outside every loop).
  for (int v = 0; v < n; v++)
    printf("%d: %d\n", v, ilh[v]);

  // Reconstruct the loop-nesting forest from ilh[]: loop h owns h plus every
  // node/subloop-header x with ilh[x]==h.
  vector<vector<int>> members(n), children(n);
  vector<int> roots;
  for (int v = 0; v < n; v++) {
    if (header[v])
      (ilh[v] < 0 ? roots : children[ilh[v]]).push_back(v);
    else if (ilh[v] >= 0)
      members[ilh[v]].push_back(v);
  }
  for (int h : roots)
    emitLoop(h, members, children, entries);
  return 0;
}

loop.py

# 《A New Algorithm for Identifying Loops in Decompilation》（反编译中循环识别的新算法）是由韦韬（Tao Wei）、毛剑、邹维等人在2007年第14届国际静态分析研讨会（SAS 2007）上发表的关于二进制逆向工程的重要论文。
#
# Input (same format as the natural-loops / irreducible-loops posts):
#   n m
#   u v        m directed edges u->v; node 0 is the entry
# Output:
#   each node's innermost loop header (-1 = outside every loop), then the headers.
import sys
sys.setrecursionlimit(1 << 20)

NIL = -1  # sentinel for "no loop header"; distinct from every node id, incl. 0


def identify_loops(n, succ, entry=0):
    traversed = [False] * n
    dfsp_pos = [0] * n          # 0 = not on the current DFS path; >0 = position
    iloop_header = [NIL] * n    # innermost loop header of each node
    self_loop = [False] * n
    counter = 0

    # Splice "b belongs to a loop headed by h" into b's header chain, keeping it
    # ordered innermost..outermost by dfsp_pos -- replaces the UNION-FIND merge.
    def tag_lhead(b, h):
        if b == h or h == NIL:
            return
        cur1, cur2 = b, h
        while iloop_header[cur1] != NIL:
            ih = iloop_header[cur1]
            if ih == cur2:
                return
            if dfsp_pos[ih] < dfsp_pos[cur2]:
                iloop_header[cur1] = cur2
                cur1, cur2 = cur2, ih
            else:
                cur1 = ih
        iloop_header[cur1] = cur2

    def traverse(b0):
        nonlocal counter
        counter += 1
        traversed[b0] = True
        dfsp_pos[b0] = counter
        for b1 in succ[b0]:
            if b1 == b0:
                self_loop[b0] = True
            if not traversed[b1]:
                traverse(b1)                    # (A) tree edge
                tag_lhead(b0, iloop_header[b1])
            elif dfsp_pos[b1] > 0:              # (B) back edge: b1 is a header
                tag_lhead(b0, b1)
            elif iloop_header[b1] != NIL:       # b1 finished, already in a loop
                h = iloop_header[b1]
                if dfsp_pos[h] > 0:             # (D) innermost header still open
                    tag_lhead(b0, h)
                else:                           # (E) re-entry (irreducible): climb
                    while iloop_header[h] != NIL:
                        h = iloop_header[h]
                        if dfsp_pos[h] > 0:
                            tag_lhead(b0, h)
                            break
            # (C) else: forward/cross edge into non-loop code -> ignore
        dfsp_pos[b0] = 0                        # leave the path

    if n > 0:
        traverse(entry)

    # A node is a loop header iff it is some node's innermost header or self-loops.
    loop_heads = set()
    for i in range(n):
        if iloop_header[i] != NIL:
            loop_heads.add(iloop_header[i])
        if self_loop[i]:
            loop_heads.add(i)
    return iloop_header, loop_heads


def main():
    data = sys.stdin.read().split()
    it = iter(data)
    n = int(next(it))
    m = int(next(it))
    succ = [[] for _ in range(n)]
    for _ in range(m):
        u = int(next(it))
        v = int(next(it))
        succ[u].append(v)
    iloop_header, loop_heads = identify_loops(n, succ)
    for v in range(n):
        print(f"{v}: {iloop_header[v]}")
    print("loop headers:", " ".join(map(str, sorted(loop_heads))) or "(none)")


if __name__ == "__main__":
    main()

wei.cc

// A New Algorithm for Identifying Loops in Decompilation
//   Tao Wei, Jian Mao, Wei Zou, Yu Chen -- SAS 2007.
//
// A single depth-first traversal tags every node with its innermost loop header
// on the fly: no dominator tree, no UNION-FIND, no second bottom-up pass.  It
// handles nested and irreducible (multi-entry) loops under Havlak's loop-nesting
// -forest definition.
//
// Input (same format as that post):
//   n m
//   u v          (m directed edges u->v; node 0 is the entry)
//
// Successors are visited in input order.  Recursion depth equals the longest
// DFS path, so extremely deep graphs need `ulimit -s unlimited`.
#include <algorithm>
#include <cstdio>
#include <utility>
#include <vector>
using namespace std;

// Per-node DFS state kept together (AoS), matching the paper's `Block` fields:
// the hot path touches a node's ilh/pos/traversed as a unit, so one cache line
// per node beats four parallel arrays.
struct Node {
  int ilh = -1;           // iloop_header: innermost loop header, -1 = none
  int pos = 0;            // DFSP_pos: 1-based depth on the current DFS path;
                          // 0 once the node leaves it (or is unvisited)
  bool traversed = false; // has the DFS reached this node?
  bool header = false;    // is this node the header of some loop?
};

vector<Node> nd;          // per-node state, indexed by node id
vector<vector<int>> succ; // successor lists, in input order
// (b, h): an edge re-enters the closed loop headed by h at b, below the header,
// so b is a non-header entry of that loop. This is all the paper's per-header
// `irreducible` bit encodes, plus which block does the re-entering.
vector<pair<int, int>> reentries;

// Weave loop header h (and its own header chain) into b's innermost-loop-header chain, ordered innermost..outermost by
// DFS-path position. Building this chain on the fly is why Wei needs no UNION-FIND at all; Havlak-Tarjan instead
// collapses inner loops with UNION-FIND in a separate reverse-preorder pass.
void tagLoopHeader(int b, int h) {
  if (h == -1)
    return;
  // Invariant, nd[b].pos >= nd[h].pos
  while (b != h) {
    int ih = nd[b].ilh;
    if (ih == -1) { // b's chain ended: append the rest of h's chain
      nd[b].ilh = h;
      return;
    }
    if (nd[ih].pos >= nd[h].pos) {
      b = ih;
    } else {
      nd[b].ilh = h;
      b = h;
      h = ih;
    }
  }
}

// Traverse b0 at path depth p; return b0's innermost loop header.
int dfs(int b0, int p) {
  nd[b0].traversed = true;
  nd[b0].pos = p;
  for (int b : succ[b0]) {
    if (!nd[b].traversed) { // (A) tree edge: recurse, absorb child's header
      tagLoopHeader(b0, dfs(b, p + 1));
    } else if (nd[b].pos > 0) { // (B) back edge: b is a loop header
      nd[b].header = true;
      tagLoopHeader(b0, b);
    } else {
      // (C)(D)(E) merged: climb b's header chain outward from its innermost header. Every header still off the DFS path
      // heads a closed loop that b0->b re-enters below the header, so b is a non-header entry of it. Stop at the first
      // header still on the path -- b0 is a genuine interior node of that loop, so attribute b0 to it. (C is the empty
      // chain; D stops on the first, still-open header; E climbs past one or more closed headers -- an irreducible
      // re-entry.)
      for (int h = nd[b].ilh; h >= 0; h = nd[h].ilh) {
        if (nd[h].pos > 0) {
          tagLoopHeader(b0, h);
          break;
        }
        reentries.push_back({b, h});
      }
    }
  }
  nd[b0].pos = 0; // b0 leaves the DFS path
  return nd[b0].ilh;
}

// Recursively print loop h after all of its subloops (children before parents), mirroring the ordering of the
// natural-loops post.
void emitLoop(int h, const vector<vector<int>> &members, const vector<vector<int>> &children,
              const vector<vector<int>> &entries) {
  for (int c : children[h])
    emitLoop(c, members, children, entries);
  printf("loop %d", h);
  if (!entries[h].empty()) { // non-header entries => irreducible loop
    printf(" (non-header entries");
    for (int e : entries[h])
      printf(" %d", e);
    printf(")");
  }
  printf(":");
  printf(" %d", h); // the header is a member of its own loop
  for (int v : members[h])
    printf(" %d", v);
  for (int c : children[h])
    printf(" (loop %d)", c);
  puts("");
}

int main() {
  int n, m;
  if (scanf("%d %d", &n, &m) != 2)
    return 0;
  succ.assign(n, {});
  for (int i = 0; i < m; i++) {
    int u, v;
    if (scanf("%d %d", &u, &v) != 2)
      return 1;
    succ[u].push_back(v);
  }
  nd.assign(n, Node{});
  if (n > 0)
    dfs(0, 1);

  // Each node's innermost loop header (-1 = outside every loop).
  for (int v = 0; v < n; v++)
    printf("%d: %d\n", v, nd[v].ilh);

  // Reconstruct the loop-nesting forest from ilh[]: loop h owns h plus every
  // node/subloop-header x with ilh[x]==h.  (header, ilh) fully determines it.
  vector<vector<int>> members(n), children(n);
  vector<int> roots;
  for (int v = 0; v < n; v++) {
    if (nd[v].header)
      (nd[v].ilh < 0 ? roots : children[nd[v].ilh]).push_back(v);
    else if (nd[v].ilh >= 0)
      members[nd[v].ilh].push_back(v);
  }

  // Collect each loop's non-header entries from the recorded re-entries.
  // sort+unique drops duplicate edges re-entering one loop at the same block.
  sort(reentries.begin(), reentries.end());
  reentries.erase(unique(reentries.begin(), reentries.end()), reentries.end());
  vector<vector<int>> entries(n);
  for (auto [b, h] : reentries)
    entries[h].push_back(b);

  for (int h : roots)
    emitLoop(h, members, children, entries);
  return 0;
}

wei_blocks.cpp

// Identify loops in one DFS pass, then lay them out as a flat block layout.
//
// "A New Algorithm for Identifying Loops in Decompilation"
//   Tao Wei, Jian Mao, Wei Zou, Yu Chen -- SAS 2007.
//
// A single depth-first traversal tags every node with its innermost loop header
// on the fly: no dominator tree, no UNION-FIND, no second bottom-up pass. It
// handles nested and irreducible (multi-entry) loops under Havlak's loop-nesting
// -forest definition.
//
// The forest is then flattened with no nested containers: loops are numbered by
// header preorder rank, sub-loops hang off a childHead/nextSibling list, and each
// loop records only its own block count. A reverse-preorder pass drops every
// in-loop block into one flat `layout` array, landing the header first in each
// loop's contiguous slice, so it need not be stored separately.
//
// Input:  n m / u v   (m directed edges u->v; node 0 is the entry)
//
// Successors are visited in input order. Recursion depth equals the longest DFS
// path, so extremely deep graphs need `ulimit -s unlimited`.
#include <algorithm>
#include <cassert>
#include <cstdio>
#include <utility>
#include <vector>
using namespace std;

// Per-node state kept together (AoS): one cache line per node beats parallel
// arrays. All fields default to 0/false, so assign() is a cheap zero-fill. `pos`
// (live only during the DFS) and `loopIdx` (live only after) never overlap, so
// they share one word.
struct Node {
  int ilh = 0;            // iloop_header: innermost loop header; dfs() sets the
                          //   -1 = none sentinel on first visit
  union {
    int pos = 0;          // during the DFS: 1-based path depth; 0 once off it
    int loopIdx;          // after the DFS: innermost-loop rank, -1 = none
  };
  bool traversed = false; // has the DFS reached this node?
  bool header = false;    // is this node the header of some loop?
};

vector<Node> nd;          // per-node state, indexed by node id
vector<vector<int>> succ; // successor lists, in input order
vector<int> preorder;     // node ids in DFS preorder
int numHeaders = 0;       // number of loop headers found by dfs()
// (h, b): an edge re-enters the closed loop headed by h at b, below the header,
// so b is a non-header entry of that (irreducible) loop.
vector<pair<int, int>> reentries;

// Weave header h (and its own header chain) into b's innermost-loop-header chain,
// ordered innermost..outermost by DFS-path position. Building it on the fly is
// why Wei needs no UNION-FIND.
void tagLoopHeader(int b, int h) {
  assert(h != -1);
  // Invariant: nd[b].pos >= nd[h].pos
  while (b != h) {
    int ih = nd[b].ilh;
    if (ih == -1) { // b's chain ended: append the rest of h's chain
      nd[b].ilh = h;
      return;
    }
    if (nd[ih].pos >= nd[h].pos)
      b = ih;
    else {
      nd[b].ilh = h;
      b = h;
      h = ih;
    }
  }
}

// Traverse b0 at path depth p; return b0's innermost loop header.
int dfs(int b0, int p) {
  nd[b0].ilh = -1; // first visit: install the "no header" sentinel
  nd[b0].pos = p;
  nd[b0].traversed = true;
  preorder.push_back(b0);
  for (int b : succ[b0]) {
    if (!nd[b].traversed) { // (A) tree edge: recurse, absorb child's header
      int h = dfs(b, p + 1);
      if (h >= 0)
        tagLoopHeader(b0, h);
    } else if (nd[b].pos > 0) { // (B) back edge: b is a loop header
      if (!nd[b].header) {      // count each distinct header once
        nd[b].header = true;
        ++numHeaders;
      }
      tagLoopHeader(b0, b);
    } else {
      // (C/D/E) climb b's header chain: each off-path header heads a closed loop
      // that b0->b re-enters below it (so b is a non-header entry). Stop at the
      // first on-path header and attribute b0 to that loop.
      for (int h = nd[b].ilh; h >= 0; h = nd[h].ilh) {
        if (nd[h].pos > 0) {
          tagLoopHeader(b0, h);
          break;
        }
        reentries.push_back({h, b});
      }
    }
  }
  nd[b0].pos = 0; // b0 leaves the DFS path
  return nd[b0].ilh;
}

// One record per loop, keyed by header preorder rank.
struct Loop {
  int header, childHead = -1, nextSibling = -1, ownCount = 1;
  int begin = 0, end = 0; // slice in `layout`
};
vector<Loop> loops;
vector<int> layout;
vector<vector<int>> entries; // header -> non-header entries
int cursor = 0;

// Reserve each loop a contiguous slice; `begin` temporarily holds the own-region
// end, which the fill pass below walks back down to the slice start.
void tour(int c) {
  cursor += loops[c].ownCount;
  loops[c].begin = cursor;
  for (int ch = loops[c].childHead; ch >= 0; ch = loops[ch].nextSibling)
    tour(ch);
  loops[c].end = cursor;
}

// Print loop c and its subloops, nested by indentation.
void dump(int c, int indent) {
  Loop &L = loops[c];
  printf("%*sloop %d  slice [%d,%d)  own blocks:", 2 * indent, "", L.header,
         L.begin, L.end);
  for (int i = L.begin; i < L.begin + L.ownCount; i++)
    printf(" %d", layout[i]); // header is layout[begin], the rest follow
  if (!entries[L.header].empty()) { // non-header entries => irreducible loop
    printf("  (non-header entries:");
    for (int e : entries[L.header])
      printf(" %d", e);
    putchar(')');
  }
  puts("");
  for (int ch = L.childHead; ch >= 0; ch = loops[ch].nextSibling)
    dump(ch, indent + 1);
}

int main() {
  int n, m;
  if (scanf("%d %d", &n, &m) != 2)
    return 0;
  succ.assign(n, {});
  for (int i = 0; i < m; i++) {
    int u, v;
    if (scanf("%d %d", &u, &v) != 2)
      return 1;
    succ[u].push_back(v);
  }
  nd.assign(n, Node{});
  if (n > 0)
    dfs(0, 1);

  int topHead = -1; // head of the root loops' sibling list
  if (numHeaders) {
    // Number loops by header preorder rank and resolve each block's innermost
    // loop in the same sweep; ilh is an already-numbered ancestor.
    loops.reserve(numHeaders);
    for (int v : preorder) {
      if (nd[v].header) {
        nd[v].loopIdx = loops.size();
        loops.push_back({v});
      } else if (nd[v].ilh >= 0) {
        nd[v].loopIdx = nd[nd[v].ilh].loopIdx;
        ++loops[nd[v].loopIdx].ownCount;
      } else {
        nd[v].loopIdx = -1;
      }
    }
    // Link loops under their parents. Ranks run in header preorder, so walking
    // them in reverse and prepending gives child lists in forward preorder.
    for (int c = (int)loops.size() - 1; c >= 0; c--) {
      int v = loops[c].header;
      int &head =
          nd[v].ilh >= 0 ? loops[nd[nd[v].ilh].loopIdx].childHead : topHead;
      loops[c].nextSibling = head;
      head = c;
    }
    // Size the slices, then fill blocks back-to-front so the header lands first.
    for (int c = topHead; c >= 0; c = loops[c].nextSibling)
      tour(c);
    layout.assign(cursor, 0);
    for (int i = (int)preorder.size() - 1; i >= 0; i--)
      if (nd[preorder[i]].loopIdx >= 0)
        layout[--loops[nd[preorder[i]].loopIdx].begin] = preorder[i];
    // Non-header entries per header, dropping duplicate re-entry edges.
    sort(reentries.begin(), reentries.end());
    reentries.erase(unique(reentries.begin(), reentries.end()), reentries.end());
    entries.assign(n, {});
    for (auto [h, b] : reentries)
      entries[h].push_back(b);
  }

  printf("block layout:");
  for (int b : layout)
    printf(" %d", b);
  puts("");
  for (int c = topHead; c >= 0; c = loops[c].nextSibling)
    dump(c, 0);
  return 0;
}