From af1f8767d858e65be85d4ce02fbd2fd3fe59dce0 Mon Sep 17 00:00:00 2001
From: Edd Barrett <vext01@gmail.com>
Date: Thu, 19 Dec 2024 12:08:10 +0000
Subject: [PATCH] Add/fix tests for new shadow stack allocation scheme.

---
 tests/c/guard_consting.c |  2 +-
 tests/c/shadow_longjmp.c | 68 ++++++++++++++++++++++++++++++++++++++++
 tests/c/simple_peeling.c |  6 ++--
 3 files changed, 72 insertions(+), 4 deletions(-)
 create mode 100644 tests/c/shadow_longjmp.c

diff --git a/tests/c/guard_consting.c b/tests/c/guard_consting.c
index 4b7a7345c..3f78a1f44 100644
--- a/tests/c/guard_consting.c
+++ b/tests/c/guard_consting.c
@@ -4,7 +4,7 @@
 //   stderr:
 //     ...
 //     %{{30}}: i1 = sgt %{{_}}, 0i32
-//     guard true, %{{30}}, [{{0}}:%{{0_2}}: %{{0}}, {{0}}:%{{0_6}}: %{{1}}, {{0}}:%{{0_7}}: %{{2}}, {{0}}:%{{0_8}}: %{{3}}, {{0}}:%{{0_9}}: %{{4}}, {{0}}:%{{9_1}}: 0i1] ; ...
+//     guard true, %{{22}}, [{{0}}:%{{0_8}}: %{{0}}, {{0}}:%{{0_9}}: %{{1}}, {{0}}:%{{0_10}}: %{{2}}, {{0}}:%{{0_11}}: %{{3}}, {{0}}:%{{9_1}}: 0i1] ; ...
 //     ...
 
 // Check that if a guard's life variables include the condition operand, that
diff --git a/tests/c/shadow_longjmp.c b/tests/c/shadow_longjmp.c
new file mode 100644
index 000000000..673730dba
--- /dev/null
+++ b/tests/c/shadow_longjmp.c
@@ -0,0 +1,68 @@
+// Run-time:
+//   env-var: YKD_LOG_IR=aot,jit-pre-opt
+//   env-var: YKD_SERIALISE_COMPILATION=1
+//   env-var: YK_LOG=4
+//   stderr:
+//     11
+//     10
+//     9
+//     longjmp
+//     8
+//     7
+//     6
+//     longjmp
+//     5
+//     4
+//     3
+//     longjmp
+//     2
+//     1
+//     0
+//     longjmp
+
+// Check that the ykllvm shadow stack pass works at runtime in the presence of
+// setjmp and longjump.
+
+#include <assert.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <yk.h>
+#include <yk_testing.h>
+#include <setjmp.h>
+
+__attribute__((noinline))
+void g(int i, jmp_buf *env) {
+  fprintf(stderr, "%d\n", i);
+  if (i % 3 == 0) {
+    fprintf(stderr, "longjmp\n");
+    longjmp(*env, 1);
+  }
+}
+
+__attribute__((noinline))
+void f(int i, jmp_buf *env) {
+  while (i >= 0) {
+    g(i, env);
+    i--;
+  }
+}
+
+int main(int argc, char **argv) {
+  YkMT *mt = yk_mt_new(NULL);
+  jmp_buf env;
+
+  int i = 11;
+  NOOPT_VAL(i);
+  while (i > 0) {
+    // Passing a NULL location. We never JIT. Just checking AOT behaviour.
+    yk_mt_control_point(mt, NULL);
+    if (setjmp(env) != 0) {
+      i -= 3;
+    }
+    f(i, &env);
+    i--;
+  }
+  yk_mt_shutdown(mt);
+  return (EXIT_SUCCESS);
+}
diff --git a/tests/c/simple_peeling.c b/tests/c/simple_peeling.c
index e962820c5..8270ab3ce 100644
--- a/tests/c/simple_peeling.c
+++ b/tests/c/simple_peeling.c
@@ -8,16 +8,16 @@
 //     ...
 //     header_start ...
 //     ...
-//     header_end [%{{0}}, %{{1}}, %{{2}}, %{{3}}, %{{4}}, %{{25}}]
+//     header_end [%{{0}}, %{{1}}, %{{2}}, %{{3}}, %{{4}}]
 //     ...
-//     body_start [%{{30}}, %{{31}}, %{{32}}, %{{33}}, %{{34}}, %{{35}}]
+//     body_start [%{{19}}, %{{20}}, %{{21}}, %{{22}}, %{{23}}]
 //     ...
 //     body_end ...
 //     ...
 //     --- End jit-post-opt ---
 //     ...
 
-// Check that basic trace compilation works.
+// Check that basic trace peeling works.
 
 #include <assert.h>
 #include <stdio.h>