From 8b03c6191c40b4fae9ed09fdc4b69b313be973a3 Mon Sep 17 00:00:00 2001
From: Martin Hoffmann <martin@nlnetlabs.nl>
Date: Fri, 11 Oct 2024 17:54:27 +0200
Subject: [PATCH] Add OnceBox::get_or_create.

---
 src/local_array/oncebox.rs | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/src/local_array/oncebox.rs b/src/local_array/oncebox.rs
index f139d41..8598e1c 100644
--- a/src/local_array/oncebox.rs
+++ b/src/local_array/oncebox.rs
@@ -44,6 +44,29 @@ impl<T> OnceBox<T> {
         };
         unsafe { &*res }
     }
+
+    pub fn get_or_create(&self, create: impl FnOnce() -> Box<T>) -> &T {
+        if let Some(res) = self.get() {
+            return res
+        }
+        let ptr = Box::leak(create());
+        let res = match self.ptr.compare_exchange(
+            null_mut(), ptr, Ordering::SeqCst, Ordering::Acquire
+        ) {
+            Ok(current) => {
+                // We set the new value, return it.
+                assert!(current.is_null());
+                ptr as *const _
+            }
+            Err(current) => {
+                // `current` is the real value we need to drop our value.
+                assert!(!current.is_null());
+                let _ = unsafe { Box::from_raw(ptr) };
+                current as *const _
+            }
+        };
+        unsafe { &*res }
+    }
 }
 
 impl<T> Drop for OnceBox<T> {