[libc-commits] [libc] [libc][__support/memory_size] fix missing branch and add related tests (PR #83016)

[Schrodinger ZHU Yifan via libc-commits]

https://github.com/SchrodingerZhu created https://github.com/llvm/llvm-project/pull/83016

fix #82644.

>From 7138ecf4994e331168e9656deaeb4cdd273572dc Mon Sep 17 00:00:00 2001
From: Schrodinger ZHU Yifan <yifanzhu at rochester.edu>
Date: Mon, 26 Feb 2024 10:19:17 -0500
Subject: [PATCH] [libc][__support/memory_size] fix missing branch and add
 related tests

---
 libc/src/__support/memory_size.h             |  6 ++-
 libc/src/string/memset_explicit.cpp          | 44 ++++++++++++++++++++
 libc/src/string/memset_explicit.h            | 20 +++++++++
 libc/test/src/__support/memory_size_test.cpp |  8 ++++
 4 files changed, 76 insertions(+), 2 deletions(-)
 create mode 100644 libc/src/string/memset_explicit.cpp
 create mode 100644 libc/src/string/memset_explicit.h

diff --git a/libc/src/__support/memory_size.h b/libc/src/__support/memory_size.h
index 4c7d2079553e88..94aee2520afaa1 100644
--- a/libc/src/__support/memory_size.h
+++ b/libc/src/__support/memory_size.h
@@ -52,9 +52,11 @@ class SafeMemSize {
 
   LIBC_INLINE SafeMemSize operator+(const SafeMemSize &other) {
     type result;
-    if (LIBC_UNLIKELY((value | other.value) < 0))
+    if (LIBC_UNLIKELY((value | other.value) < 0)) {
       result = -1;
-    result = value + other.value;
+    } else {
+      result = value + other.value;
+    }
     return SafeMemSize{result};
   }
 
diff --git a/libc/src/string/memset_explicit.cpp b/libc/src/string/memset_explicit.cpp
new file mode 100644
index 00000000000000..8500216cd3d079
--- /dev/null
+++ b/libc/src/string/memset_explicit.cpp
@@ -0,0 +1,44 @@
+//===-- Implementation of memset_explicit ---------------------------------===//
+//
+// 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/atomic.h"
+#include "src/__support/CPP/cstddef.h"
+#include "src/__support/common.h"
+#include "src/string/memory_utils/inline_memset.h"
+
+namespace LIBC_NAMESPACE {
+
+// Requirement:
+// - the function cannot be inlined
+// - no store to any byte shall be elided
+// - not read can be reordered before the return of the function
+// - all cache shall be invalidated
+// - all byte shall be stored in a uniform manner to avoid side channel attacks
+[[gnu::noinline]] LLVM_LIBC_FUNCTION(void *, memset_explicit,
+                                     (void *dst, int value, size_t count)) {
+  auto ptr = reinterpret_cast<cpp::byte *>(dst);
+  auto byte_value = static_cast<cpp::byte>(value);
+  for (size_t i = 0; i < count; ++i) {
+    ptr[i] = byte_value;
+// Since store should be uniform, the compiler is not allowed to reorder
+// combine or unroll the loop. Put a compiler barrier to effectively
+// stop the compiler from doing such optimization. Unfortunately, clang
+// ethusiastically unroll the loop even if there is a signal fence.
+#ifdef __clang__
+    asm volatile("" : : : "memory");
+#else
+    cpp::atomic_signal_fence(cpp::MemoryOrder::SEQ_CST);
+#endif
+  }
+  // All the store before the fence shall be visible to other threads
+  // No one can hold a valid copy of the original data in their cache.
+  cpp::atomic_thread_fence(cpp::MemoryOrder::SEQ_CST);
+  return dst;
+}
+
+} // namespace LIBC_NAMESPACE
diff --git a/libc/src/string/memset_explicit.h b/libc/src/string/memset_explicit.h
new file mode 100644
index 00000000000000..f6c189761a123c
--- /dev/null
+++ b/libc/src/string/memset_explicit.h
@@ -0,0 +1,20 @@
+//===-- Implementation header for memset_explicit ---------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIBC_SRC_STRING_MEMSET_EXPLICIT_H
+#define LLVM_LIBC_SRC_STRING_MEMSET_EXPLICIT_H
+
+#include <stddef.h> // size_t
+
+namespace LIBC_NAMESPACE {
+
+[[gnu::noinline]] void *memset_explicit(void *ptr, int value, size_t count);
+
+} // namespace LIBC_NAMESPACE
+
+#endif // LLVM_LIBC_SRC_STRING_MEMSET_EXPLICIT_H
diff --git a/libc/test/src/__support/memory_size_test.cpp b/libc/test/src/__support/memory_size_test.cpp
index 93ef3711d40e00..cb7e4f43cc5e2f 100644
--- a/libc/test/src/__support/memory_size_test.cpp
+++ b/libc/test/src/__support/memory_size_test.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "src/__support/memory_size.h"
+#include "test/UnitTest/LibcTest.h"
 #include "test/UnitTest/Test.h"
 
 namespace LIBC_NAMESPACE {
@@ -49,6 +50,13 @@ TEST(LlvmLibcMemSizeTest, Addition) {
   ASSERT_FALSE((max + SafeMemSize{static_cast<size_t>(1)}).valid());
   ASSERT_FALSE((third + third + third + third).valid());
   ASSERT_FALSE((half + half + half).valid());
+
+  ASSERT_FALSE((SafeMemSize{static_cast<size_t>(-1)} +
+                SafeMemSize{static_cast<size_t>(2)})
+                   .valid());
+  ASSERT_FALSE((SafeMemSize{static_cast<size_t>(2)} +
+                SafeMemSize{static_cast<size_t>(-1)})
+                   .valid());
 }
 
 TEST(LlvmLibcMemSizeTest, Multiplication) {