Review comments addressed

library/std/src/sys/pal/sgx/abi/entry.S

@@ -98,7 +98,7 @@ IMAGE_BASE:
 .equ tcsls_user_r14,            0x50
 .equ tcsls_user_r15,            0x58
 .equ tcsls_tcs_addr,            0x60
-.equ tcsls_tls_ptr,             0x68
+.equ tcsls_tls_data,            0x68
 
 .macro load_tcsls_flag_secondary_bool reg:req comments:vararg
     .ifne tcsls_flag_secondary /* to convert to a bool, must be the first bit */
@@ -353,16 +353,20 @@ get_tcs_addr:
     lfence
     jmp *%r11
 
-.global get_tls_ptr
-get_tls_ptr:
-    mov %gs:tcsls_tls_ptr(,%rdi,8),%rax
+/* Fetches the data available at an offset from tcsls_tls_data. Refer to TlsIndex enum */
+/* for details on the offsets that can be used as an input to this function. */
+.global get_tls_data
+get_tls_data:
+    mov %gs:tcsls_tls_data(,%rdi,8),%rax
     pop %r11
     lfence
     jmp *%r11
 
-.global set_tls_ptr
-set_tls_ptr:
-    mov %rsi,%gs:tcsls_tls_ptr(,%rdi,8)
+/* Sets input data at an offset from tcsls_tls_data. Refer to TlsIndex enum for details on */
+/* the offsets that can be used as an input to this function. */
+.global set_tls_data
+set_tls_data:
+    mov %rsi,%gs:tcsls_tls_data(,%rdi,8)
     pop %r11
     lfence
     jmp *%r11

library/std/src/sys/pal/sgx/abi/mod.rs

@@ -2,7 +2,7 @@
 
 use crate::io::Write;
 use core::arch::global_asm;
-use core::sync::atomic::{AtomicUsize, Ordering, AtomicBool};
+use core::sync::atomic::{AtomicUsize, Ordering};
 use snmalloc_edp::*;
 
 // runtime features
@@ -19,16 +19,12 @@ pub mod usercalls;
 #[cfg(not(test))]
 global_asm!(include_str!("entry.S"), options(att_syntax));
 
-// To initialize global allocator once per program
-static INIT: AtomicBool = AtomicBool::new(false);
-
 #[repr(C)]
 struct EntryReturn(u64, u64);
 
 #[cfg(not(test))]
 #[no_mangle]
 unsafe extern "C" fn tcs_init(secondary: bool) {
-
     // Be very careful when changing this code: it runs before the binary has been
     // relocated. Any indirect accesses to symbols will likely fail.
     const UNINIT: usize = 0;
@@ -48,9 +44,7 @@ unsafe extern "C" fn tcs_init(secondary: bool) {
             reloc::relocate_elf_rela();
 
             // Snmalloc global allocator initialization
-            if !INIT.swap(true, Ordering::Relaxed) {
-                unsafe { sn_global_init(mem::heap_base(), mem::heap_size()); }
-            }
+            unsafe { sn_global_init(mem::heap_base(), mem::heap_size()); }
 
             RELOC_STATE.store(DONE, Ordering::Release);
         }
@@ -77,7 +71,7 @@ extern "C" fn entry(p1: u64, p2: u64, p3: u64, secondary: bool, p4: u64, p5: u64
     let mut allocator = core::mem::MaybeUninit::<snmalloc_edp::Alloc>::uninit();
     unsafe {
         sn_thread_init(allocator.as_mut_ptr());
-        tls::set_tls_ptr(tls::TlsIndex::AllocPtr, allocator.as_mut_ptr() as *const u8);
+        tls::set_tls_data(tls::TlsIndex::AllocPtr, allocator.as_mut_ptr() as *const u8);
     }
 
     // FIXME: how to support TLS in library mode?
@@ -90,7 +84,7 @@ extern "C" fn entry(p1: u64, p2: u64, p3: u64, secondary: bool, p4: u64, p5: u64
         drop(join_notifier);
         drop(tls);
 
-        alloc_thread_cleanup(allocator.as_mut_ptr());
+        alloc_thread_cleanup();
 
         EntryReturn(0, 0)
     } else {
@@ -112,16 +106,17 @@ extern "C" fn entry(p1: u64, p2: u64, p3: u64, secondary: bool, p4: u64, p5: u64
             if ret == 0 {
                 drop(tls_guard);
                 drop(tls);
-                alloc_thread_cleanup(allocator.as_mut_ptr());
+                alloc_thread_cleanup();
             }
             exit_with_code(ret)
         }
     }
 }
 
-pub(super) fn alloc_thread_cleanup(allocator: *mut snmalloc_edp::Alloc) {
+pub(super) fn alloc_thread_cleanup() {
     unsafe {
-        tls::set_tls_ptr(tls::TlsIndex::AllocPtr, core::ptr::null_mut());
+        let allocator = tls::get_tls_data(crate::sys::abi::tls::TlsIndex::AllocPtr) as *mut Alloc;
+        tls::set_tls_data(tls::TlsIndex::AllocPtr, core::ptr::null_mut());
         sn_thread_cleanup(allocator);
     }
 }

library/std/src/sys/pal/sgx/abi/tls/mod.rs

@@ -23,17 +23,20 @@ macro_rules! dup {
 #[export_name = "_ZN16__rust_internals3std3sys3sgx3abi3tls14TLS_DESTRUCTORE"]
 static TLS_DESTRUCTOR: [AtomicUsize; TLS_KEYS] = dup!((* * * * * * *) (AtomicUsize::new(0)));
 
-// Use TCSLS to store both the thread specific allocator
-// and TLS address at different indices
+// Use TCS Local Storage to store various thread specific data at
+// indices definied by the TlsIndex enum below
 #[repr(u8)]
 pub enum TlsIndex {
-    AllocPtr = 0,
-    TlsPtr = 1,
+    AllocPtr = 0, /* Thread specific allocator address */
+    TlsPtr = 1,   /* TLS address */
 }
 
+// Functions to access TLS data from the TCSLS. Use TlsIndex
+// enum to reference the desired field.
+// Function definition available in entry.S
 extern "C" {
-    pub fn get_tls_ptr(index: TlsIndex) -> *const u8;
-    pub fn set_tls_ptr(index: TlsIndex, tls: *const u8);
+    pub fn get_tls_data(index: TlsIndex) -> *const u8;
+    pub fn set_tls_data(index: TlsIndex, tls: *const u8);
 }
 
 #[derive(Copy, Clone)]
@@ -95,21 +98,21 @@ impl Tls {
     }
 
     pub unsafe fn activate(&self) -> ActiveTls<'_> {
-        // FIXME: Needs safety information. See entry.S for `set_tls_ptr` definition.
-        unsafe { set_tls_ptr(TlsIndex::TlsPtr, self as *const Tls as _) };
+        // FIXME: Needs safety information. See entry.S for `set_tls_data` definition.
+        unsafe { set_tls_data(TlsIndex::TlsPtr, self as *const Tls as _) };
         ActiveTls { tls: self }
     }
 
     #[allow(unused)]
     pub unsafe fn activate_persistent(self: Box<Self>) {
-        // FIXME: Needs safety information. See entry.S for `set_tls_ptr` definition.
-        unsafe { set_tls_ptr(TlsIndex::TlsPtr, core::ptr::addr_of!(*self) as _) };
+        // FIXME: Needs safety information. See entry.S for `set_tls_data` definition.
+        unsafe { set_tls_data(TlsIndex::TlsPtr, core::ptr::addr_of!(*self) as _) };
         mem::forget(self);
     }
 
     unsafe fn current<'a>() -> &'a Tls {
-        // FIXME: Needs safety information. See entry.S for `set_tls_ptr` definition.
-        unsafe { &*(get_tls_ptr(TlsIndex::TlsPtr) as *const Tls) }
+        // FIXME: Needs safety information. See entry.S for `set_tls_data` definition.
+        unsafe { &*(get_tls_data(TlsIndex::TlsPtr) as *const Tls) }
     }
 
     pub fn create(dtor: Option<unsafe extern "C" fn(*mut u8)>) -> Key {

library/std/src/sys/pal/sgx/alloc.rs

@@ -44,5 +44,5 @@ pub unsafe extern "C" fn __rust_c_dealloc(ptr: *mut u8, size: usize, align: usiz
 #[cfg(not(test))]
 #[no_mangle]
 pub extern "C" fn __rust_get_thread_allocator() -> *mut Alloc {
-    unsafe{ crate::sys::abi::tls::get_tls_ptr(crate::sys::abi::tls::TlsIndex::AllocPtr) as *mut Alloc }
+    unsafe{ crate::sys::abi::tls::get_tls_data(crate::sys::abi::tls::TlsIndex::AllocPtr) as *mut Alloc }
 }