From fb4df412cee4143bcf7bd4a4d81009fcd9907cf2 Mon Sep 17 00:00:00 2001
From: Nicolas Kruse <nicolas@nonan.net>
Date: Sun, 26 Oct 2025 14:09:45 +0100
Subject: [PATCH] extract_code tool: added patching of function call
 relocations

---
 src/copapy/_compiler.py  |  2 +-
 stencils/aux_functions.c | 15 ++++++++++++++-
 tests/test_vector.py     |  9 +++++----
 tools/extract_code.py    |  2 ++
 4 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/src/copapy/_compiler.py b/src/copapy/_compiler.py
index 329c301..3d697c4 100644
--- a/src/copapy/_compiler.py
+++ b/src/copapy/_compiler.py
@@ -318,7 +318,7 @@ def compile_to_instruction_list(node_list: Iterable[Node], sdb: stencil_database
             else:
                 raise ValueError(f"Unsupported: {name} {patch.target_symbol_info} {patch.target_symbol_name}")
 
-    assert False, aux_function_mem_layout
+    #assert False, aux_function_mem_layout
 
     # write entry function code
     dw.write_com(binw.Command.COPY_CODE)
diff --git a/stencils/aux_functions.c b/stencils/aux_functions.c
index 0cd5cb6..c166acf 100644
--- a/stencils/aux_functions.c
+++ b/stencils/aux_functions.c
@@ -12,7 +12,7 @@ __attribute__((noinline)) int floor_div(float arg1, float arg2) {
     return i;
 }
 
-__attribute__((noinline)) float fast_sqrt(float n) {
+__attribute__((noinline)) float fast_sqrt2(float n) {
     if (n < 0) return -1;
 
     float x = n;             // initial guess
@@ -25,6 +25,10 @@ __attribute__((noinline)) float fast_sqrt(float n) {
     return x;
 }
 
+__attribute__((noinline)) float fast_sqrt(float n) {
+    return n * 3.5 + 4.5;
+}
+
 float fast_pow_float(float base, float exponent) {
     union {
         float f;
@@ -36,4 +40,13 @@ float fast_pow_float(float base, float exponent) {
     int32_t y = (int32_t)(exponent * (x - 1072632447) + 1072632447);
     u.i = (uint32_t)y;
     return u.f;
+}
+
+int main() {
+    // Test aux functions
+    float a = 16.0f;
+    float sqrt_a = fast_sqrt(a);
+    float pow_a = fast_pow_float(a, 0.5f);
+    float sqrt2_a = fast_sqrt2(a);
+    return 0;
 }
\ No newline at end of file
diff --git a/tests/test_vector.py b/tests/test_vector.py
index 90635d2..51ba3a9 100644
--- a/tests/test_vector.py
+++ b/tests/test_vector.py
@@ -15,16 +15,17 @@ def test_compiled_vectors():
     t2 = t1.sum()
 
     t3 = cp.vector(cp.variable(1 / (v + 1)) for v in range(3))
-    #t4 = ((t3 * t1) * 2).magnitude()
-    t4 = ((t3 * t1) * 2).sum()
-
+    t4 = ((t3 * t1) * 2).magnitude()
+    #t4 = ((t3 * t1) * 2).sum()
+    t5 = cp.sqrt(cp.variable(8.0))
 
     tg = cp.Target()
-    tg.compile(t2, t4)
+    tg.compile(t2, t4, t5)
     tg.run()
 
     assert isinstance(t2, cp.variable) and tg.read_value(t2) == 10 + 11 + 12 + 0 + 1 + 2
     #assert isinstance(t4, cp.variable) and tg.read_value(t4) == ((1/1*10 + 1/2*11 + 1/3*12) * 2)**0.5
+    assert isinstance(t5, cp.variable) and tg.read_value(t5) == 8.0 * 3.5 + 4.5
 
 if __name__ == "__main__":
     test_compiled_vectors()
diff --git a/tools/extract_code.py b/tools/extract_code.py
index 71646f7..eaca3bf 100644
--- a/tools/extract_code.py
+++ b/tools/extract_code.py
@@ -47,6 +47,8 @@ if __name__ == "__main__":
             offs = dr.read_int()
             reloc_type = dr.read_int()
             value = dr.read_int(signed=True)
+            assert reloc_type == RelocationType.RELOC_RELATIVE_32.value
+            program_data[offs:offs + 4] = value.to_bytes(4, byteorder, signed=True)
             print(f"PATCH_FUNC patch_offs={offs} reloc_type={reloc_type} value={value}")
         elif com == Command.PATCH_OBJECT:
             offs = dr.read_int()