#
# Copyright (C) 2024-2024 Intel Corporation.
# SPDX-License-Identifier: MIT
#

# A simple call trace for probe mode PinGlue Service

import json, pin

# Maps to keep track of call counts and locations
call_trace_count_map = {}
call_trace_location_map = {}


def rtn_cb(name):
    """Callback function that increments the call count for a given routine name.
    Args:
    name (str): The name of the routine being called and it's image (in format of <routine name>@<image name> ).
    """
    global call_trace_count_map
    assert(name in call_trace_count_map)
    call_trace_count_map[name] += 1


def exit_cb():
    """Callback function that is called before the exit function.
    It collects the call trace data and sends it as a JSON string.
    """
    global call_trace_count_map, call_trace_location_map
    result_json = {}
    
    for name in call_trace_count_map:
        assert(name in call_trace_location_map)
        result_json[name] = {"CallCount": call_trace_count_map[name], "Location": call_trace_location_map[name]}
    result_json_str = json.dumps(result_json)
    Glue_SendServiceResultCallback(result_json_str)
    

def image_instrumentation_cb(img):
    """Callback function for instrumenting images.
    Args:
    img: The image object to be instrumented.
    """

    # Instrument before the exit function
    rtn = pin.RTN_FindByName(img, "_exit")
    if pin.RTN_Valid(rtn):
        pin.RTN_InsertCallProbed(rtn, pin.IPOINT_BEFORE, exit_cb)
    
    # Iterate through sections and routines within the image
    img_name = pin.IMG_Name(img)
    sec = pin.IMG_SecHead(img)
    while(pin.SEC_Valid(sec)):
        rtn = pin.SEC_RtnHead(sec)
        while(pin.RTN_Valid(rtn)):
            rtn_name = pin.RTN_Name(rtn)
            
            # Check with the user-defined client callback whether to instrument the routine
            if(is_to_instrument_rtn(img_name, rtn_name)):
                assert(pin.RTN_IsSafeForProbedInsertion(rtn))
                name =  rtn_name + "@" + img_name
                
                # Insert a call to the routine callback before the routine execution
                pin.RTN_InsertCallProbed(rtn, pin.IPOINT_BEFORE, rtn_cb, pin.IARG_PYOBJ, name)
                
                # Initialize the call count and location information for the routine
                global call_trace_count_map, call_trace_location_map
                call_trace_count_map[name] = 0
                loc_info = pin.PIN_GetSourceLocation(pin.RTN_Address(rtn))
                call_trace_location_map[name] = f"srcfile:{loc_info[2]}, line: {loc_info[1]}, column: {loc_info[0]}"
            
            rtn = pin.RTN_Next(rtn)
        sec = pin.SEC_Next(sec)
    
# Register the image instrumentation callback with Pin
pin.IMG_AddInstrumentFunction(image_instrumentation_cb)
