GrammaTech / gtirb-rewriting

Python API for rewriting GTIRB files
GNU General Public License v3.0
16 stars 3 forks source link

Pass argument to patch function #6

Closed avncharlie closed 1 year ago

avncharlie commented 1 year ago

Is there a way to create a patch and pass arguments to the patch function? Currently I am using lambda functions to emulate this as such:

context.register_insert(
    SingleBlockScope(block, BlockPosition.ENTRY),
    Patch.from_function(lambda _: f"movq ${ARGUMENT}, %rax", Constraints(x86_syntax=X86Syntax.INTEL))
)

However as the lambda function seems to be evaluated at a later point, the string will also be created at that point. This causes issues when creating patches in a loop. For example:

for i, block in enumerate(blocks):
    context.register_insert(
        SingleBlockScope(block, BlockPosition.ENTRY),
        Patch.from_function(lambda _: f"movq ${i}, %rax", Constraints(x86_syntax=X86Syntax.INTEL))
    )

In the above code, all patches would be generated with the exact same instruction, with the value of i being its final value as the loop finishes iterating.

To avoid this I am using a recursive function as such:

def create_patches(blocks, i):
    if len(blocks) == 0: return
    block = blocks[0]
    context.register_insert(
        SingleBlockScope(block, BlockPosition.ENTRY),
        Patch.from_function(lambda _: f"movq ${i}, %rax", Constraints(x86_syntax=X86Syntax.INTEL))
    )
    create_patches(blocks[1:], i+1)

create_patches(blocks, 0)

But I feel that I am probably missing an easier way to do this. Thanks

jranieri-grammatech commented 1 year ago

gtirb-rewriting understands how to peek through partial and get the underlying function's constraints, so you could do something like this:

from functools import partial

@patch_constraints(x86_syntax=X86Syntax.INTEL)
def mov_patch(ctx, value):
    return f"movq ${value}, %rax"

for i, block in enumerate(blocks):
    context.register_insert(
        SingleBlockScope(block, BlockPosition.ENTRY),
        Patch.from_function(partial(mov_patch, value=i))
    )

Alternatively you can subclass Patch instead of using Patch.from_function:

class MovRAXPatch(Patch):
    def __init__(self, value: int):
        self._value = value
        super().__init__(Constraints(x86_syntax=X86Syntax.INTEL))

    def get_asm(self, insertion_context: InsertionContext) -> str:
        return f"movq ${self._value}, %rax"

for i, block in enumerate(blocks):
    context.register_insert(
        SingleBlockScope(block, BlockPosition.ENTRY),
        MovRAXPatch(value=i)
    )

My browser has been trying to autocorrect the code, so it's possible something minor is amiss above, but that should get you in the right direction. My personal thought for when to use a class versus partial comes down to how much state you need to pass and how reusable the patch is. The more of either of those, the more I'm likely to use a class.

avncharlie commented 1 year ago

Thanks, that makes sense!