rama_utils/str/arcstr/arc_str.rs
1#![expect(
2 // We follow libstd's lead and prefer to define both.
3 clippy::partialeq_ne_impl,
4 // This is a really annoying clippy lint, since it's required for so many cases...
5 clippy::cast_ptr_alignment,
6 // Vendored from upstream `arcstr`: matches stdlib panicking conventions
7 // (capacity overflow, layout failure) and uses inner `#[allow]` attributes
8 // and `# Safety` doc sections in the upstream-idiomatic style.
9 clippy::panic,
10 clippy::panic_in_result_fn,
11 clippy::multiple_unsafe_ops_per_block,
12 clippy::unnecessary_safety_doc,
13 clippy::allow_attributes,
14 reason = "vendored from upstream arcstr; preserve upstream idioms"
15)]
16use core::alloc::Layout;
17use core::mem::{MaybeUninit, align_of, size_of};
18use core::ptr::NonNull;
19#[cfg(not(all(loom, test)))]
20pub(crate) use core::sync::atomic::{AtomicUsize, Ordering};
21#[cfg(all(loom, test))]
22pub(crate) use loom::sync::atomic::{AtomicUsize, Ordering};
23use smol_str::SmolStr;
24
25use crate::std::borrow::Cow;
26use crate::std::borrow::ToOwned as _;
27use crate::std::boxed::Box;
28use crate::std::string::String;
29
30use super::Substr;
31
32/// A better atomically-reference counted string type.
33///
34/// ## Benefits of `ArcStr` over `Arc<str>`
35///
36/// - It's possible to create a const `ArcStr` from a literal via the
37/// [`arcstr!`][crate::str::arcstr::arcstr] macro. This is probably the killer
38/// feature, to be honest.
39///
40/// These "static" `ArcStr`s are zero cost, take no heap allocation, and don't
41/// even need to perform atomic reads/writes when being cloned or dropped (nor
42/// at any other time).
43///
44/// They even get stored in the read-only memory of your executable, which can
45/// be beneficial for performance and memory usage. (In theory your linker may
46/// even dedupe these for you, but usually not)
47///
48/// - `ArcStr`s from `arcstr!` can be turned into `&'static str` safely
49/// at any time using [`ArcStr::as_static`]. (This returns an Option, which is
50/// `None` if the `ArcStr` was not static)
51///
52/// - This should be unsurprising given the literal functionality, but
53/// [`ArcStr::new`] is able to be a `const` function.
54///
55/// - `ArcStr` is thin, e.g. only a single pointer. Great for cases where you
56/// want to keep the data structure lightweight or need to do some FFI stuff
57/// with it.
58///
59/// - `ArcStr` is totally immutable. No need to lose sleep because you're afraid
60/// of code which thinks it has a right to mutate your `Arc`s just because it
61/// holds the only reference...
62///
63/// - Lower reference counting operations are lower overhead because we don't
64/// support `Weak` references. This can be a drawback for some use cases, but
65/// improves performance for the common case of no-weak-refs.
66///
67/// ## What does "zero-cost literals" mean?
68///
69/// In a few places I call the literal arcstrs "zero-cost". No overhead most
70/// accesses accesses (aside from stuff like `as_static` which obviously
71/// requires it). and it imposes a extra branch in both `clone` and `drop`.
72///
73/// This branch in `clone`/`drop` is not on the result of an atomic load, and is
74/// just a normal memory read. This is actually what allows literal/static
75/// `ArcStr`s to avoid needing to perform any atomic operations in those
76/// functions, which seems likely more than cover the cost.
77///
78/// (Additionally, it's almost certain that in the future we'll be able to
79/// reduce the synchronization required for atomic instructions. This is due to
80/// our guarantee of immutability and lack of support for `Weak`.)
81///
82/// # Usage
83///
84/// ## As a `const`
85///
86/// The big unique feature of `ArcStr` is the ability to create static/const
87/// `ArcStr`s. (See [the macro](crate::str::arcstr::arcstr) docs or the [feature
88/// overview][feats]
89///
90/// [feats]: index.html#feature-overview
91///
92/// ```
93/// # use rama_utils::str::arcstr::{ArcStr, arcstr};
94/// const WOW: ArcStr = arcstr!("cool robot!");
95/// assert_eq!(WOW, "cool robot!");
96/// ```
97///
98/// ## As a `str`
99///
100/// (This is not unique to `ArcStr`, but is a frequent source of confusion I've
101/// seen): `ArcStr` implements `Deref<Target = str>`, and so all functions and
102/// methods from `str` work on it, even though we don't expose them on `ArcStr`
103/// directly.
104///
105/// ```
106/// # use rama_utils::str::arcstr::ArcStr;
107/// let s = ArcStr::from("something");
108/// // These go through `Deref`, so they work even though
109/// // there is no `ArcStr::eq_ignore_ascii_case` function
110/// assert!(s.eq_ignore_ascii_case("SOMETHING"));
111/// ```
112///
113/// Additionally, `&ArcStr` can be passed to any function which accepts `&str`.
114/// For example:
115///
116/// ```
117/// # use rama_utils::str::arcstr::ArcStr;
118/// fn accepts_str(s: &str) {
119/// # _ = s;
120/// // s...
121/// }
122///
123/// let test_str: ArcStr = "test".into();
124/// // This works even though `&test_str` is normally an `&ArcStr`
125/// accepts_str(&test_str);
126///
127/// // Of course, this works for functionality from the standard library as well.
128/// let test_but_loud = ArcStr::from("TEST");
129/// assert!(test_str.eq_ignore_ascii_case(&test_but_loud));
130/// ```
131#[repr(transparent)]
132pub struct ArcStr(NonNull<ThinInner>);
133
134unsafe impl Sync for ArcStr {}
135unsafe impl Send for ArcStr {}
136
137impl ArcStr {
138 /// Construct a new empty string.
139 ///
140 /// # Examples
141 ///
142 /// ```
143 /// # use rama_utils::str::arcstr::ArcStr;
144 /// let s = ArcStr::new();
145 /// assert_eq!(s, "");
146 /// ```
147 #[inline]
148 #[must_use]
149 pub const fn new() -> Self {
150 EMPTY
151 }
152
153 /// Attempt to copy the provided string into a newly allocated `ArcStr`, but
154 /// return `None` if we cannot allocate the required memory.
155 ///
156 /// # Examples
157 ///
158 /// ```
159 /// # use rama_utils::str::arcstr::ArcStr;
160 ///
161 /// # fn do_stuff_with(s: ArcStr) {}
162 ///
163 /// let some_big_str = "please pretend this is a very long string";
164 /// if let Some(s) = ArcStr::try_alloc(some_big_str) {
165 /// do_stuff_with(s);
166 /// } else {
167 /// // Complain about allocation failure, somehow.
168 /// }
169 /// ```
170 #[inline]
171 #[must_use]
172 pub fn try_alloc(copy_from: &str) -> Option<Self> {
173 if let Ok(inner) = ThinInner::try_allocate(copy_from, false) {
174 Some(Self(inner))
175 } else {
176 None
177 }
178 }
179
180 /// Attempt to allocate memory for an [`ArcStr`] of length `n`, and use the
181 /// provided callback to fully initialize the provided buffer with valid
182 /// UTF-8 text.
183 ///
184 /// This function returns `None` if memory allocation fails, see
185 /// [`ArcStr::init_with_unchecked`] for a version which calls
186 /// [`handle_alloc_error`](std::alloc::handle_alloc_error).
187 ///
188 /// # Safety
189 /// The provided `initializer` callback must fully initialize the provided
190 /// buffer with valid UTF-8 text.
191 ///
192 /// # Examples
193 ///
194 /// ```
195 /// # use rama_utils::str::arcstr::ArcStr;
196 /// # use core::mem::MaybeUninit;
197 /// let arcstr = unsafe {
198 /// ArcStr::try_init_with_unchecked(10, |s: &mut [MaybeUninit<u8>]| {
199 /// s.fill(MaybeUninit::new(b'a'));
200 /// }).unwrap()
201 /// };
202 /// assert_eq!(arcstr, "aaaaaaaaaa")
203 /// ```
204 #[inline]
205 pub unsafe fn try_init_with_unchecked<F>(n: usize, initializer: F) -> Option<Self>
206 where
207 F: FnOnce(&mut [MaybeUninit<u8>]),
208 {
209 if let Ok(inner) =
210 // SAFETY: contract requests callee to ensure buffer is fully initialized
211 unsafe { ThinInner::try_allocate_with(n, false, AllocInit::Uninit, initializer) }
212 {
213 Some(Self(inner))
214 } else {
215 None
216 }
217 }
218
219 /// Allocate memory for an [`ArcStr`] of length `n`, and use the provided
220 /// callback to fully initialize the provided buffer with valid UTF-8 text.
221 ///
222 /// This function calls
223 /// [`handle_alloc_error`](std::alloc::handle_alloc_error) if memory
224 /// allocation fails, see [`ArcStr::try_init_with_unchecked`] for a version
225 /// which returns `None`
226 ///
227 /// # Safety
228 /// The provided `initializer` callback must fully initialize the provided
229 /// buffer with valid UTF-8 text.
230 ///
231 /// # Examples
232 ///
233 /// ```
234 /// # use rama_utils::str::arcstr::ArcStr;
235 /// # use core::mem::MaybeUninit;
236 /// let arcstr = unsafe {
237 /// ArcStr::init_with_unchecked(10, |s: &mut [MaybeUninit<u8>]| {
238 /// s.fill(MaybeUninit::new(b'a'));
239 /// })
240 /// };
241 /// assert_eq!(arcstr, "aaaaaaaaaa")
242 /// ```
243 #[inline]
244 pub unsafe fn init_with_unchecked<F>(n: usize, initializer: F) -> Self
245 where
246 F: FnOnce(&mut [MaybeUninit<u8>]),
247 {
248 // SAFETY: contract requests callee to ensure buffer is fully initialized
249 match unsafe { ThinInner::try_allocate_with(n, false, AllocInit::Uninit, initializer) } {
250 Ok(inner) => Self(inner),
251 Err(None) => panic!("capacity overflow"),
252 Err(Some(layout)) => crate::std::alloc::handle_alloc_error(layout),
253 }
254 }
255
256 /// Attempt to allocate memory for an [`ArcStr`] of length `n`, and use the
257 /// provided callback to initialize the provided (initially-zeroed) buffer
258 /// with valid UTF-8 text.
259 ///
260 /// Note: This function is provided with a zeroed buffer, and performs UTF-8
261 /// validation after calling the initializer. While both of these are fast
262 /// operations, some high-performance use cases will be better off using
263 /// [`ArcStr::try_init_with_unchecked`] as the building block.
264 ///
265 /// # Errors
266 /// The provided `initializer` callback must initialize the provided buffer
267 /// with valid UTF-8 text, or a UTF-8 error will be returned.
268 ///
269 /// # Examples
270 ///
271 /// ```
272 /// # use rama_utils::str::arcstr::ArcStr;
273 ///
274 /// let s = ArcStr::init_with(5, |slice| {
275 /// slice
276 /// .iter_mut()
277 /// .zip(b'0'..b'5')
278 /// .for_each(|(db, sb)| *db = sb);
279 /// }).unwrap();
280 /// assert_eq!(s, "01234");
281 /// ```
282 #[inline]
283 pub fn init_with<F>(n: usize, initializer: F) -> Result<Self, core::str::Utf8Error>
284 where
285 F: FnOnce(&mut [u8]),
286 {
287 let mut failed = None::<core::str::Utf8Error>;
288 let wrapper = |zeroed_slice: &mut [MaybeUninit<u8>]| {
289 debug_assert_eq!(n, zeroed_slice.len());
290 // Safety: we pass `AllocInit::Zero`, so this is actually initialized
291 let slice = unsafe {
292 core::slice::from_raw_parts_mut(zeroed_slice.as_mut_ptr().cast::<u8>(), n)
293 };
294 initializer(slice);
295 if let Err(e) = core::str::from_utf8(slice) {
296 failed = Some(e);
297 }
298 };
299 match unsafe { ThinInner::try_allocate_with(n, false, AllocInit::Zero, wrapper) } {
300 Ok(inner) => {
301 // Ensure we clean up the allocation even on error.
302 let this = Self(inner);
303 if let Some(e) = failed {
304 Err(e)
305 } else {
306 Ok(this)
307 }
308 }
309 Err(None) => panic!("capacity overflow"),
310 Err(Some(layout)) => crate::std::alloc::handle_alloc_error(layout),
311 }
312 }
313
314 /// Extract a string slice containing our data.
315 ///
316 /// Note: This is an equivalent to our `Deref` implementation, but can be
317 /// more readable than `&*s` in the cases where a manual invocation of
318 /// `Deref` would be required.
319 ///
320 /// # Examples
321 // TODO: find a better example where `&*` would have been required.
322 /// ```
323 /// # use rama_utils::str::arcstr::ArcStr;
324 /// let s = ArcStr::from("abc");
325 /// assert_eq!(s.as_str(), "abc");
326 /// ```
327 #[inline]
328 #[must_use]
329 pub fn as_str(&self) -> &str {
330 self
331 }
332
333 /// Returns the length of this `ArcStr` in bytes.
334 ///
335 /// # Examples
336 ///
337 /// ```
338 /// # use rama_utils::str::arcstr::ArcStr;
339 /// let a = ArcStr::from("foo");
340 /// assert_eq!(a.len(), 3);
341 /// ```
342 #[inline]
343 #[must_use]
344 pub fn len(&self) -> usize {
345 self.get_inner_len_flag().uint_part()
346 }
347
348 #[inline]
349 fn get_inner_len_flag(&self) -> PackedFlagUint {
350 unsafe { ThinInner::get_len_flag(self.0.as_ptr()) }
351 }
352
353 /// Returns true if this `ArcStr` is empty.
354 ///
355 /// # Examples
356 ///
357 /// ```
358 /// # use rama_utils::str::arcstr::ArcStr;
359 /// assert!(!ArcStr::from("foo").is_empty());
360 /// assert!(ArcStr::new().is_empty());
361 /// ```
362 #[inline]
363 #[must_use]
364 pub fn is_empty(&self) -> bool {
365 self.len() == 0
366 }
367
368 /// Convert us to a `core::string::String`.
369 ///
370 /// This is provided as an inherent method to avoid needing to route through
371 /// the `Display` machinery, but is equivalent to `ToString::to_string`.
372 ///
373 /// # Examples
374 ///
375 /// ```
376 /// # use rama_utils::str::arcstr::ArcStr;
377 /// let s = ArcStr::from("abc");
378 /// assert_eq!(s.to_string(), "abc");
379 /// ```
380 #[inline]
381 #[allow(clippy::inherent_to_string_shadow_display)]
382 #[must_use]
383 pub fn to_string(&self) -> String {
384 self.as_str().to_owned()
385 }
386
387 /// Extract a byte slice containing the string's data.
388 ///
389 /// # Examples
390 ///
391 /// ```
392 /// # use rama_utils::str::arcstr::ArcStr;
393 /// let foobar = ArcStr::from("foobar");
394 /// assert_eq!(foobar.as_bytes(), b"foobar");
395 /// ```
396 #[inline]
397 #[must_use]
398 pub fn as_bytes(&self) -> &[u8] {
399 let len = self.len();
400 let p = self.0.as_ptr();
401 unsafe {
402 let data = p.cast::<u8>().add(OFFSET_DATA);
403 debug_assert_eq!(core::ptr::addr_of!((*p).data).cast::<u8>(), data);
404 core::slice::from_raw_parts(data, len)
405 }
406 }
407
408 /// Return the raw pointer this `ArcStr` wraps, for advanced use cases.
409 ///
410 /// Note that in addition to the `NonNull` constraint expressed in the type
411 /// signature, we also guarantee the pointer has an alignment of at least 8
412 /// bytes, even on platforms where a lower alignment would be acceptable.
413 ///
414 /// # Examples
415 ///
416 /// ```
417 /// # use rama_utils::str::arcstr::ArcStr;
418 /// let s = ArcStr::from("abcd");
419 /// let p = ArcStr::into_raw(s);
420 /// // Some time later...
421 /// let s = unsafe { ArcStr::from_raw(p) };
422 /// assert_eq!(s, "abcd");
423 /// ```
424 #[inline]
425 #[must_use]
426 pub fn into_raw(this: Self) -> NonNull<()> {
427 let p = this.0;
428 #[allow(clippy::mem_forget)]
429 core::mem::forget(this);
430 p.cast()
431 }
432
433 /// The opposite version of [`Self::into_raw`]. Still intended only for
434 /// advanced use cases.
435 ///
436 /// # Safety
437 ///
438 /// This function must be used on a valid pointer returned from
439 /// [`ArcStr::into_raw`]. Additionally, you must ensure that a given `ArcStr`
440 /// instance is only dropped once.
441 ///
442 /// # Examples
443 ///
444 /// ```
445 /// # use rama_utils::str::arcstr::ArcStr;
446 /// let s = ArcStr::from("abcd");
447 /// let p = ArcStr::into_raw(s);
448 /// // Some time later...
449 /// let s = unsafe { ArcStr::from_raw(p) };
450 /// assert_eq!(s, "abcd");
451 /// ```
452 #[inline]
453 #[must_use]
454 pub unsafe fn from_raw(ptr: NonNull<()>) -> Self {
455 Self(ptr.cast())
456 }
457
458 /// Returns true if the two `ArcStr`s point to the same allocation.
459 ///
460 /// Note that functions like `PartialEq` check this already, so there's
461 /// no performance benefit to doing something like `ArcStr::ptr_eq(&a1, &a2) || (a1 == a2)`.
462 ///
463 /// Caveat: `const`s aren't guaranteed to only occur in an executable a
464 /// single time, and so this may be non-deterministic for `ArcStr` defined
465 /// in a `const` with [`arcstr!`][crate::str::arcstr::arcstr], unless one
466 /// was created by a `clone()` on the other.
467 ///
468 /// # Examples
469 ///
470 /// ```
471 /// use rama_utils::str::arcstr::{ArcStr, arcstr};
472 ///
473 /// let foobar = ArcStr::from("foobar");
474 /// let same_foobar = foobar.clone();
475 /// let other_foobar = ArcStr::from("foobar");
476 /// assert!(ArcStr::ptr_eq(&foobar, &same_foobar));
477 /// assert!(!ArcStr::ptr_eq(&foobar, &other_foobar));
478 ///
479 /// const YET_AGAIN_A_DIFFERENT_FOOBAR: ArcStr = arcstr!("foobar");
480 /// let strange_new_foobar = YET_AGAIN_A_DIFFERENT_FOOBAR.clone();
481 /// let wild_blue_foobar = strange_new_foobar.clone();
482 /// assert!(ArcStr::ptr_eq(&strange_new_foobar, &wild_blue_foobar));
483 /// ```
484 #[inline]
485 #[must_use]
486 pub fn ptr_eq(lhs: &Self, rhs: &Self) -> bool {
487 core::ptr::eq(lhs.0.as_ptr(), rhs.0.as_ptr())
488 }
489
490 /// Returns the number of references that exist to this `ArcStr`. If this is
491 /// a static `ArcStr` (For example, one from
492 /// [`arcstr!`][crate::str::arcstr::arcstr]), returns `None`.
493 ///
494 /// Despite the difference in return type, this is named to match the method
495 /// from the stdlib's Arc:
496 /// [`Arc::strong_count`][std::sync::Arc::strong_count].
497 ///
498 /// If you aren't sure how to handle static `ArcStr` in the context of this
499 /// return value, `ArcStr::strong_count(&s).unwrap_or(usize::MAX)` is
500 /// frequently reasonable.
501 ///
502 /// # Safety
503 ///
504 /// This method by itself is safe, but using it correctly requires extra
505 /// care. Another thread can change the strong count at any time, including
506 /// potentially between calling this method and acting on the result.
507 ///
508 /// However, it may never change from `None` to `Some` or from `Some` to
509 /// `None` for a given `ArcStr` — whether or not it is static is determined
510 /// at construction, and never changes.
511 ///
512 /// # Examples
513 ///
514 /// ### Dynamic ArcStr
515 /// ```
516 /// # use rama_utils::str::arcstr::ArcStr;
517 /// let foobar = ArcStr::from("foobar");
518 /// assert_eq!(Some(1), ArcStr::strong_count(&foobar));
519 /// let also_foobar = ArcStr::clone(&foobar);
520 /// assert_eq!(Some(2), ArcStr::strong_count(&foobar));
521 /// assert_eq!(Some(2), ArcStr::strong_count(&also_foobar));
522 /// ```
523 ///
524 /// ### Static ArcStr
525 /// ```
526 /// # use rama_utils::str::arcstr::{ArcStr, arcstr};
527 /// let baz = arcstr!("baz");
528 /// assert_eq!(None, ArcStr::strong_count(&baz));
529 /// // Similarly:
530 /// assert_eq!(None, ArcStr::strong_count(&ArcStr::default()));
531 /// ```
532 #[inline]
533 #[must_use]
534 pub fn strong_count(this: &Self) -> Option<usize> {
535 let cf = Self::load_count_flag(this, Ordering::Acquire)?;
536 if cf.flag_part() {
537 None
538 } else {
539 Some(cf.uint_part())
540 }
541 }
542
543 /// Safety: Unsafe to use `this` is stored in static memory (check
544 /// `Self::has_static_lenflag`)
545 #[inline]
546 unsafe fn load_count_flag_raw(this: &Self, ord_if_needed: Ordering) -> PackedFlagUint {
547 PackedFlagUint::from_encoded(unsafe { (*this.0.as_ptr()).count_flag.load(ord_if_needed) })
548 }
549
550 #[inline]
551 fn load_count_flag(this: &Self, ord_if_needed: Ordering) -> Option<PackedFlagUint> {
552 if Self::has_static_lenflag(this) {
553 None
554 } else {
555 let count_and_flag = PackedFlagUint::from_encoded(unsafe {
556 (*this.0.as_ptr()).count_flag.load(ord_if_needed)
557 });
558 Some(count_and_flag)
559 }
560 }
561
562 /// Convert the `ArcStr` into a "static" `ArcStr`, even if it was originally
563 /// created from runtime values. The `&'static str` is returned.
564 ///
565 /// This is useful if you want to use [`ArcStr::as_static`] or
566 /// [`ArcStr::is_static`] on a value only known at runtime.
567 ///
568 /// If the `ArcStr` is already static, then this is a noop.
569 ///
570 /// # Caveats
571 /// Calling this function on an ArcStr will cause us to never free it, thus
572 /// leaking it's memory. Doing this excessively can lead to problems.
573 ///
574 /// # Examples
575 /// ```no_run
576 /// # // This isn't run because it needs a leakcheck suppression,
577 /// # // which I can't seem to make work in CI (no symbols for
578 /// # // doctests?). Instead, we test this in tests/arc_str.rs
579 /// # use rama_utils::str::arcstr::ArcStr;
580 /// let s = ArcStr::from("foobar");
581 /// assert!(!ArcStr::is_static(&s));
582 /// assert!(ArcStr::as_static(&s).is_none());
583 ///
584 /// let leaked: &'static str = s.leak();
585 /// assert_eq!(leaked, s);
586 /// assert!(ArcStr::is_static(&s));
587 /// assert_eq!(ArcStr::as_static(&s), Some("foobar"));
588 /// ```
589 #[inline]
590 #[must_use]
591 pub fn leak(&self) -> &'static str {
592 if Self::has_static_lenflag(self) {
593 return unsafe { Self::to_static_unchecked(self) };
594 }
595 let is_static_count = unsafe {
596 // Not sure about ordering, maybe relaxed would be fine.
597 Self::load_count_flag_raw(self, Ordering::Acquire)
598 };
599 if is_static_count.flag_part() {
600 return unsafe { Self::to_static_unchecked(self) };
601 }
602 unsafe { Self::become_static(self, is_static_count.uint_part() == 1) };
603 debug_assert!(Self::is_static(self));
604 unsafe { Self::to_static_unchecked(self) }
605 }
606
607 unsafe fn become_static(this: &Self, is_unique: bool) {
608 if is_unique {
609 // SAFETY: inner pointer is per contract always valid
610 unsafe {
611 core::ptr::addr_of_mut!((*this.0.as_ptr()).count_flag).write(AtomicUsize::new(
612 PackedFlagUint::new_raw(true, 1).encoded_value(),
613 ));
614 }
615 // SAFETY: inner pointer is per contract always valid
616 let lenp = unsafe { core::ptr::addr_of_mut!((*this.0.as_ptr()).len_flag) };
617 // SAFETY: packed flag is per contract always valid,
618 // so reading is fine
619 debug_assert!(!unsafe { lenp.read() }.flag_part());
620 // SAFETY: packed flag is per contract always valid,
621 // so reading & writing is fine
622 unsafe { lenp.write(lenp.read().with_flag(true)) };
623 } else {
624 let flag_bit = PackedFlagUint::new_raw(true, 0).encoded_value();
625 // SAFETY: inner pointer is per contract always valid
626 let atomic_count_flag = unsafe { &*core::ptr::addr_of!((*this.0.as_ptr()).count_flag) };
627 atomic_count_flag.fetch_or(flag_bit, Ordering::Release);
628 }
629 }
630
631 #[inline]
632 unsafe fn to_static_unchecked(this: &Self) -> &'static str {
633 // SAFETY: by mutual contract this operation is fine
634 unsafe { &*Self::str_ptr(this) }
635 }
636
637 #[inline]
638 fn bytes_ptr(this: &Self) -> *const [u8] {
639 let len = this.get_inner_len_flag().uint_part();
640 unsafe {
641 let p: *const ThinInner = this.0.as_ptr();
642 let data = p.cast::<u8>().add(OFFSET_DATA);
643 debug_assert_eq!(core::ptr::addr_of!((*p).data).cast::<u8>(), data,);
644 core::ptr::slice_from_raw_parts(data, len)
645 }
646 }
647
648 #[inline]
649 fn str_ptr(this: &Self) -> *const str {
650 Self::bytes_ptr(this) as *const str
651 }
652
653 /// Returns true if `this` is a "static" ArcStr. For example, if it was
654 /// created from a call to [`arcstr!`][crate::str::arcstr::arcstr]),
655 /// returned by `ArcStr::new`, etc.
656 ///
657 /// Static `ArcStr`s can be converted to `&'static str` for free using
658 /// [`ArcStr::as_static`], without leaking memory — they're static constants
659 /// in the program (somewhere).
660 ///
661 /// # Examples
662 ///
663 /// ```
664 /// # use rama_utils::str::arcstr::{ArcStr, arcstr};
665 /// const STATIC: ArcStr = arcstr!("Electricity!");
666 /// assert!(ArcStr::is_static(&STATIC));
667 ///
668 /// let still_static = arcstr!("Shocking!");
669 /// assert!(ArcStr::is_static(&still_static));
670 /// assert!(
671 /// ArcStr::is_static(&still_static.clone()),
672 /// "Cloned statics are still static"
673 /// );
674 ///
675 /// let nonstatic = ArcStr::from("Grounded...");
676 /// assert!(!ArcStr::is_static(&nonstatic));
677 /// ```
678 #[inline]
679 #[must_use]
680 pub fn is_static(this: &Self) -> bool {
681 // We align this to 16 bytes and keep the `is_static` flags in the same
682 // place. In theory this means that if `cfg(target_feature = "avx")`
683 // (where aligned 16byte loads are atomic), the compiler *could*
684 // implement this function using the equivalent of:
685 // ```
686 // let vec = _mm_load_si128(self.0.as_ptr().cast());
687 // let mask = _mm_movemask_pd(_mm_srli_epi64(vac, 63));
688 // mask != 0
689 // ```
690 // and that's all; one load, no branching. (I don't think it *does*, but
691 // I haven't checked so I'll be optimistic and keep the `#[repr(align)]`
692 // -- hey, maybe the CPU can peephole-optimize it).
693 //
694 // That said, unless I did it in asm, *I* can't implement it that way,
695 // since Rust's semantics don't allow me to make that change
696 // optimization on my own (that load isn't considered atomic, for
697 // example).
698 this.get_inner_len_flag().flag_part()
699 || unsafe { Self::load_count_flag_raw(this, Ordering::Relaxed).flag_part() }
700 }
701
702 /// This is true for any `ArcStr` that has been static from the time when it
703 /// was created. It's cheaper than `has_static_rcflag`.
704 #[inline]
705 fn has_static_lenflag(this: &Self) -> bool {
706 this.get_inner_len_flag().flag_part()
707 }
708
709 /// Returns true if `this` is a "static"/`"literal"` ArcStr. For example, if
710 /// it was created from a call to [`arcstr!`][crate::str::arcstr::arcstr]), returned by
711 /// `ArcStr::new`, etc.
712 ///
713 /// Static `ArcStr`s can be converted to `&'static str` for free using
714 /// [`ArcStr::as_static`], without leaking memory — they're static constants
715 /// in the program (somewhere).
716 ///
717 /// # Examples
718 ///
719 /// ```
720 /// # use rama_utils::str::arcstr::{ArcStr, arcstr};
721 /// const STATIC: ArcStr = arcstr!("Electricity!");
722 /// assert_eq!(ArcStr::as_static(&STATIC), Some("Electricity!"));
723 ///
724 /// // Note that they don't have to be consts, just made using `arcstr!`:
725 /// let still_static = arcstr!("Shocking!");
726 /// assert_eq!(ArcStr::as_static(&still_static), Some("Shocking!"));
727 /// // Cloning a static still produces a static.
728 /// assert_eq!(ArcStr::as_static(&still_static.clone()), Some("Shocking!"));
729 ///
730 /// // But it won't work for strings from other sources.
731 /// let nonstatic = ArcStr::from("Grounded...");
732 /// assert_eq!(ArcStr::as_static(&nonstatic), None);
733 /// ```
734 #[inline]
735 #[must_use]
736 pub fn as_static(this: &Self) -> Option<&'static str> {
737 if Self::is_static(this) {
738 // We know static strings live forever, so they can have a static lifetime.
739 Some(unsafe { &*(this.as_str() as *const str) })
740 } else {
741 None
742 }
743 }
744
745 // Not public API. Exists so the `arcstr!` macro can call it.
746 #[inline]
747 #[doc(hidden)]
748 pub const unsafe fn _private_new_from_static_data<B>(
749 ptr: &'static StaticArcStrInner<B>,
750 ) -> Self {
751 // SAFETY: ThisInner's contract upholds the needed guarantees
752 Self(unsafe { NonNull::new_unchecked(ptr as *const _ as *mut ThinInner) })
753 }
754
755 /// Returns a substr of `self` over the given range.
756 ///
757 /// # Examples
758 ///
759 /// ```
760 /// use rama_utils::str::arcstr::{ArcStr, Substr};
761 ///
762 /// let a = ArcStr::from("abcde");
763 /// let b: Substr = a.substr(2..);
764 ///
765 /// assert_eq!(b, "cde");
766 /// ```
767 ///
768 /// # Panics
769 /// If any of the following are untrue, we panic
770 /// - `range.start() <= range.end()`
771 /// - `range.end() <= self.len()`
772 /// - `self.is_char_boundary(start) && self.is_char_boundary(end)`
773 /// - These can be conveniently verified in advance using
774 /// `self.get(start..end).is_some()` if needed.
775 #[inline]
776 pub fn substr(&self, range: impl core::ops::RangeBounds<usize>) -> Substr {
777 Substr::from_parts(self, range)
778 }
779
780 /// Returns a [`Substr`] of self over the given `&str`.
781 ///
782 /// It is not rare to end up with a `&str` which holds a view into a
783 /// `ArcStr`'s backing data. A common case is when using functionality that
784 /// takes and returns `&str` and are entirely unaware of `arcstr`, for
785 /// example: `str::trim()`.
786 ///
787 /// This function allows you to reconstruct a [`Substr`] from a `&str` which
788 /// is a view into this `ArcStr`'s backing string.
789 ///
790 /// # Examples
791 ///
792 /// ```
793 /// use rama_utils::str::arcstr::{ArcStr, Substr};
794 /// let text = ArcStr::from(" abc");
795 /// let trimmed = text.trim();
796 /// let substr: Substr = text.substr_from(trimmed);
797 /// assert_eq!(substr, "abc");
798 /// // for illustration
799 /// assert!(ArcStr::ptr_eq(substr.parent(), &text));
800 /// assert_eq!(substr.range(), 3..6);
801 /// ```
802 ///
803 /// # Panics
804 ///
805 /// Panics if `substr` isn't a view into our memory.
806 ///
807 /// Also panics if `substr` is a view into our memory but is >= `u32::MAX`
808 /// bytes away from our start, if we're a 64-bit machine and
809 /// `substr-usize-indices` is not enabled.
810 #[must_use]
811 pub fn substr_from(&self, substr: &str) -> Substr {
812 if substr.is_empty() {
813 return Substr::new();
814 }
815
816 let self_start = self.as_ptr() as usize;
817 let self_end = self_start + self.len();
818
819 let substr_start = substr.as_ptr() as usize;
820 let substr_end = substr_start + substr.len();
821 if substr_start < self_start || substr_end > self_end {
822 out_of_range(self, &substr);
823 }
824
825 let index = substr_start - self_start;
826 let end = index + substr.len();
827 self.substr(index..end)
828 }
829
830 /// If possible, returns a [`Substr`] of self over the
831 /// given `&str`.
832 ///
833 /// This is a fallible version of [`ArcStr::substr_from`].
834 ///
835 /// It is not rare to end up with a `&str` which holds a view into a
836 /// `ArcStr`'s backing data. A common case is when using functionality that
837 /// takes and returns `&str` and are entirely unaware of `arcstr`, for
838 /// example: `str::trim()`.
839 ///
840 /// This function allows you to reconstruct a [`Substr`] from a `&str` which
841 /// is a view into this `ArcStr`'s backing string.
842 ///
843 /// # Examples
844 ///
845 /// ```
846 /// use rama_utils::str::arcstr::{ArcStr, Substr};
847 /// let text = ArcStr::from(" abc");
848 /// let trimmed = text.trim();
849 /// let substr: Option<Substr> = text.try_substr_from(trimmed);
850 /// assert_eq!(substr.unwrap(), "abc");
851 /// // `&str`s not derived from `self` will return None.
852 /// let not_substr = text.try_substr_from("abc");
853 /// assert!(not_substr.is_none());
854 /// ```
855 ///
856 /// # Panics
857 ///
858 /// Panics if `substr` is a view into our memory but is >= `u32::MAX` bytes
859 /// away from our start, if we're a 64-bit machine and
860 /// `substr-usize-indices` is not enabled.
861 #[must_use]
862 pub fn try_substr_from(&self, substr: &str) -> Option<Substr> {
863 if substr.is_empty() {
864 return Some(Substr::new());
865 }
866
867 let self_start = self.as_ptr() as usize;
868 let self_end = self_start + self.len();
869
870 let substr_start = substr.as_ptr() as usize;
871 let substr_end = substr_start + substr.len();
872 if substr_start < self_start || substr_end > self_end {
873 return None;
874 }
875
876 let index = substr_start - self_start;
877 let end = index + substr.len();
878 debug_assert!(self.get(index..end).is_some());
879 Some(self.substr(index..end))
880 }
881
882 /// Compute a derived `&str` a function of `&str` =>
883 /// `&str`, and produce a Substr of the result if possible.
884 ///
885 /// The function may return either a derived string, or any empty string.
886 ///
887 /// This function is mainly a wrapper around [`ArcStr::try_substr_from`]. If
888 /// you're coming to `arcstr` from the `shared_string` crate, this is the
889 /// moral equivalent of the `slice_with` function.
890 ///
891 /// # Examples
892 ///
893 /// ```
894 /// use rama_utils::str::arcstr::{ArcStr, Substr};
895 /// let text = ArcStr::from(" abc");
896 /// let trimmed: Option<Substr> = text.try_substr_using(str::trim);
897 /// assert_eq!(trimmed.unwrap(), "abc");
898 /// let other = text.try_substr_using(|_s| "different string!");
899 /// assert_eq!(other, None);
900 /// // As a special case, this is allowed.
901 /// let empty = text.try_substr_using(|_s| "");
902 /// assert_eq!(empty.unwrap(), "");
903 /// ```
904 pub fn try_substr_using(&self, f: impl FnOnce(&str) -> &str) -> Option<Substr> {
905 self.try_substr_from(f(self.as_str()))
906 }
907
908 /// Compute a derived `&str` a function of `&str` =>
909 /// `&str`, and produce a Substr of the result.
910 ///
911 /// The function may return either a derived string, or any empty string.
912 /// Returning anything else will result in a panic.
913 ///
914 /// This function is mainly a wrapper around [`ArcStr::try_substr_from`]. If
915 /// you're coming to `arcstr` from the `shared_string` crate, this is the
916 /// likely closest to the `slice_with_unchecked` function, but this panics
917 /// instead of UB on dodginess.
918 ///
919 /// # Examples
920 ///
921 /// ```
922 /// use rama_utils::str::arcstr::{ArcStr, Substr};
923 /// let text = ArcStr::from(" abc");
924 /// let trimmed: Substr = text.substr_using(str::trim);
925 /// assert_eq!(trimmed, "abc");
926 /// // As a special case, this is allowed.
927 /// let empty = text.substr_using(|_s| "");
928 /// assert_eq!(empty, "");
929 /// ```
930 pub fn substr_using(&self, f: impl FnOnce(&str) -> &str) -> Substr {
931 self.substr_from(f(self.as_str()))
932 }
933
934 /// Creates an `ArcStr` by repeating the source string `n` times
935 ///
936 /// # Errors
937 ///
938 /// This function returns `None` if the capacity overflows or allocation
939 /// fails.
940 ///
941 /// # Examples
942 ///
943 /// ```
944 /// use rama_utils::str::arcstr::ArcStr;
945 ///
946 /// let source = "A";
947 /// let repeated = ArcStr::try_repeat(source, 10);
948 /// assert_eq!(repeated.unwrap(), "AAAAAAAAAA");
949 /// ```
950 #[must_use]
951 pub fn try_repeat(source: &str, n: usize) -> Option<Self> {
952 // If the source string is empty or the user asked for zero repetitions,
953 // return an empty string
954 if source.is_empty() || n == 0 {
955 return Some(Self::new());
956 }
957
958 // Calculate the capacity for the allocated string
959 let capacity = source.len().checked_mul(n)?;
960 let inner =
961 ThinInner::try_allocate_maybe_uninit(capacity, false, AllocInit::Uninit).ok()?;
962
963 unsafe {
964 let mut data_ptr = ThinInner::data_ptr(inner);
965 let data_end = data_ptr.add(capacity);
966
967 // Copy `source` into the allocated string `n` times
968 while data_ptr < data_end {
969 core::ptr::copy_nonoverlapping(source.as_ptr(), data_ptr, source.len());
970 data_ptr = data_ptr.add(source.len());
971 }
972 }
973
974 Some(Self(inner))
975 }
976}
977
978#[cold]
979#[inline(never)]
980fn out_of_range(arc: &ArcStr, substr: &&str) -> ! {
981 let arc_start = arc.as_ptr();
982 let arc_end = arc_start.wrapping_add(arc.len());
983 let substr_start = substr.as_ptr();
984 let substr_end = substr_start.wrapping_add(substr.len());
985 panic!(
986 "ArcStr over ({arc_start:p}..{arc_end:p}) does not contain substr over ({substr_start:p}..{substr_end:p})",
987 );
988}
989
990impl Clone for ArcStr {
991 #[inline]
992 fn clone(&self) -> Self {
993 if !Self::is_static(self) {
994 // From libstd's impl:
995 //
996 // > Using a relaxed ordering is alright here, as knowledge of the
997 // > original reference prevents other threads from erroneously deleting
998 // > the object.
999 //
1000 // See: https://doc.rust-lang.org/src/alloc/sync.rs.html#1073
1001 let n: PackedFlagUint = PackedFlagUint::from_encoded(unsafe {
1002 let step = PackedFlagUint::FALSE_ONE.encoded_value();
1003 (*self.0.as_ptr())
1004 .count_flag
1005 .fetch_add(step, Ordering::Relaxed)
1006 });
1007 // Refcount saturation: if more than `RC_MAX` (~`usize::MAX/2`) live clones of
1008 // this `ArcStr` exist, transition it into the "static" state by setting the
1009 // flag bit. From that point on `Clone` and `Drop` early-return on `is_static`,
1010 // so the backing `ThinInner` is **intentionally leaked** rather than aborting
1011 // the process — rama is proxy-first and prefers staying up over crashing on a
1012 // pathological refcount. Reaching this branch is effectively unreachable in
1013 // practice (it requires ~`usize::MAX/2` live clones of a single ArcStr; on
1014 // 64-bit that is ~2^62 — far beyond any realistic workload), and the upstream
1015 // `arcstr` crate `abort()`s here instead. We emit a `tracing::error!` so the
1016 // (vanishingly unlikely) event leaves a forensic trail.
1017 //
1018 // Note: this also means the `RC_MAX = MAX/2` headroom argument from upstream
1019 // no longer guarantees that concurrent `fetch_add`s cannot wrap past
1020 // `UINT_PART_MAX`; in practice the headroom is still astronomical and any
1021 // workload that exhausts it has bigger problems than a memory leak.
1022 if n.uint_part() > RC_MAX && !n.flag_part() {
1023 let val = PackedFlagUint::new_raw(true, 0).encoded_value();
1024 unsafe {
1025 (*self.0.as_ptr())
1026 .count_flag
1027 .fetch_or(val, Ordering::Release)
1028 };
1029 tracing::error!(
1030 rc = n.uint_part(),
1031 rc_max = RC_MAX,
1032 "ArcStr refcount saturated; transitioning to static (leaking backing storage) instead of aborting",
1033 );
1034 }
1035 }
1036 Self(self.0)
1037 }
1038}
1039const RC_MAX: usize = PackedFlagUint::UINT_PART_MAX / 2;
1040
1041impl Drop for ArcStr {
1042 #[inline]
1043 fn drop(&mut self) {
1044 if Self::is_static(self) {
1045 return;
1046 }
1047 unsafe {
1048 let this = self.0.as_ptr();
1049 let enc = PackedFlagUint::from_encoded(
1050 (*this)
1051 .count_flag
1052 .fetch_sub(PackedFlagUint::FALSE_ONE.encoded_value(), Ordering::Release),
1053 );
1054 // Note: `enc == PackedFlagUint::FALSE_ONE`
1055 if enc == PackedFlagUint::FALSE_ONE {
1056 _ = (*this).count_flag.load(Ordering::Acquire);
1057 ThinInner::destroy_cold(this)
1058 }
1059 }
1060 }
1061}
1062// Caveat on the `static`/`strong` fields: "is_static" indicates if we're
1063// located in static data (as with empty string). is_static being false meanse
1064// we are a normal arc-ed string.
1065//
1066// While `ArcStr` claims to hold a pointer to a `ThinInner`, for the static case
1067// we actually are using a pointer to a `StaticArcStrInner<[u8; N]>`. These have
1068// almost identical layouts, except the static contains a explicit trailing
1069// array, and does not have a `AtomicUsize` The issue is: We kind of want the
1070// static ones to not have any interior mutability, so that `const`s can use
1071// them, and so that they may be stored in read-only memory.
1072//
1073// We do this by keeping a flag in `len_flag` flag to indicate which case we're
1074// in, and maintaining the invariant that if we're a `StaticArcStrInner` **we
1075// may never access `.strong` in any way or produce a `&ThinInner` pointing to
1076// our data**.
1077//
1078// This is more subtle than you might think, sinc AFAIK we're not legally
1079// allowed to create an `&ThinInner` until we're 100% sure it's nonstatic, and
1080// prior to determining it, we are forced to work from entirely behind a raw
1081// pointer...
1082//
1083// That said, a bit of this hoop jumping might be not required in the future,
1084// but for now what we're doing works and is apparently sound:
1085// https://github.com/rust-lang/unsafe-code-guidelines/issues/246
1086#[repr(C, align(8))]
1087struct ThinInner {
1088 // Both of these are `PackedFlagUint`s that store `is_static` as the flag.
1089 //
1090 // The reason it's not just stored in len is because an ArcStr may become
1091 // static after creation (via `ArcStr::leak`) and we don't need to do an
1092 // atomic load to access the length (and not only because it would mess with
1093 // optimization).
1094 //
1095 // The reason it's not just stored in the count is because it may be UB to
1096 // do atomic loads from read-only memory. This is also the reason it's not
1097 // stored in a separate atomic, and why doing an atomic load to access the
1098 // length wouldn't be acceptable even if compilers were really good.
1099 len_flag: PackedFlagUint,
1100 count_flag: AtomicUsize,
1101 data: [u8; 0],
1102}
1103
1104const OFFSET_LENFLAGS: usize = 0;
1105const OFFSET_COUNTFLAGS: usize = size_of::<PackedFlagUint>();
1106const OFFSET_DATA: usize = OFFSET_COUNTFLAGS + size_of::<AtomicUsize>();
1107
1108// Not public API, exists for macros.
1109#[repr(C, align(8))]
1110#[doc(hidden)]
1111pub struct StaticArcStrInner<Buf> {
1112 pub len_flag: usize,
1113 pub count_flag: usize,
1114 pub data: Buf,
1115}
1116
1117#[doc(hidden)]
1118pub const STATIC_COUNT_VALUE: usize = PackedFlagUint::new_raw(true, 1).encoded_value();
1119
1120impl<Buf> StaticArcStrInner<Buf> {
1121 #[doc(hidden)]
1122 #[inline]
1123 #[must_use]
1124 pub const fn encode_len(v: usize) -> Option<usize> {
1125 match PackedFlagUint::new(true, v) {
1126 Some(v) => Some(v.encoded_value()),
1127 None => None,
1128 }
1129 }
1130}
1131
1132const _: [(); size_of::<StaticArcStrInner<[u8; 0]>>()] = [(); 2 * size_of::<usize>()];
1133const _: [(); align_of::<StaticArcStrInner<[u8; 0]>>()] = [(); 8];
1134
1135const _: [(); size_of::<StaticArcStrInner<[u8; 2 * size_of::<usize>()]>>()] =
1136 [(); 4 * size_of::<usize>()];
1137const _: [(); align_of::<StaticArcStrInner<[u8; 2 * size_of::<usize>()]>>()] = [(); 8];
1138
1139const _: [(); size_of::<ThinInner>()] = [(); 2 * size_of::<usize>()];
1140const _: [(); align_of::<ThinInner>()] = [(); 8];
1141
1142const _: [(); align_of::<AtomicUsize>()] = [(); align_of::<usize>()];
1143const _: [(); align_of::<AtomicUsize>()] = [(); size_of::<usize>()];
1144const _: [(); size_of::<AtomicUsize>()] = [(); size_of::<usize>()];
1145
1146const _: [(); align_of::<PackedFlagUint>()] = [(); align_of::<usize>()];
1147const _: [(); size_of::<PackedFlagUint>()] = [(); size_of::<usize>()];
1148
1149#[derive(Clone, Copy, PartialEq, Eq)]
1150#[repr(transparent)]
1151struct PackedFlagUint(usize);
1152impl PackedFlagUint {
1153 const UINT_PART_MAX: usize = (1 << (usize::BITS - 1)) - 1;
1154 /// Encodes `false` as the flag and `1` as the uint. Used for a few things,
1155 /// such as the amount we `fetch_add` by for refcounting, and so on.
1156 const FALSE_ONE: Self = Self::new_raw(false, 1);
1157
1158 #[inline]
1159 const fn new(flag_part: bool, uint_part: usize) -> Option<Self> {
1160 if uint_part > Self::UINT_PART_MAX {
1161 None
1162 } else {
1163 Some(Self::new_raw(flag_part, uint_part))
1164 }
1165 }
1166
1167 #[inline(always)]
1168 const fn new_raw(flag_part: bool, uint_part: usize) -> Self {
1169 Self(flag_part as usize | (uint_part << 1))
1170 }
1171
1172 #[inline(always)]
1173 const fn uint_part(self) -> usize {
1174 self.0 >> 1
1175 }
1176
1177 #[inline(always)]
1178 const fn flag_part(self) -> bool {
1179 (self.0 & 1) != 0
1180 }
1181
1182 #[inline(always)]
1183 const fn from_encoded(v: usize) -> Self {
1184 Self(v)
1185 }
1186
1187 #[inline(always)]
1188 const fn encoded_value(self) -> usize {
1189 self.0
1190 }
1191
1192 #[inline(always)]
1193 #[must_use]
1194 const fn with_flag(self, v: bool) -> Self {
1195 Self(v as usize | self.0)
1196 }
1197}
1198
1199const EMPTY: ArcStr = super::arcstr!("");
1200
1201impl ThinInner {
1202 #[inline]
1203 fn allocate(data: &str, initially_static: bool) -> NonNull<Self> {
1204 match Self::try_allocate(data, initially_static) {
1205 Ok(v) => v,
1206 Err(None) => alloc_overflow(),
1207 Err(Some(layout)) => crate::std::alloc::handle_alloc_error(layout),
1208 }
1209 }
1210
1211 #[inline]
1212 fn data_ptr(this: NonNull<Self>) -> *mut u8 {
1213 unsafe { this.as_ptr().cast::<u8>().add(OFFSET_DATA) }
1214 }
1215
1216 /// Allocates a `ThinInner` where the data segment is uninitialized or
1217 /// zeroed.
1218 ///
1219 /// Returns `Err(Some(layout))` if we failed to allocate that layout, and
1220 /// `Err(None)` for integer overflow when computing layout
1221 fn try_allocate_maybe_uninit(
1222 capacity: usize,
1223 initially_static: bool,
1224 init_how: AllocInit,
1225 ) -> Result<NonNull<Self>, Option<Layout>> {
1226 const ALIGN: usize = align_of::<ThinInner>();
1227
1228 debug_assert_ne!(capacity, 0);
1229 if capacity >= (isize::MAX as usize) - (OFFSET_DATA + ALIGN) {
1230 return Err(None);
1231 }
1232
1233 debug_assert!(Layout::from_size_align(capacity + OFFSET_DATA, ALIGN).is_ok());
1234 let layout = unsafe { Layout::from_size_align_unchecked(capacity + OFFSET_DATA, ALIGN) };
1235 let ptr = match init_how {
1236 AllocInit::Uninit => unsafe { crate::std::alloc::alloc(layout) as *mut Self },
1237 AllocInit::Zero => unsafe { crate::std::alloc::alloc_zeroed(layout) as *mut Self },
1238 };
1239 if ptr.is_null() {
1240 return Err(Some(layout));
1241 }
1242
1243 // we actually already checked this above...
1244 debug_assert!(PackedFlagUint::new(initially_static, capacity).is_some());
1245
1246 let len_flag = PackedFlagUint::new_raw(initially_static, capacity);
1247 debug_assert_eq!(len_flag.uint_part(), capacity);
1248 debug_assert_eq!(len_flag.flag_part(), initially_static);
1249
1250 unsafe {
1251 core::ptr::addr_of_mut!((*ptr).len_flag).write(len_flag);
1252
1253 let initial_count_flag = PackedFlagUint::new_raw(initially_static, 1);
1254 let count_flag: AtomicUsize = AtomicUsize::new(initial_count_flag.encoded_value());
1255 core::ptr::addr_of_mut!((*ptr).count_flag).write(count_flag);
1256
1257 debug_assert_eq!(
1258 (ptr as *const u8).wrapping_add(OFFSET_DATA),
1259 (*ptr).data.as_ptr(),
1260 );
1261
1262 Ok(NonNull::new_unchecked(ptr))
1263 }
1264 }
1265
1266 // returns `Err(Some(l))` if we failed to allocate that layout, and
1267 // `Err(None)` for integer overflow when computing layout.
1268 #[inline]
1269 fn try_allocate(data: &str, initially_static: bool) -> Result<NonNull<Self>, Option<Layout>> {
1270 // Safety: we initialize the whole buffer by copying `data` into it.
1271 unsafe {
1272 // Allocate a enough space to hold the given string
1273 Self::try_allocate_with(
1274 data.len(),
1275 initially_static,
1276 AllocInit::Uninit,
1277 // Copy the given string into the allocation
1278 |uninit_slice| {
1279 debug_assert_eq!(uninit_slice.len(), data.len());
1280 core::ptr::copy_nonoverlapping(
1281 data.as_ptr(),
1282 uninit_slice.as_mut_ptr().cast::<u8>(),
1283 data.len(),
1284 )
1285 },
1286 )
1287 }
1288 }
1289
1290 /// Safety: caller must fully initialize the provided buffer with valid
1291 /// UTF-8 in the `initializer` function (well, you at least need to handle
1292 /// it before giving it back to the user).
1293 #[inline]
1294 unsafe fn try_allocate_with(
1295 len: usize,
1296 initially_static: bool,
1297 init_style: AllocInit,
1298 initializer: impl FnOnce(&mut [core::mem::MaybeUninit<u8>]),
1299 ) -> Result<NonNull<Self>, Option<Layout>> {
1300 // Allocate a enough space to hold the given string
1301 let this = Self::try_allocate_maybe_uninit(len, initially_static, init_style)?;
1302
1303 // SAFETY: initialised above
1304 initializer(unsafe {
1305 core::slice::from_raw_parts_mut(Self::data_ptr(this).cast::<MaybeUninit<u8>>(), len)
1306 });
1307
1308 Ok(this)
1309 }
1310
1311 #[inline]
1312 unsafe fn get_len_flag(p: *const Self) -> PackedFlagUint {
1313 debug_assert_eq!(OFFSET_LENFLAGS, 0);
1314 // SAFETY: valid by mutual conract of ThisInner and the callee
1315 unsafe { *p.cast() }
1316 }
1317
1318 #[cold]
1319 unsafe fn destroy_cold(p: *mut Self) {
1320 // SAFETY: valid by mutual conract of ThisInner and the callee
1321 let lf = unsafe { Self::get_len_flag(p) };
1322 let (is_static, len) = (lf.flag_part(), lf.uint_part());
1323 debug_assert!(!is_static);
1324 let layout = {
1325 let size = len + OFFSET_DATA;
1326 let align = align_of::<Self>();
1327 // SAFETY: valid by mutual conract of ThisInner and the callee
1328 unsafe { Layout::from_size_align_unchecked(size, align) }
1329 };
1330 // SAFETY: valid by mutual conract of ThisInner and the callee
1331 unsafe { crate::std::alloc::dealloc(p as *mut _, layout) };
1332 }
1333}
1334
1335#[derive(Clone, Copy, PartialEq)]
1336enum AllocInit {
1337 Uninit,
1338 Zero,
1339}
1340
1341#[inline(never)]
1342#[cold]
1343fn alloc_overflow() -> ! {
1344 panic!("overflow during Layout computation")
1345}
1346
1347impl From<&str> for ArcStr {
1348 #[inline]
1349 fn from(s: &str) -> Self {
1350 if s.is_empty() {
1351 Self::new()
1352 } else {
1353 Self(ThinInner::allocate(s, false))
1354 }
1355 }
1356}
1357
1358impl TryFrom<&[u8]> for ArcStr {
1359 type Error = core::str::Utf8Error;
1360
1361 #[inline(always)]
1362 fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
1363 let s = core::str::from_utf8(value)?;
1364 Ok(s.into())
1365 }
1366}
1367
1368impl TryFrom<crate::std::vec::Vec<u8>> for ArcStr {
1369 type Error = crate::std::string::FromUtf8Error;
1370
1371 #[inline(always)]
1372 fn try_from(value: crate::std::vec::Vec<u8>) -> Result<Self, Self::Error> {
1373 let s = String::from_utf8(value)?;
1374 Ok(s.into())
1375 }
1376}
1377
1378impl TryFrom<&crate::std::vec::Vec<u8>> for ArcStr {
1379 type Error = core::str::Utf8Error;
1380
1381 #[inline(always)]
1382 fn try_from(value: &crate::std::vec::Vec<u8>) -> Result<Self, Self::Error> {
1383 let s = core::str::from_utf8(value)?;
1384 Ok(s.into())
1385 }
1386}
1387
1388impl From<SmolStr> for ArcStr {
1389 #[inline(always)]
1390 fn from(s: SmolStr) -> Self {
1391 Self::from(s.as_str())
1392 }
1393}
1394
1395impl From<&SmolStr> for ArcStr {
1396 #[inline(always)]
1397 fn from(s: &SmolStr) -> Self {
1398 Self::from(s.as_str())
1399 }
1400}
1401
1402impl core::ops::Deref for ArcStr {
1403 type Target = str;
1404 #[inline]
1405 fn deref(&self) -> &str {
1406 unsafe { core::str::from_utf8_unchecked(self.as_bytes()) }
1407 }
1408}
1409
1410impl Default for ArcStr {
1411 #[inline]
1412 fn default() -> Self {
1413 Self::new()
1414 }
1415}
1416
1417impl From<String> for ArcStr {
1418 #[inline]
1419 fn from(v: String) -> Self {
1420 v.as_str().into()
1421 }
1422}
1423
1424impl From<&mut str> for ArcStr {
1425 #[inline]
1426 fn from(s: &mut str) -> Self {
1427 let s: &str = s;
1428 Self::from(s)
1429 }
1430}
1431
1432impl From<Box<str>> for ArcStr {
1433 #[inline]
1434 fn from(s: Box<str>) -> Self {
1435 Self::from(&s[..])
1436 }
1437}
1438impl From<ArcStr> for Box<str> {
1439 #[inline]
1440 fn from(s: ArcStr) -> Self {
1441 s.as_str().into()
1442 }
1443}
1444impl From<ArcStr> for crate::std::rc::Rc<str> {
1445 #[inline]
1446 fn from(s: ArcStr) -> Self {
1447 s.as_str().into()
1448 }
1449}
1450impl From<ArcStr> for crate::std::Arc<str> {
1451 #[inline]
1452 fn from(s: ArcStr) -> Self {
1453 s.as_str().into()
1454 }
1455}
1456impl From<crate::std::rc::Rc<str>> for ArcStr {
1457 #[inline]
1458 fn from(s: crate::std::rc::Rc<str>) -> Self {
1459 Self::from(&*s)
1460 }
1461}
1462impl From<crate::std::Arc<str>> for ArcStr {
1463 #[inline]
1464 fn from(s: crate::std::Arc<str>) -> Self {
1465 Self::from(&*s)
1466 }
1467}
1468impl<'a> From<Cow<'a, str>> for ArcStr {
1469 #[inline]
1470 fn from(s: Cow<'a, str>) -> Self {
1471 Self::from(&*s)
1472 }
1473}
1474impl<'a> From<&'a ArcStr> for Cow<'a, str> {
1475 #[inline]
1476 fn from(s: &'a ArcStr) -> Self {
1477 Cow::Borrowed(s)
1478 }
1479}
1480
1481impl<'a> From<ArcStr> for Cow<'a, str> {
1482 #[inline]
1483 fn from(s: ArcStr) -> Self {
1484 if let Some(st) = ArcStr::as_static(&s) {
1485 Cow::Borrowed(st)
1486 } else {
1487 Cow::Owned(s.to_string())
1488 }
1489 }
1490}
1491
1492impl From<&String> for ArcStr {
1493 #[inline]
1494 fn from(s: &String) -> Self {
1495 Self::from(s.as_str())
1496 }
1497}
1498impl From<&Self> for ArcStr {
1499 #[inline]
1500 fn from(s: &Self) -> Self {
1501 s.clone()
1502 }
1503}
1504
1505impl core::fmt::Debug for ArcStr {
1506 #[inline]
1507 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
1508 core::fmt::Debug::fmt(self.as_str(), f)
1509 }
1510}
1511
1512impl core::fmt::Display for ArcStr {
1513 #[inline]
1514 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
1515 core::fmt::Display::fmt(self.as_str(), f)
1516 }
1517}
1518
1519impl PartialEq for ArcStr {
1520 #[inline]
1521 fn eq(&self, o: &Self) -> bool {
1522 Self::ptr_eq(self, o) || PartialEq::eq(self.as_str(), o.as_str())
1523 }
1524 #[inline]
1525 fn ne(&self, o: &Self) -> bool {
1526 !Self::ptr_eq(self, o) && PartialEq::ne(self.as_str(), o.as_str())
1527 }
1528}
1529
1530impl Eq for ArcStr {}
1531
1532macro_rules! impl_peq {
1533 (@one $a:ty, $b:ty) => {
1534 #[allow(clippy::extra_unused_lifetimes)]
1535 impl<'a> PartialEq<$b> for $a {
1536 #[inline]
1537 fn eq(&self, s: &$b) -> bool {
1538 PartialEq::eq(&self[..], &s[..])
1539 }
1540 #[inline]
1541 fn ne(&self, s: &$b) -> bool {
1542 PartialEq::ne(&self[..], &s[..])
1543 }
1544 }
1545 };
1546 ($(($a:ty, $b:ty),)+) => {$(
1547 impl_peq!(@one $a, $b);
1548 impl_peq!(@one $b, $a);
1549 )+};
1550}
1551
1552impl_peq! {
1553 (ArcStr, str),
1554 (ArcStr, &'a str),
1555 (ArcStr, String),
1556 (ArcStr, Cow<'a, str>),
1557 (ArcStr, Box<str>),
1558 (ArcStr, crate::std::Arc<str>),
1559 (ArcStr, crate::std::rc::Rc<str>),
1560 (ArcStr, crate::std::Arc<String>),
1561 (ArcStr, crate::std::rc::Rc<String>),
1562}
1563
1564impl PartialOrd for ArcStr {
1565 #[inline]
1566 #[allow(clippy::non_canonical_partial_ord_impl)]
1567 fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
1568 Some(self.cmp(other))
1569 }
1570}
1571
1572impl Ord for ArcStr {
1573 #[inline]
1574 fn cmp(&self, other: &Self) -> core::cmp::Ordering {
1575 self.as_str().cmp(other.as_str())
1576 }
1577}
1578
1579impl core::hash::Hash for ArcStr {
1580 #[inline]
1581 fn hash<H: core::hash::Hasher>(&self, h: &mut H) {
1582 self.as_str().hash(h)
1583 }
1584}
1585
1586macro_rules! impl_index {
1587 ($($IdxT:ty,)*) => {$(
1588 impl core::ops::Index<$IdxT> for ArcStr {
1589 type Output = str;
1590 #[inline]
1591 fn index(&self, i: $IdxT) -> &Self::Output {
1592 &self.as_str()[i]
1593 }
1594 }
1595 )*};
1596}
1597
1598impl_index! {
1599 core::ops::RangeFull,
1600 core::ops::Range<usize>,
1601 core::ops::RangeFrom<usize>,
1602 core::ops::RangeTo<usize>,
1603 core::ops::RangeInclusive<usize>,
1604 core::ops::RangeToInclusive<usize>,
1605}
1606
1607impl AsRef<str> for ArcStr {
1608 #[inline]
1609 fn as_ref(&self) -> &str {
1610 self
1611 }
1612}
1613
1614impl AsRef<[u8]> for ArcStr {
1615 #[inline]
1616 fn as_ref(&self) -> &[u8] {
1617 self.as_bytes()
1618 }
1619}
1620
1621impl core::borrow::Borrow<str> for ArcStr {
1622 #[inline]
1623 fn borrow(&self) -> &str {
1624 self
1625 }
1626}
1627
1628impl core::str::FromStr for ArcStr {
1629 type Err = core::convert::Infallible;
1630 #[inline]
1631 fn from_str(s: &str) -> Result<Self, Self::Err> {
1632 Ok(Self::from(s))
1633 }
1634}
1635
1636#[cfg(test)]
1637mod test {
1638 use super::*;
1639
1640 fn sasi_layout_check<Buf>() {
1641 assert!(align_of::<StaticArcStrInner<Buf>>() >= 8);
1642 assert_eq!(
1643 core::mem::offset_of!(StaticArcStrInner<Buf>, count_flag),
1644 OFFSET_COUNTFLAGS
1645 );
1646 assert_eq!(
1647 core::mem::offset_of!(StaticArcStrInner<Buf>, len_flag),
1648 OFFSET_LENFLAGS
1649 );
1650 assert_eq!(
1651 core::mem::offset_of!(StaticArcStrInner<Buf>, data),
1652 OFFSET_DATA
1653 );
1654 assert_eq!(
1655 core::mem::offset_of!(ThinInner, count_flag),
1656 core::mem::offset_of!(StaticArcStrInner::<Buf>, count_flag),
1657 );
1658 assert_eq!(
1659 core::mem::offset_of!(ThinInner, len_flag),
1660 core::mem::offset_of!(StaticArcStrInner::<Buf>, len_flag),
1661 );
1662 assert_eq!(
1663 core::mem::offset_of!(ThinInner, data),
1664 core::mem::offset_of!(StaticArcStrInner::<Buf>, data),
1665 );
1666 }
1667
1668 #[test]
1669 fn verify_type_pun_offsets_sasi_big_bufs() {
1670 assert_eq!(
1671 core::mem::offset_of!(ThinInner, count_flag),
1672 OFFSET_COUNTFLAGS,
1673 );
1674 assert_eq!(core::mem::offset_of!(ThinInner, len_flag), OFFSET_LENFLAGS);
1675 assert_eq!(core::mem::offset_of!(ThinInner, data), OFFSET_DATA);
1676
1677 assert!(align_of::<ThinInner>() >= 8);
1678
1679 sasi_layout_check::<[u8; 0]>();
1680 sasi_layout_check::<[u8; 1]>();
1681 sasi_layout_check::<[u8; 2]>();
1682 sasi_layout_check::<[u8; 3]>();
1683 sasi_layout_check::<[u8; 4]>();
1684 sasi_layout_check::<[u8; 5]>();
1685 sasi_layout_check::<[u8; 15]>();
1686 sasi_layout_check::<[u8; 16]>();
1687 sasi_layout_check::<[u8; 64]>();
1688 sasi_layout_check::<[u8; 128]>();
1689 sasi_layout_check::<[u8; 1024]>();
1690 sasi_layout_check::<[u8; 4095]>();
1691 sasi_layout_check::<[u8; 4096]>();
1692 }
1693}
1694
1695#[cfg(all(test, loom))]
1696mod loomtest {
1697 use super::ArcStr;
1698 use loom::sync::Arc;
1699 use loom::thread;
1700 #[test]
1701 fn cloning_threads() {
1702 loom::model(|| {
1703 let a = ArcStr::from("abcdefgh");
1704 let addr = a.as_ptr() as usize;
1705
1706 let a1 = Arc::new(a);
1707 let a2 = a1.clone();
1708
1709 let t1 = thread::spawn(move || {
1710 let b: ArcStr = (*a1).clone();
1711 assert_eq!(b.as_ptr() as usize, addr);
1712 });
1713 let t2 = thread::spawn(move || {
1714 let b: ArcStr = (*a2).clone();
1715 assert_eq!(b.as_ptr() as usize, addr);
1716 });
1717
1718 t1.join().unwrap();
1719 t2.join().unwrap();
1720 });
1721 }
1722 #[test]
1723 fn drop_timing() {
1724 loom::model(|| {
1725 let a1 = std::vec![
1726 ArcStr::from("s1"),
1727 ArcStr::from("s2"),
1728 ArcStr::from("s3"),
1729 ArcStr::from("s4"),
1730 ];
1731 let a2 = a1.clone();
1732
1733 let t1 = thread::spawn(move || {
1734 let mut a1 = a1;
1735 while let Some(s) = a1.pop() {
1736 assert!(s.starts_with("s"));
1737 }
1738 });
1739 let t2 = thread::spawn(move || {
1740 let mut a2 = a2;
1741 while let Some(s) = a2.pop() {
1742 assert!(s.starts_with("s"));
1743 }
1744 });
1745
1746 t1.join().unwrap();
1747 t2.join().unwrap();
1748 });
1749 }
1750
1751 #[test]
1752 fn leak_drop() {
1753 loom::model(|| {
1754 let a1 = ArcStr::from("foo");
1755 let a2 = a1.clone();
1756
1757 let t1 = thread::spawn(move || {
1758 drop(a1);
1759 });
1760 let t2 = thread::spawn(move || a2.leak());
1761 t1.join().unwrap();
1762 let leaked: &'static str = t2.join().unwrap();
1763 assert_eq!(leaked, "foo");
1764 });
1765 }
1766}