Convert `heap_allocate` null pointer check to valid tag checks.

Pointers returned by `heap_allocate` should be checked for validity
through the tag bit, not by comparing them with `nullptr`.

Go through all allocation sites of the repository and update where
needed.

Signed-off-by: Hugo Lefeuvre <hugo.lefeuvre@scisemi.com>
diff --git a/sdk/core/scheduler/common.h b/sdk/core/scheduler/common.h
index a0744ed..565152d 100644
--- a/sdk/core/scheduler/common.h
+++ b/sdk/core/scheduler/common.h
@@ -247,6 +247,10 @@
 		HeapObject(struct SObjStruct *heapCapability, T *allocatedObject)
 		  : pointer(allocatedObject, heapCapability)
 		{
+			if (!__builtin_cheri_tag_get(allocatedObject))
+			{
+				pointer = nullptr;
+			}
 		}
 
 		/**
@@ -264,10 +268,14 @@
 		              heap_allocate(timeout, heapCapability, sizeof(T))),
 		            {heapCapability})
 		{
-			if (pointer)
+			if (__builtin_cheri_tag_get(pointer.raw()))
 			{
 				new (pointer.get()) T(std::forward<Args>(args)...);
 			}
+			else
+			{
+				pointer = nullptr;
+			}
 		}
 
 		/**
diff --git a/sdk/core/scheduler/multiwait.h b/sdk/core/scheduler/multiwait.h
index d5630d0..23710dc 100644
--- a/sdk/core/scheduler/multiwait.h
+++ b/sdk/core/scheduler/multiwait.h
@@ -194,7 +194,7 @@
 			                        heapCapability,
 			                        sizeof(MultiWaiterInternal) +
 			                          (length * sizeof(EventWaiter)));
-			if (q == nullptr)
+			if (!__builtin_cheri_tag_get(q))
 			{
 				error = -ENOMEM;
 				return {};
diff --git a/sdk/include/microvium/microvium_port.h b/sdk/include/microvium/microvium_port.h
index 5028346..1b2e4b1 100644
--- a/sdk/include/microvium/microvium_port.h
+++ b/sdk/include/microvium/microvium_port.h
@@ -298,8 +298,13 @@
  */
 #define MVM_CONTEXTUAL_MALLOC(size, context)                                   \
 	({                                                                         \
-		Timeout t = {0, 0};                                                    \
-		heap_allocate(&t, context, size);                                      \
+		Timeout t   = {0, 0};                                                  \
+		void   *ret = heap_allocate(&t, context, size);                        \
+		if (!__builtin_cheri_tag_get(ret))                                     \
+		{                                                                      \
+			ret = NULL;                                                        \
+		}                                                                      \
+		ret;                                                                   \
 	})
 #define MVM_CONTEXTUAL_FREE(ptr, context) heap_free(context, ptr)
 
diff --git a/sdk/include/stdlib.h b/sdk/include/stdlib.h
index da9f2bf..0d802ee 100644
--- a/sdk/include/stdlib.h
+++ b/sdk/include/stdlib.h
@@ -238,13 +238,23 @@
 #ifndef CHERIOT_NO_AMBIENT_MALLOC
 static inline void *malloc(size_t size)
 {
-	Timeout t = {0, 0};
-	return heap_allocate(&t, MALLOC_CAPABILITY, size);
+	Timeout t   = {0, 0};
+	void   *ptr = heap_allocate(&t, MALLOC_CAPABILITY, size);
+	if (!__builtin_cheri_tag_get(ptr))
+	{
+		ptr = NULL;
+	}
+	return ptr;
 }
 static inline void *calloc(size_t nmemb, size_t size)
 {
-	Timeout t = {0, 0};
-	return heap_allocate_array(&t, MALLOC_CAPABILITY, nmemb, size);
+	Timeout t   = {0, 0};
+	void   *ptr = heap_allocate_array(&t, MALLOC_CAPABILITY, nmemb, size);
+	if (!__builtin_cheri_tag_get(ptr))
+	{
+		ptr = NULL;
+	}
+	return ptr;
 }
 static inline int free(void *ptr)
 {
diff --git a/sdk/lib/event_group/event_group.cc b/sdk/lib/event_group/event_group.cc
index dafa1eb..947b2f4 100644
--- a/sdk/lib/event_group/event_group.cc
+++ b/sdk/lib/event_group/event_group.cc
@@ -46,7 +46,7 @@
 	auto   group =
 	  static_cast<EventGroup *>(heap_allocate(timeout, heapCapability, size));
 	*outGroup = group;
-	if (!group)
+	if (!__builtin_cheri_tag_get(group))
 	{
 		return -ENOMEM;
 	}
diff --git a/tests/allocator-test.cc b/tests/allocator-test.cc
index 2bf3926..2a8c26a 100644
--- a/tests/allocator-test.cc
+++ b/tests/allocator-test.cc
@@ -5,6 +5,7 @@
 #define TEST_NAME "Allocator"
 
 #include "tests.hh"
+#include <cheri.hh>
 #include <cheriot-atomic.hh>
 #include <cstdlib>
 #include <debug.hh>
@@ -72,7 +73,7 @@
 				Timeout t{AllocTimeout};
 				allocation = heap_allocate(&t, MALLOC_CAPABILITY, AllocSize);
 				TEST(
-				  allocation != nullptr,
+				  __builtin_cheri_tag_get(allocation),
 				  "Cannot make allocations anymore. Either the revoker is not "
 				  "working or it's too slow");
 			}
@@ -84,7 +85,7 @@
 			for (auto allocation : allocations)
 			{
 				TEST(
-				  __builtin_cheri_tag_get(allocation) == 0,
+				  !__builtin_cheri_tag_get(allocation),
 				  "tag for freed memory {} from allocation {} should be clear",
 				  allocation);
 			}
@@ -140,6 +141,8 @@
 			Timeout t{0};
 			allocation =
 			  heap_allocate(&noWait, MALLOC_CAPABILITY, BigAllocSize);
+			// here test for nullptr (as opposed to the valid tag
+			// bit) because we specifically want to check for OOM
 			if (allocation == nullptr)
 			{
 				memoryExhausted = true;
@@ -148,11 +151,13 @@
 		}
 		TEST(memoryExhausted, "Failed to exhaust memory");
 		debug_log("Trying a non-blocking allocation");
+		// nullptr check because we explicitly want to check for OOM
 		TEST(heap_allocate(&noWait, MALLOC_CAPABILITY, BigAllocSize) == nullptr,
 		     "Non-blocking heap allocation did not return failure with memory "
 		     "exhausted");
 		debug_log("Trying a huge allocation");
 		Timeout forever{UnlimitedTimeout};
+		// nullptr check because we explicitly want to check for OOM
 		TEST(heap_allocate(&forever, MALLOC_CAPABILITY, 1024 * 1024 * 1024) ==
 		       nullptr,
 		     "Non-blocking heap allocation did not return failure on huge "
@@ -165,7 +170,7 @@
 		debug_log("Entering blocking malloc");
 		Timeout t{AllocTimeout};
 		void   *ptr = heap_allocate(&t, MALLOC_CAPABILITY, BigAllocSize);
-		TEST(ptr != nullptr,
+		TEST(__builtin_cheri_tag_get(ptr),
 		     "Failed to make progress on blocking allocation, allocation "
 		     "returned {}",
 		     ptr);
@@ -190,13 +195,12 @@
 		auto                  t    = Timeout(0); /* don't sleep */
 
 		auto doAlloc = [&](size_t sz) {
-			auto p = heap_allocate(&t, MALLOC_CAPABILITY, sz);
+			CHERI::Capability p{heap_allocate(&t, MALLOC_CAPABILITY, sz)};
 
-			if (p != nullptr)
+			if (p.is_valid())
 			{
-				CHERI::Capability pwrap{p};
 				// dlmalloc can give you one granule more.
-				TEST(pwrap.length() == sz || pwrap.length() == sz + 8,
+				TEST(p.length() == sz || p.length() == sz + 8,
 				     "Bad return length");
 				memset(p, 0xCA, sz);
 				allocations.push_back(p);
@@ -328,9 +332,10 @@
 		for (size_t i = 16; i < 256; i <<= 1)
 		{
 			allocated += i;
-			TEST(heap_allocate(&noWait, SECOND_HEAP, i) != nullptr,
-			     "Allocating {} bytes failed",
-			     i);
+			TEST(
+			  __builtin_cheri_tag_get(heap_allocate(&noWait, SECOND_HEAP, i)),
+			  "Allocating {} bytes failed",
+			  i);
 		}
 		debug_log("Quota left after allocating {} bytes: {}",
 		          allocated,
@@ -354,8 +359,10 @@
 		debug_log("Before allocating, quota left: {}",
 		          heap_quota_remaining(SECOND_HEAP));
 		Timeout longTimeout{1000};
-		void   *ptr  = heap_allocate(&longTimeout, SECOND_HEAP, 16);
-		void   *ptr2 = heap_allocate(&longTimeout, SECOND_HEAP, 16);
+		void   *ptr = heap_allocate(&longTimeout, SECOND_HEAP, 16);
+		TEST(__builtin_cheri_tag_get(ptr), "Failed to allocate 16 bytes");
+		void *ptr2 = heap_allocate(&longTimeout, SECOND_HEAP, 16);
+		TEST(__builtin_cheri_tag_get(ptr2), "Failed to allocate 16 bytes");
 		debug_log("After allocating, quota left: {}",
 		          heap_quota_remaining(SECOND_HEAP));
 		static cheriot::atomic<int> state = 0;
@@ -518,7 +525,7 @@
 	Timeout t{5};
 	test_free_all();
 	void *ptr = heap_allocate(&t, STATIC_SEALED_VALUE(secondHeap), 32);
-	TEST(ptr, "Failed to allocate 32 bytes");
+	TEST(__builtin_cheri_tag_get(ptr), "Failed to allocate 32 bytes");
 	TEST(heap_address_is_valid(ptr) == true,
 	     "Heap object incorrectly reported as not heap address");
 	int ret = heap_free(MALLOC_CAPABILITY, ptr);