From 6416df1335c502b1ae562952a1a37f5b4da2dbf4 Mon Sep 17 00:00:00 2001
From: Nicolas <Nicolas@nonan.net>
Date: Tue, 11 Nov 2025 15:54:11 +0100
Subject: [PATCH] allowing branching in stencils by patching all
 passing-on-jump instructions and stripping only the last jump instruction if
 its the last instruction in the function

---
 src/copapy/_compiler.py |  5 +++--
 src/copapy/_stencils.py | 11 +++++++++--
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/src/copapy/_compiler.py b/src/copapy/_compiler.py
index 519b99d..32b86ef 100644
--- a/src/copapy/_compiler.py
+++ b/src/copapy/_compiler.py
@@ -314,7 +314,7 @@ def compile_to_dag(node_list: Iterable[Node], sdb: stencil_database) -> tuple[bi
         #print(f"* {node.name} ({offset}) " + ' '.join(f'{d:02X}' for d in data))
 
         for reloc in sdb.get_relocations(node.name, stencil=True):
-            if reloc.target_symbol_info in {'STT_OBJECT', 'STT_NOTYPE'}:
+            if reloc.target_symbol_info in ('STT_OBJECT', 'STT_NOTYPE', 'STT_SECTION'):
                 #print('-- ' + reloc.target_symbol_name + ' // ' + node.name)
                 if reloc.target_symbol_name.startswith('dummy_'):
                     # Patch for write and read addresses to/from heap variables
@@ -323,7 +323,8 @@ def compile_to_dag(node_list: Iterable[Node], sdb: stencil_database) -> tuple[bi
                     obj_addr = object_addr_lookup[associated_net]
                     patch = sdb.get_patch(reloc, obj_addr, offset, binw.Command.PATCH_OBJECT.value)
                 elif reloc.target_symbol_name.startswith('result_'):
-                    raise Exception(f"Stencil {node.name} seems to branch to multiple result_* calls.")
+                    # Set return jump address to address of following stencil
+                    patch = sdb.get_patch(reloc, offset + len(data), offset, binw.Command.PATCH_FUNC.value)
                 else:
                     # Patch constants addresses on heap
                     obj_addr = reloc.target_symbol_offset + section_addr_lookup[reloc.target_section_index]
diff --git a/src/copapy/_stencils.py b/src/copapy/_stencils.py
index ca63547..f060cee 100644
--- a/src/copapy/_stencils.py
+++ b/src/copapy/_stencils.py
@@ -78,7 +78,14 @@ def strip_function(func: elf_symbol) -> bytes:
 
 def get_stencil_position(func: elf_symbol) -> tuple[int, int]:
     start_index = 0  # There must be no prolog
-    end_index = get_last_call_in_function(func)
+    # Find last relocation in function
+    last_instr = get_last_call_in_function(func)
+    function_size = func.fields['st_size']
+    if last_instr + 5 >= function_size:  # Check if jump is last instruction
+        end_index = last_instr  # Jump can be striped
+    else:
+        end_index = function_size
+
     return start_index, end_index
 
 
@@ -216,7 +223,7 @@ class stencil_database():
             symbol_type = symbol_type + 0x03  # Relative to data section
             #print(f" *> {pr.type} {patch_value=} {symbol_address=} {pr.fields['r_addend']=} {pr.bits=}, {function_offset=} {patch_offset=}")
 
-        elif pr.type.endswith('_CALL26'):
+        elif pr.type.endswith('_CALL26') or pr.type.endswith('_JUMP26'):
             # R_AARCH64_CALL26
             # ((S + A) - P) >> 2
             assert pr.file.byteorder == 'little', "Big endian not supported for ARM64"