#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Open-IOV contributors
# Print the distinct pages of the semaphore VAs a process references
# (reuses memprobe.py's scanner; run in the guest as root). The VA window is
# derived from the process's /dev/nvidia* mappings unless lo/hi are given.
# usage: sudo python3 sempages.py <pid> [lo] [hi]
import os
import sys

sys.path.insert(0, '/tmp')
import memprobe as m

pid = int(sys.argv[1])
if len(sys.argv) > 3:
    lo = int(sys.argv[2], 16)
    hi = int(sys.argv[3], 16)
else:
    lo = 1 << 64
    hi = 0
    with open('/proc/%d/maps' % pid) as f:
        for line in f:
            parts = line.split()
            if len(parts) < 6 or not parts[5].startswith('/dev/nvidia'):
                continue
            a, b = parts[0].split('-')
            lo = min(lo, int(a, 16))
            hi = max(hi, int(b, 16))
    if hi == 0:
        print('no /dev/nvidia mappings')
        raise SystemExit(1)

fd = os.open('/proc/%d/mem' % pid, os.O_RDONLY)
# Pin what libcuda actually polls: the semaphore objects referenced by the
# client's trackers (and the semaphore memory they point at). Scanning for
# bare sem-object shapes matches far too many pages once the ring/pushbuffer
# windows are in range; trackers are a structural reference, so pinning is
# limited to pages the client really waits on. Fall back to the shape scan
# only when no tracker resolves.
pages = set()
trackers = m.find_trackers(pid, fd, lo, hi, limit=256)
for _p, _sub, _last, _cur, sem, mem in trackers:
    if sem:
        pages.add(sem & ~0xfff)
    if mem:
        pages.add(mem & ~0xfff)
if not pages:
    sems = m.sem_objs(pid, fd, lo, hi)
    pages = {va & ~0xfff for va in sems}
os.close(fd)
pages = sorted(pages)
print('trackers: %d pages: %d (window 0x%x..0x%x)' %
      (len(trackers), len(pages), lo, hi))
for pg in pages:
    print('0x%x' % pg)