swift-mirror/utils/lldb/lldbToolBox.py

"""
LLDB Helpers for working with the swift compiler.

Load into LLDB with 'command script import /path/to/lldbToolBox.py'

This will also import LLVM data formatters as well, assuming that llvm is next
to the swift checkout.
"""
import argparse
import os
import shlex
import subprocess
import sys
import tempfile

import lldb

REPO_BASE = os.path.abspath(os.path.join(__file__, os.pardir, os.pardir,
                                         os.pardir, os.pardir))
SWIFT_REPO = os.path.join(REPO_BASE, "swift")
LLVM_REPO = os.path.join(REPO_BASE, "llvm-project")
LLVM_DATAFORMATTER_PATH = os.path.join(LLVM_REPO, "llvm", "utils",
                                       "lldbDataFormatters.py")
SWIFT_DATAFORMATTER_PATH = os.path.join(SWIFT_REPO, "utils",
                                        "lldb", "lldbSwiftDataFormatters.py")


def import_llvm_dataformatters(debugger):
    if not os.access(LLVM_DATAFORMATTER_PATH, os.F_OK):
        print("WARNING! Could not find LLVM data formatters!")
        return
    cmd = 'command script import {}'.format(LLVM_DATAFORMATTER_PATH)
    debugger.HandleCommand(cmd)
    print("Loaded LLVM data formatters.")


def import_swift_dataformatters(debugger):
    if not os.access(SWIFT_DATAFORMATTER_PATH, os.F_OK):
        print("WARNING! Could not find Swift data formatters!")
        return
    cmd = 'command script import {}'.format(SWIFT_DATAFORMATTER_PATH)
    debugger.HandleCommand(cmd)
    print("Loaded Swift data formatters.")


VIEWCFG_PATH = os.path.join(SWIFT_REPO, "utils", "viewcfg")
BLOCKIFYASM_PATH = os.path.join(SWIFT_REPO, "utils", "dev-scripts",
                                "blockifyasm")


def disassemble_asm_cfg(debugger, command, exec_ctx, result, internal_dict):
    """
    This function disassembles the current assembly frame into a temporary file
    and then uses that temporary file as input to blockifyasm | viewcfg. This
    will cause a pdf of the cfg to be opened on Darwin.
    """
    d = exec_ctx.frame.Disassemble()

    with tempfile.TemporaryFile() as f:
        f.write(bytes(d, 'utf-8'))
        f.flush()
        f.seek(0)
        p1 = subprocess.Popen([BLOCKIFYASM_PATH], stdin=f,
                              stdout=subprocess.PIPE)
        subprocess.Popen([VIEWCFG_PATH], stdin=p1.stdout)
        p1.stdout.close()  # Allow p1 to receive a SIGPIPE if p2 exits.


def disassemble_to_file(debugger, command, exec_ctx, result, internal_dict):
    """This function disassembles the current assembly frame into a file specified
    by the user.
    """
    parser = argparse.ArgumentParser(prog='disassemble-to-file', description="""
    Dump the disassembly of the current frame or specified function to the
    specified file.
    """)
    parser.add_argument('file', type=argparse.FileType('w'),
                        default=sys.stdout)
    parser.add_argument('-n', dest='func_name', help="""
    Function name to disassembly. Frame used if unset.""")
    args = parser.parse_args(shlex.split(command))
    if args.func_name is None:
        args.file.write(exec_ctx.frame.disassembly)
    else:
        name = args.func_name
        result = exec_ctx.target.FindFunctions(name)
        if result is None:
            raise RuntimeError('No function with name: {}'.format(name))
        if len(result) > 1:
            errorStr = 'Matched multiple functions to name: {}'
            raise RuntimeError(errorStr.format(name))
        f = result[0].GetFunction()
        args.file.write(str(f.GetInstructions(exec_ctx.target)) + "\n")


def sequence(debugger, command, exec_ctx, result, internal_dict):
    """
    Combine multiple semicolon separated lldb commands into one command.

    This command is particularly useful for defining aliases and breakpoint
    commands. Some examples:

        # Define an alias that prints rax and also steps one instruction.
        command alias xs sequence p/x $rax; stepi

        # Breakpoint command to show the frame's info and arguments.
        breakpoint command add -o 'seq frame info; reg read arg1 arg2 arg3'

        # Override `b` to allow a condition to be specified. For example:
        #     b someMethod if someVar > 2
        command regex b
        s/(.+) if (.+)/seq _regexp-break %1; break mod -c "%2"/
        s/(.*)/_regexp-break %1/
    """
    interpreter = debugger.GetCommandInterpreter()
    for subcommand in command.split(';'):
        subcommand = subcommand.strip()
        if not subcommand:
            continue  # skip empty commands

        ret = lldb.SBCommandReturnObject()
        interpreter.HandleCommand(subcommand, exec_ctx, ret)
        if ret.GetOutput():
            print(ret.GetOutput().strip(), file=result)

        if not ret.Succeeded():
            result.SetError(ret.GetError())
            result.SetStatus(ret.GetStatus())
            return


def __lldb_init_module(debugger, internal_dict):
    import_llvm_dataformatters(debugger)
    import_swift_dataformatters(debugger)
    debugger.HandleCommand('command script add disassemble-asm-cfg '
                           '-f lldbToolBox.disassemble_asm_cfg')
    debugger.HandleCommand('command script add disassemble-to-file '
                           '-f lldbToolBox.disassemble_to_file')
    debugger.HandleCommand('command script add sequence '
                           '-h "Run multiple semicolon separated commands" '
                           '-f lldbToolBox.sequence')