Add SmallStr

4 files changed, 90 insertions, 1 deletions

diff --git a/compiler/rustc_data_structures/Cargo.toml b/compiler/rustc_data_structures/Cargo.toml
index ad296c97659..b82e9717261 100644
--- a/compiler/rustc_data_structures/Cargo.toml
+++ b/compiler/rustc_data_structures/Cargo.toml
@@ -20,7 +20,7 @@ stable_deref_trait = "1.0.0"
 rayon = { version = "0.3.2", package = "rustc-rayon" }
 rayon-core = { version = "0.3.2", package = "rustc-rayon-core" }
 rustc-hash = "1.1.0"
-smallvec = { version = "1.6.1", features = ["union", "may_dangle"] }
+smallvec = { version = "1.6.1", features = ["const_generics", "union", "may_dangle"] }
 rustc_index = { path = "../rustc_index", package = "rustc_index" }
 bitflags = "1.2.1"
 measureme = "10.0.0"
diff --git a/compiler/rustc_data_structures/src/lib.rs b/compiler/rustc_data_structures/src/lib.rs
index b1e59d65029..1a3fe652521 100644
--- a/compiler/rustc_data_structures/src/lib.rs
+++ b/compiler/rustc_data_structures/src/lib.rs
@@ -80,6 +80,7 @@ pub mod obligation_forest;
 pub mod owning_ref;
 pub mod sip128;
 pub mod small_c_str;
+pub mod small_str;
 pub mod snapshot_map;
 pub mod stable_map;
 pub mod svh;
diff --git a/compiler/rustc_data_structures/src/small_str.rs b/compiler/rustc_data_structures/src/small_str.rs
new file mode 100644
index 00000000000..800acb1b03e
--- /dev/null
+++ b/compiler/rustc_data_structures/src/small_str.rs
@@ -0,0 +1,68 @@
+use smallvec::SmallVec;
+
+#[cfg(test)]
+mod tests;
+
+/// Like SmallVec but for strings.
+#[derive(Default)]
+pub struct SmallStr<const N: usize>(SmallVec<[u8; N]>);
+
+impl<const N: usize> SmallStr<N> {
+    #[inline]
+    pub fn new() -> Self {
+        SmallStr(SmallVec::default())
+    }
+
+    #[inline]
+    pub fn push_str(&mut self, s: &str) {
+        self.0.extend_from_slice(s.as_bytes());
+    }
+
+    #[inline]
+    pub fn empty(&self) -> bool {
+        self.0.is_empty()
+    }
+
+    #[inline]
+    pub fn spilled(&self) -> bool {
+        self.0.spilled()
+    }
+
+    #[inline]
+    pub fn as_str(&self) -> &str {
+        unsafe { std::str::from_utf8_unchecked(self.0.as_slice()) }
+    }
+}
+
+impl<const N: usize> std::ops::Deref for SmallStr<N> {
+    type Target = str;
+
+    #[inline]
+    fn deref(&self) -> &str {
+        self.as_str()
+    }
+}
+
+impl<const N: usize, A: AsRef<str>> FromIterator<A> for SmallStr<N> {
+    #[inline]
+    fn from_iter<T>(iter: T) -> Self
+    where
+        T: IntoIterator<Item = A>,
+    {
+        let mut s = SmallStr::default();
+        s.extend(iter);
+        s
+    }
+}
+
+impl<const N: usize, A: AsRef<str>> Extend<A> for SmallStr<N> {
+    #[inline]
+    fn extend<T>(&mut self, iter: T)
+    where
+        T: IntoIterator<Item = A>,
+    {
+        for a in iter.into_iter() {
+            self.push_str(a.as_ref());
+        }
+    }
+}
diff --git a/compiler/rustc_data_structures/src/small_str/tests.rs b/compiler/rustc_data_structures/src/small_str/tests.rs
new file mode 100644
index 00000000000..7635a9b7204
--- /dev/null
+++ b/compiler/rustc_data_structures/src/small_str/tests.rs
@@ -0,0 +1,20 @@
+use super::*;
+
+#[test]
+fn empty() {
+    let s = SmallStr::<1>::new();
+    assert!(s.empty());
+    assert_eq!("", s.as_str());
+    assert!(!s.spilled());
+}
+
+#[test]
+fn from_iter() {
+    let s = ["aa", "bb", "cc"].iter().collect::<SmallStr<6>>();
+    assert_eq!("aabbcc", s.as_str());
+    assert!(!s.spilled());
+
+    let s = ["aa", "bb", "cc", "dd"].iter().collect::<SmallStr<6>>();
+    assert_eq!("aabbccdd", s.as_str());
+    assert!(s.spilled());
+}