# -*- coding: utf-8 -*-
#
# Part of ChipBench (AGPLv3). SPDX-FileCopyrightText: 2026 Starting Line Productions LLC
# AutoFuncs.py -- create functions at every LCALL/LJMP target + scheduler
# jumptable entries, iterating until stable. Handles call tables (JMP @A+DPTR)
# where the table entries are LJMPs by treating table bodies as code.
from ghidra.program.model.address import AddressSet
from ghidra.app.cmd.function import CreateFunctionCmd
from ghidra.app.cmd.disassemble import DisassembleCommand
from ghidra.app.plugin.core.analysis import AutoAnalysisManager

prog = currentProgram
af = prog.getAddressFactory()
cs = af.getDefaultAddressSpace()
fm = prog.getFunctionManager()
mem = prog.getMemory()

def A(x): return cs.getAddress(x)

def mk(a):
    if a < 0x10000 and fm.getFunctionAt(a) is None:
        CreateFunctionCmd(a).applyTo(prog, monitor)

rounds = 0
prev = -1
while rounds < 4:
    n = 0
    for i in range(0x0000, 0xEE00):
        try:
            b = mem.getByte(A(i)) & 0xFF
        except:
            continue
        if b == 0x12 or b == 0x02:  # LCALL/LJMP
            try:
                hi = mem.getByte(A(i+1)) & 0xFF
                lo = mem.getByte(A(i+2)) & 0xFF
            except:
                continue
            t = (hi << 8) | lo
            if 0x0008 <= t < 0xEE00:
                if fm.getFunctionAt(A(t)) is None:
                    if CreateFunctionCmd(A(t)).applyTo(prog, monitor):
                        n += 1
    # reanalyze
    mgr = AutoAnalysisManager.getAnalysisManager(prog)
    mgr.reAnalyzeAll(None)
    mgr.startAnalysis(monitor)
    cur = fm.getFunctionCount()
    print("round %d: +%d new, total %d" % (rounds, n, cur))
    if cur == prev:
        break
    prev = cur
    rounds += 1

print("AutoFuncs done: %d functions, %d instructions" % (
    fm.getFunctionCount(), prog.getListing().getNumInstructions()))