Use thread_local! for json/repr stack

Summary:
`thread_local!` with `const` initializer generates slightly inefficient but more or or less good code ([issue in rust](rust-lang/rust#104033)) and roughtly equivalent to hack done here before.

`#[thread_local]` is no longer needed (earlier this was not possible because `thread_local!` `const` initializer was implemented after this code was implemented.)

Now we are down to four features preventing using starlark-rust on stable:

```
#![cfg_attr(rust_nightly, feature(const_type_id))]
#![cfg_attr(rust_nightly, feature(core_intrinsics))]
#![feature(const_mut_refs)]
#![feature(generic_associated_types)]
#![feature(maybe_uninit_write_slice)]
#![feature(ptr_metadata)]
```

Reviewed By: ndmitchell

Differential Revision: D41063194

fbshipit-source-id: 2d2e340c987676887c4f185bd563ac5cb90749f0

Author: stepancheg

starlark/src/lib.rs

@@ -356,7 +356,6 @@
 #![feature(generic_associated_types)]
 #![feature(maybe_uninit_write_slice)]
 #![feature(ptr_metadata)]
-#![feature(thread_local)]
 // Plugins
 #![cfg_attr(feature = "gazebo_lint", feature(plugin))]
 #![cfg_attr(feature = "gazebo_lint", allow(deprecated))] // :(

starlark/src/values/recursive_repr_or_json_guard.rs

@@ -32,94 +32,64 @@ pub(crate) struct JsonStackGuard;
 
 impl Drop for ReprStackGuard {
     fn drop(&mut self) {
-        let mut stack = REPR_STACK.take();
-        let popped = stack.pop();
-        debug_assert!(popped.is_some());
-        REPR_STACK.set(stack);
+        REPR_STACK.with(|repr_stack| {
+            let mut stack = Cell::take(repr_stack);
+            let popped = stack.pop();
+            debug_assert!(popped.is_some());
+            repr_stack.set(stack);
+        })
     }
 }
 
 impl Drop for JsonStackGuard {
     fn drop(&mut self) {
-        let mut stack = JSON_STACK.take();
-        let popped = stack.pop();
-        debug_assert!(popped.is_some());
-        JSON_STACK.set(stack);
+        JSON_STACK.with(|json_stack| {
+            let mut stack = Cell::take(json_stack);
+            let popped = stack.pop();
+            debug_assert!(popped.is_some());
+            json_stack.set(stack);
+        })
     }
 }
 
-/// Release per-thread memory.
-///
-/// `#[thread_local]` is fast, but it [does not call destructor on thread exit][1].
-/// `thread_local!` is somewhat slower.
-///
-/// So we use `#[thread_local]` for fast access and `thread_local!` for releasing memory.
-///
-/// Note this is similar to how `thread_local!` actually implemented,
-/// but this is somewhat more efficient.
-///
-/// [1]: https://play.rust-lang.org/?version=nightly&mode=debug&edition=2021&gist=874a250ed9642d3e1559a8ba7d26abfb
-struct ReleaseMemoryOnThreadExit;
-
-impl Drop for ReleaseMemoryOnThreadExit {
-    fn drop(&mut self) {
-        // We replace per-thread data with empty sets.
-        // So destructors for these fields won't be called,
-        // but that's fine because we don't need to release memory of empty sets.
-        REPR_STACK.take();
-        JSON_STACK.take();
-    }
-}
-
-thread_local! {
-    static RELEASE_MEMORY_ON_THREAD_EXIT: ReleaseMemoryOnThreadExit = ReleaseMemoryOnThreadExit;
-}
-
-/// Register a callback to release memory on thread exit.
-#[cold]
-#[inline(never)]
-fn init_release_memory_on_thread_exit() {
-    RELEASE_MEMORY_ON_THREAD_EXIT.with(|_| {});
-}
-
 /// Returned when `repr` is called recursively and a cycle is detected.
 pub(crate) struct ReprCycle;
 
 /// Returned when `to_json` is called recursively and a cycle is detected.
 pub(crate) struct JsonCycle;
 
-#[thread_local]
-static REPR_STACK: Cell<SmallSet<RawPointer>> = Cell::new(SmallSet::new());
+thread_local! {
+    static REPR_STACK: Cell<SmallSet<RawPointer>> = const { Cell::new(SmallSet::new()) };
+}
 
-#[thread_local]
-static JSON_STACK: Cell<SmallSet<RawPointer>> = Cell::new(SmallSet::new());
+thread_local! {
+    static JSON_STACK: Cell<SmallSet<RawPointer>> = const { Cell::new(SmallSet::new()) };
+}
 
 /// Push a value to the stack, return error if it is already on the stack.
 pub(crate) fn repr_stack_push(value: Value) -> Result<ReprStackGuard, ReprCycle> {
-    let mut stack = REPR_STACK.take();
-    if unlikely(stack.capacity() == 0) {
-        init_release_memory_on_thread_exit();
-    }
-    if unlikely(!stack.insert(value.ptr_value())) {
-        REPR_STACK.set(stack);
-        Err(ReprCycle)
-    } else {
-        REPR_STACK.set(stack);
-        Ok(ReprStackGuard)
-    }
+    REPR_STACK.with(|repr_stack| {
+        let mut stack = Cell::take(repr_stack);
+        if unlikely(!stack.insert(value.ptr_value())) {
+            repr_stack.set(stack);
+            Err(ReprCycle)
+        } else {
+            repr_stack.set(stack);
+            Ok(ReprStackGuard)
+        }
+    })
 }
 
 /// Push a value to the stack, return error if it is already on the stack.
 pub(crate) fn json_stack_push(value: Value) -> Result<JsonStackGuard, JsonCycle> {
-    let mut stack = JSON_STACK.take();
-    if unlikely(stack.capacity() == 0) {
-        init_release_memory_on_thread_exit();
-    }
-    if unlikely(!stack.insert(value.ptr_value())) {
-        JSON_STACK.set(stack);
-        Err(JsonCycle)
-    } else {
-        JSON_STACK.set(stack);
-        Ok(JsonStackGuard)
-    }
+    JSON_STACK.with(|json_stack| {
+        let mut stack = Cell::take(json_stack);
+        if unlikely(!stack.insert(value.ptr_value())) {
+            json_stack.set(stack);
+            Err(JsonCycle)
+        } else {
+            json_stack.set(stack);
+            Ok(JsonStackGuard)
+        }
+    })
 }