#!/usr/bin/env python3
"""对照两种前沿策略,同样的求解器、同样的预算:
  fifo  : 朴素 FIFO,前沿超限则截断(无覆盖率反馈)
  cov   : 覆盖率制导 —— 用 afl-showmap 在 AFL 插桩版上量边,只保留【带来新边】的产出
这正是本框架 Worker/Master 两级 bitmap 去重所做的事。"""
import hashlib, json, os, shutil, subprocess, sys, tempfile, time

SYMBIN, AFLBIN, DEPTH, MAXGEN, POLICY, CAP = (
    sys.argv[1], sys.argv[2], int(sys.argv[3]), int(sys.argv[4]), sys.argv[5], int(sys.argv[6]))
W = tempfile.mkdtemp(prefix="n2_")

def edges(path):
    mp = os.path.join(W, "m.txt")
    subprocess.run(["afl-showmap", "-t", "3000", "-m", "none", "-q", "-o", mp, "--", AFLBIN, path],
                   stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                   env={**os.environ, "AFL_MAP_SIZE": "65536"})
    try:
        with open(mp) as f:
            return {int(l.split(":")[0]) for l in f if ":" in l}
    except OSError:
        return set()

def depth_of(p):
    d = open(p, "rb").read(); n = 0
    while n < DEPTH and n < len(d) and d[n] == 0x41 + n: n += 1
    return n

seed = os.path.join(W, "seed"); open(seed, "wb").write(b"\x00" * DEPTH)
virgin = set(edges(seed))                     # 全局 bitmap(边集)
seen = {hashlib.sha256(open(seed,'rb').read()).hexdigest()}
frontier = [seed]; best = 0
execs = out_total = queries = 0; solve_us = 0; t0 = time.time()
for gen in range(1, MAXGEN + 1):
    nxt = []
    for inp in frontier:
        o = tempfile.mkdtemp(dir=W); tel = o + ".json"
        subprocess.run([SYMBIN, inp], env={**os.environ, "SYMCC_OUTPUT_DIR": o,
                       "SYMCC_INPUT_FILE": inp, "SYMCC_TELEMETRY_OUT": tel},
                       stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=300)
        execs += 1
        if os.path.exists(tel):
            t = json.load(open(tel)); queries += t["solver_queries"]; solve_us += t["solver_time_us"]
        for fn in sorted(os.listdir(o)):
            if fn.endswith(".hints"): continue
            p = os.path.join(o, fn); out_total += 1
            h = hashlib.sha256(open(p,"rb").read()).hexdigest()
            if h in seen: continue
            seen.add(h)
            keep = os.path.join(W, h[:16]); shutil.copy(p, keep)
            if POLICY == "cov":
                e = edges(keep)
                if not (e - virgin):           # 无新边 → 丢弃(这就是 bitmap 去重)
                    continue
                virgin |= e
            nxt.append(keep); best = max(best, depth_of(keep))
    if best >= DEPTH:
        print(f"  [{POLICY:4}] 深度{DEPTH} 第{gen}代解出 | 执行={execs} 产出={out_total} 保留={len(nxt)} "
              f"Z3查询={queries} 求解={solve_us/1000:.0f}ms 墙钟={time.time()-t0:.1f}s"); break
    frontier = nxt[:CAP]
    if not frontier:
        print(f"  [{POLICY:4}] 前沿枯竭 止步深度{best}/{DEPTH} 第{gen}代 | 执行={execs} 产出={out_total} "
              f"Z3查询={queries} 求解={solve_us/1000:.0f}ms 墙钟={time.time()-t0:.1f}s"); break
else:
    print(f"  [{POLICY:4}] {MAXGEN}代未解出 止步深度{best}/{DEPTH} | 执行={execs} 产出={out_total} "
          f"Z3查询={queries} 求解={solve_us/1000:.0f}ms 墙钟={time.time()-t0:.1f}s")
shutil.rmtree(W, ignore_errors=True)