From 67578f14b730a2856e3a7218727aa64bd774012d Mon Sep 17 00:00:00 2001
From: Ali Mashtizadeh <mashti@uwaterloo.ca>
Date: Wed, 30 Aug 2023 12:45:51 -0400
Subject: [PATCH] Add lock and malloc annotations to the code.

---
 sys/include/cdefs.h    | 23 ++++++++++++++++++++++-
 sys/include/kmem.h     |  2 +-
 sys/include/spinlock.h | 11 +++++++----
 sys/kern/spinlock.c    |  4 ++--
 sys/kern/thread.c      |  2 +-
 5 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/sys/include/cdefs.h b/sys/include/cdefs.h
index 5445f23..4236ee3 100644
--- a/sys/include/cdefs.h
+++ b/sys/include/cdefs.h
@@ -8,12 +8,33 @@
 #define PERCORE		__attribute__((section(".data.percore")))
 
 #define INLINE		inline
-#define NO_RETURN	__attribute__((noreturn))
+#define ALWAYS_INLINE	__attribute__((__always_inline__))
 
+#define NO_RETURN	__attribute__((noreturn))
 #define UNREACHABLE	__builtin_unreachable
 
 #define ROUNDUP(_x, _n) (((_x) + (_n) - 1) & ~((_n) - 1))
 
+#define __MALLOC	__attribute__((__malloc__))
+
+#if __has_extension(c_thread_safety_attributes)
+#define __NO_LOCK_ANALYSIS	__attribute__((no_thread_safety_analysis))
+#define __LOCKABLE		__attribute__((lockable))
+#define __LOCK_EX(_x)		__attribute__((exclusive_lock_function(_x)))
+#define __TRYLOCK_EX(_x)	__attribute__((exclusive_trylock_function(_x)))
+#define __UNLOCK_EX(_x)		__attribute__((unlock_function(_x)))
+#define __LOCK_EX_ASSERT(_x)	__attribute__((assert_exclusive_lock(_x)))
+#define __GUARDED_BY(_x)	__attribute__((guarded_by(_x)))
+#else
+#define __NO_LOCK_ANALYSIS
+#define __LOCKABLE
+#define __LOCK_EX(_x)
+#define __TRYLOCK_EX(_x)
+#define __UNLOCK_EX(_x)
+#define __LOCK_EX_ASSERT(_x)
+#define __GUARDED_BY(_x)
+#endif
+
 #define __hidden	__attribute__((__visibility__("hidden")))
 
 #define __printflike(_fmt, _var) __attribute__((__format__(__printf__, _fmt, _var)))
diff --git a/sys/include/kmem.h b/sys/include/kmem.h
index bff08d1..42c0cc0 100644
--- a/sys/include/kmem.h
+++ b/sys/include/kmem.h
@@ -59,7 +59,7 @@ typedef struct Slab {
 } Slab;
 
 void Slab_Init(Slab *slab, const char *name, uintptr_t objsz, uintptr_t align);
-void *Slab_Alloc(Slab *slab);
+void *Slab_Alloc(Slab *slab) __attribute__((malloc));
 void Slab_Free(Slab *slab, void *obj);
 
 #define DECLARE_SLAB(_type) \
diff --git a/sys/include/spinlock.h b/sys/include/spinlock.h
index 7d7d0d9..b5de060 100644
--- a/sys/include/spinlock.h
+++ b/sys/include/spinlock.h
@@ -2,6 +2,9 @@
 #ifndef __SPINLOCK_H__
 #define __SPINLOCK_H__
 
+#include <stdint.h>
+
+#include <sys/cdefs.h>
 #include <sys/queue.h>
 
 #define SPINLOCK_NAMELEN    32
@@ -22,7 +25,7 @@ typedef struct Spinlock
     char		    name[SPINLOCK_NAMELEN];
     LIST_ENTRY(Spinlock)    lockList;
     TAILQ_ENTRY(Spinlock)   lockStack;
-} Spinlock;
+} __LOCKABLE Spinlock;
 
 void Critical_Init();
 void Critical_Enter();
@@ -31,9 +34,9 @@ void Critical_Exit();
 void Spinlock_EarlyInit();
 void Spinlock_Init(Spinlock *lock, const char *name, uint64_t type);
 void Spinlock_Destroy(Spinlock *lock);
-void Spinlock_Lock(Spinlock *lock);
-void Spinlock_Unlock(Spinlock *lock);
-bool Spinlock_IsHeld(Spinlock *lock);
+void Spinlock_Lock(Spinlock *lock) __LOCK_EX(*lock);
+void Spinlock_Unlock(Spinlock *lock) __UNLOCK_EX(*lock);
+bool Spinlock_IsHeld(Spinlock *lock) __LOCK_EX_ASSERT(*lock);
 
 #endif /* __SPINLOCK_H__ */
 
diff --git a/sys/kern/spinlock.c b/sys/kern/spinlock.c
index 984dfee..59d33ad 100644
--- a/sys/kern/spinlock.c
+++ b/sys/kern/spinlock.c
@@ -62,7 +62,7 @@ Spinlock_Destroy(Spinlock *lock)
 }
 
 void
-Spinlock_Lock(Spinlock *lock)
+Spinlock_Lock(Spinlock *lock) __NO_LOCK_ANALYSIS
 {
     uint64_t startTSC;
     Critical_Enter();
@@ -91,7 +91,7 @@ Spinlock_Lock(Spinlock *lock)
 }
 
 void
-Spinlock_Unlock(Spinlock *lock)
+Spinlock_Unlock(Spinlock *lock) __NO_LOCK_ANALYSIS
 {
     ASSERT(lock->cpu == CPU());
 
diff --git a/sys/kern/thread.c b/sys/kern/thread.c
index 2f2bc1c..0ac88fb 100644
--- a/sys/kern/thread.c
+++ b/sys/kern/thread.c
@@ -277,7 +277,7 @@ Thread_Wait(Thread *thr, uint64_t tid)
 extern TaskStateSegment64 TSS[MAX_CPUS];
 
 void
-ThreadKThreadEntry(TrapFrame *tf)
+ThreadKThreadEntry(TrapFrame *tf) __NO_LOCK_ANALYSIS
 {
     TSS[CPU()].rsp0 = curProc[CPU()]->kstack + 4096;