diff --git a/misc/test_runner/CMakeLists.txt b/misc/test_runner/CMakeLists.txt
index 158402c0c..e49cfcd6b 100644
--- a/misc/test_runner/CMakeLists.txt
+++ b/misc/test_runner/CMakeLists.txt
@@ -1,3 +1,3 @@
 add_library(ia2_test_runner STATIC test_runner.c)
 target_include_directories(ia2_test_runner
-    PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include)
+    PRIVATE ${CMAKE_SOURCE_DIR}/runtime/libia2/include)
diff --git a/misc/test_runner/include/ia2_test_runner.h b/misc/test_runner/include/ia2_test_runner.h
index c7e217b51..730bff34b 100644
--- a/misc/test_runner/include/ia2_test_runner.h
+++ b/misc/test_runner/include/ia2_test_runner.h
@@ -1,22 +1,33 @@
 #pragma once
+#include <ia2.h>
 #include <stdio.h>
 #include <stdbool.h>
 #include <stdlib.h>
 #include <unistd.h>
 #include <assert.h>
 
+#if !defined(IA2_TEST_RUNNER_SOURCE)
+typedef void *ia2_test_fn;
+#else
+typedef void (*ia2_test_fn)(void);
+#endif
+
 struct fake_criterion_test {
-  void (*test)(void);
-  void (*init)(void);
+  ia2_test_fn test;
+  ia2_test_fn init;
   int exit_code;
 };
 
+#if !defined(IA2_TEST_RUNNER_SOURCE)
 #define Test(suite, name, ...)                                                                                                 \
+  IA2_BEGIN_NO_WRAP                                                                                                            \
   void fake_criterion_##suite##_##name(void);                                                                                  \
+  IA2_END_NO_WRAP                                                                                                              \
   __attribute__((__section__("fake_criterion_tests"))) struct fake_criterion_test fake_criterion_##suite##_##name##_##test = { \
       .test = fake_criterion_##suite##_##name,                                                                                 \
       ##__VA_ARGS__};                                                                                                          \
   void fake_criterion_##suite##_##name(void)
+#endif
 
 #define cr_log_info(f, ...) printf(f "\n", ##__VA_ARGS__)
 #define cr_log_error(f, ...) fprintf(stderr, f "\n", ##__VA_ARGS__)
diff --git a/misc/test_runner/test_runner.c b/misc/test_runner/test_runner.c
index 85b5cec49..ed568fd6a 100644
--- a/misc/test_runner/test_runner.c
+++ b/misc/test_runner/test_runner.c
@@ -1,3 +1,4 @@
+#define IA2_TEST_RUNNER_SOURCE
 #include "include/ia2_test_runner.h"
 #include <stdbool.h>
 #include <stdio.h>
diff --git a/tools/rewriter/SourceRewriter.cpp b/tools/rewriter/SourceRewriter.cpp
index 89bd8e5e3..70509ef85 100644
--- a/tools/rewriter/SourceRewriter.cpp
+++ b/tools/rewriter/SourceRewriter.cpp
@@ -614,6 +614,7 @@ class FnPtrExpr : public RefactoringCallback {
     auto *fn_ptr_expr = result.Nodes.getNodeAs<clang::DeclRefExpr>("fnPtrExpr");
     assert(fn_ptr_expr != nullptr);
 
+    auto annotation = fn_ptr_expr->getDecl()->getAttr<clang::AnnotateAttr>();
     assert(result.SourceManager != nullptr);
     auto &sm = *result.SourceManager;