[libc] at_quick_exit function implemented (#94317)

- added at_quick_exit function 
- used helper file exit_handler which reuses code from atexit
- atexit now calls helper functions from exit_handler
- test cases and dependencies are added

---------

Co-authored-by: Aaryan Shukla <aaryanshukla@google.com>

2 files changed, 104 insertions, 1 deletions

diff --git a/libc/test/src/stdlib/CMakeLists.txt b/libc/test/src/stdlib/CMakeLists.txt
index 6a7faed..3848877 100644
--- a/libc/test/src/stdlib/CMakeLists.txt
+++ b/libc/test/src/stdlib/CMakeLists.txt
@@ -354,7 +354,20 @@ if(LLVM_LIBC_FULL_BUILD)
       libc.src.stdlib.exit
       libc.src.stdlib.atexit
       libc.src.__support.CPP.array
-      libc.src.__support.CPP.utility
+  )
+
+  add_libc_test(
+    at_quick_exit_test
+    # The EXPECT_EXITS test is only availible for unit tests.
+    UNIT_TEST_ONLY
+    SUITE
+      libc-stdlib-tests
+    SRCS
+      at_quick_exit_test.cpp
+    DEPENDS
+      libc.src.stdlib.quick_exit
+      libc.src.stdlib.at_quick_exit
+      libc.src.__support.CPP.array
   )
 
   add_libc_test(
diff --git a/libc/test/src/stdlib/at_quick_exit_test.cpp b/libc/test/src/stdlib/at_quick_exit_test.cpp
new file mode 100644
index 0000000..e0a258d
--- /dev/null
+++ b/libc/test/src/stdlib/at_quick_exit_test.cpp
@@ -0,0 +1,90 @@
+//===-- Unittests for at_quick_exit ---------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "src/__support/CPP/array.h"
+#include "src/__support/CPP/utility.h"
+#include "src/stdlib/at_quick_exit.h"
+#include "src/stdlib/quick_exit.h"
+#include "test/UnitTest/Test.h"
+
+static int a;
+TEST(LlvmLibcAtQuickExit, Basic) {
+  // In case tests ever run multiple times.
+  a = 0;
+
+  auto test = [] {
+    int status = LIBC_NAMESPACE::at_quick_exit(+[] {
+      if (a != 1)
+        __builtin_trap();
+    });
+    status |= LIBC_NAMESPACE::at_quick_exit(+[] { a++; });
+    if (status)
+      __builtin_trap();
+
+    LIBC_NAMESPACE::quick_exit(0);
+  };
+  EXPECT_EXITS(test, 0);
+}
+
+TEST(LlvmLibcAtQuickExit, AtQuickExitCallsSysExit) {
+  auto test = [] {
+    LIBC_NAMESPACE::at_quick_exit(+[] { _Exit(1); });
+    LIBC_NAMESPACE::quick_exit(0);
+  };
+  EXPECT_EXITS(test, 1);
+}
+
+static int size;
+static LIBC_NAMESPACE::cpp::array<int, 256> arr;
+
+template <int... Ts>
+void register_at_quick_exit_handlers(
+    LIBC_NAMESPACE::cpp::integer_sequence<int, Ts...>) {
+  (LIBC_NAMESPACE::at_quick_exit(+[] { arr[size++] = Ts; }), ...);
+}
+
+template <int count> constexpr auto get_test() {
+  return [] {
+    LIBC_NAMESPACE::at_quick_exit(+[] {
+      if (size != count)
+        __builtin_trap();
+      for (int i = 0; i < count; i++)
+        if (arr[i] != count - 1 - i)
+          __builtin_trap();
+    });
+    register_at_quick_exit_handlers(
+        LIBC_NAMESPACE::cpp::make_integer_sequence<int, count>{});
+    LIBC_NAMESPACE::quick_exit(0);
+  };
+}
+
+TEST(LlvmLibcAtQuickExit, ReverseOrder) {
+  // In case tests ever run multiple times.
+  size = 0;
+
+  auto test = get_test<32>();
+  EXPECT_EXITS(test, 0);
+}
+
+TEST(LlvmLibcAtQuickExit, Many) {
+  // In case tests ever run multiple times.
+  size = 0;
+
+  auto test = get_test<256>();
+  EXPECT_EXITS(test, 0);
+}
+
+TEST(LlvmLibcAtQuickExit, HandlerCallsAtQuickExit) {
+  auto test = [] {
+    LIBC_NAMESPACE::at_quick_exit(+[] {
+      LIBC_NAMESPACE::at_quick_exit(+[] { LIBC_NAMESPACE::quick_exit(1); });
+    });
+    LIBC_NAMESPACE::quick_exit(0);
+  };
+  EXPECT_EXITS(test, 1);
+}