From 9c5640886d95cba73b10a59a43692c9bae4037fb Mon Sep 17 00:00:00 2001
From: Feoramund <161657516+Feoramund@users.noreply.github.com>
Date: Thu, 5 Jun 2025 14:36:55 -0400
Subject: [PATCH] Add `@(no_sanitize_memory)` proc attribute with MSan
 additions to `base:sanitizer`

---
 base/sanitizer/memory.odin | 74 ++++++++++++++++++++++++++++++++++++++
 src/check_decl.cpp         |  1 +
 src/checker.cpp            |  6 ++++
 src/checker.hpp            |  1 +
 src/entity.cpp             |  1 +
 src/llvm_backend_proc.cpp  |  2 +-
 6 files changed, 84 insertions(+), 1 deletion(-)
 create mode 100644 base/sanitizer/memory.odin

diff --git a/base/sanitizer/memory.odin b/base/sanitizer/memory.odin
new file mode 100644
index 000000000..b16309a49
--- /dev/null
+++ b/base/sanitizer/memory.odin
@@ -0,0 +1,74 @@
+#+no-instrumentation
+package sanitizer
+
+@(private="file")
+MSAN_ENABLED :: .Memory in ODIN_SANITIZER_FLAGS
+
+@(private="file")
+@(default_calling_convention="system")
+foreign {
+	__msan_unpoison :: proc(addr: rawptr, size: uint) ---
+}
+
+/*
+Marks a slice as fully initialized.
+
+Code instrumented with `-sanitize:memory` will be permitted to access any
+address within the slice as if it had already been initialized.
+
+When msan is not enabled this procedure does nothing.
+*/
+memory_unpoison_slice :: proc "contextless" (region: $T/[]$E) {
+	when MSAN_ENABLED {
+		__msan_unpoison(raw_data(region),  size_of(E) * len(region))
+	}
+}
+
+/*
+Marks a pointer as fully initialized.
+
+Code instrumented with `-sanitize:memory` will be permitted to access memory
+within the region the pointer points to as if it had already been initialized.
+
+When msan is not enabled this procedure does nothing.
+*/
+memory_unpoison_ptr :: proc "contextless" (ptr: ^$T) {
+	when MSAN_ENABLED {
+		__msan_unpoison(ptr, size_of(T))
+	}
+}
+
+/*
+Marks the region covering `[ptr, ptr+len)` as fully initialized.
+
+Code instrumented with `-sanitize:memory` will be permitted to access memory
+within this range as if it had already been initialized.
+
+When msan is not enabled this procedure does nothing.
+*/
+memory_unpoison_rawptr :: proc "contextless" (ptr: rawptr, len: int) {
+	when MSAN_ENABLED {
+		__msan_unpoison(ptr, uint(len))
+	}
+}
+
+/*
+Marks the region covering `[ptr, ptr+len)` as fully initialized.
+
+Code instrumented with `-sanitize:memory` will be permitted to access memory
+within this range as if it had already been initialized.
+
+When msan is not enabled this procedure does nothing.
+*/
+memory_unpoison_rawptr_uint :: proc "contextless" (ptr: rawptr, len: uint) {
+	when MSAN_ENABLED {
+		__msan_unpoison(ptr, len)
+	}
+}
+
+memory_unpoison :: proc {
+	memory_unpoison_slice,
+	memory_unpoison_ptr,
+	memory_unpoison_rawptr,
+	memory_unpoison_rawptr_uint,
+}
diff --git a/src/check_decl.cpp b/src/check_decl.cpp
index d53c3c6b7..c696fc4c1 100644
--- a/src/check_decl.cpp
+++ b/src/check_decl.cpp
@@ -1370,6 +1370,7 @@ gb_internal void check_proc_decl(CheckerContext *ctx, Entity *e, DeclInfo *d) {
 	e->Procedure.has_instrumentation = has_instrumentation;
 
 	e->Procedure.no_sanitize_address = ac.no_sanitize_address;
+	e->Procedure.no_sanitize_memory  = ac.no_sanitize_memory;
 
 	e->deprecated_message = ac.deprecated_message;
 	e->warning_message = ac.warning_message;
diff --git a/src/checker.cpp b/src/checker.cpp
index 9bc02cd87..4a5381014 100644
--- a/src/checker.cpp
+++ b/src/checker.cpp
@@ -3776,6 +3776,12 @@ gb_internal DECL_ATTRIBUTE_PROC(proc_decl_attribute) {
 		}
 		ac->no_sanitize_address = true;
 		return true;
+	} else if (name == "no_sanitize_memory") {
+		if (value != nullptr) {
+			error(value, "'%.*s' expects no parameter", LIT(name));
+		}
+		ac->no_sanitize_memory = true;
+		return true;
 	}
 	return false;
 }
diff --git a/src/checker.hpp b/src/checker.hpp
index 0cdfd69ab..dabb7330a 100644
--- a/src/checker.hpp
+++ b/src/checker.hpp
@@ -140,6 +140,7 @@ struct AttributeContext {
 	bool    instrumentation_enter : 1;
 	bool    instrumentation_exit  : 1;
 	bool    no_sanitize_address   : 1;
+	bool    no_sanitize_memory    : 1;
 	bool    rodata                : 1;
 	bool    ignore_duplicates     : 1;
 	u32 optimization_mode; // ProcedureOptimizationMode
diff --git a/src/entity.cpp b/src/entity.cpp
index a16779419..6c0aa6ace 100644
--- a/src/entity.cpp
+++ b/src/entity.cpp
@@ -263,6 +263,7 @@ struct Entity {
 			bool    uses_branch_location       : 1;
 			bool    is_anonymous               : 1;
 			bool    no_sanitize_address        : 1;
+			bool    no_sanitize_memory         : 1;
 		} Procedure;
 		struct {
 			Array<Entity *> entities;
diff --git a/src/llvm_backend_proc.cpp b/src/llvm_backend_proc.cpp
index f51ed2b4d..0a51b5cb5 100644
--- a/src/llvm_backend_proc.cpp
+++ b/src/llvm_backend_proc.cpp
@@ -345,7 +345,7 @@ gb_internal lbProcedure *lb_create_procedure(lbModule *m, Entity *entity, bool i
 		if (build_context.sanitizer_flags & SanitizerFlag_Address && !entity->Procedure.no_sanitize_address) {
 			lb_add_attribute_to_proc(m, p->value, "sanitize_address");
 		}
-		if (build_context.sanitizer_flags & SanitizerFlag_Memory) {
+		if (build_context.sanitizer_flags & SanitizerFlag_Memory && !entity->Procedure.no_sanitize_memory) {
 			lb_add_attribute_to_proc(m, p->value, "sanitize_memory");
 		}
 		if (build_context.sanitizer_flags & SanitizerFlag_Thread) {