// GhidraScript: FindD3DSetTransformSites.java
// Usage:
//   analyzeHeadless <projDir> <projName> -process <progName> -noanalysis -readOnly \
//     -scriptPath tools/ghidra_scripts -postScript FindD3DSetTransformSites.java [minAddr] [maxAddr]

import ghidra.app.script.GhidraScript;
import ghidra.program.model.address.Address;
import ghidra.program.model.listing.Function;
import ghidra.program.model.listing.Instruction;
import ghidra.program.model.scalar.Scalar;

public class FindD3DSetTransformSites extends GhidraScript {
    private boolean hasScalar(Instruction ins, long want) {
        for (int op = 0; op < ins.getNumOperands(); op++) {
            for (Object obj : ins.getOpObjects(op)) {
                if (obj instanceof Scalar) {
                    Scalar s = (Scalar)obj;
                    if (s.getUnsignedValue() == want || s.getSignedValue() == want) return true;
                }
            }
        }
        return false;
    }

    private boolean inRange(Address addr, Address min, Address max) {
        if (min != null && addr.compareTo(min) < 0) return false;
        if (max != null && addr.compareTo(max) > 0) return false;
        return true;
    }

    private boolean followedByCall(Instruction ins, int maxSteps) {
        Instruction cur = ins;
        for (int i = 0; i < maxSteps; i++) {
            cur = cur.getNext();
            if (cur == null) return false;
            if ("CALL".equals(cur.getMnemonicString())) return true;
        }
        return false;
    }

    private boolean nearbyPushState(Instruction ins) {
        Instruction cur = ins;
        for (int i = 0; i < 10; i++) {
            cur = cur.getPrevious();
            if (cur == null) break;
            if (!"PUSH".equals(cur.getMnemonicString())) continue;
            if (hasScalar(cur, 2) || hasScalar(cur, 3)) return true;
        }
        return false;
    }

    @Override
    public void run() throws Exception {
        String[] args = getScriptArgs();
        Address min = args.length > 0 ? toAddr(Long.decode(args[0])) : null;
        Address max = args.length > 1 ? toAddr(Long.decode(args[1])) : null;

        println("=== D3D SetTransform VIEW/PROJ candidates ===");
        int hits = 0;
        for (Instruction ins : currentProgram.getListing().getInstructions(true)) {
            if (!inRange(ins.getAddress(), min, max)) continue;
            if (!hasScalar(ins, 0xb0)) continue;
            if (!followedByCall(ins, 3)) continue;
            if (!nearbyPushState(ins)) continue;

            hits++;
            Function f = getFunctionContaining(ins.getAddress());
            String fs = f == null ? "(no func)" : f.getName() + " @ " + f.getEntryPoint();
            println("");
            println("HIT " + hits + "  " + ins.getAddress() + "  " + fs);

            Instruction cur = ins;
            for (int i = 0; i < 10; i++) {
                Instruction prev = cur.getPrevious();
                if (prev == null) break;
                cur = prev;
            }
            for (int i = 0; i < 18 && cur != null; i++) {
                String mark = cur.getAddress().equals(ins.getAddress()) ? "=>" : "  ";
                println(mark + " " + cur.getAddress() + "  " + cur);
                cur = cur.getNext();
            }
            if (monitor.isCancelled()) break;
        }
        println("hits: " + hits);
    }
}