Small String Optimization + string sharing for reggie.

reggie/src/machine.rs

@@ -46,15 +46,15 @@ impl std::fmt::Display for DataType {
 pub struct ArgumentRegisters {
     pub f: [f64; ARG_REG_COUNT],
     pub i: [i32; ARG_REG_COUNT],
-    pub s: [Option<LuaString>; ARG_REG_COUNT],
+    pub s: [LuaString; ARG_REG_COUNT],
     pub t: [Option<TableRef>; ARG_REG_COUNT],
     pub d: [LuaValue; ARG_REG_COUNT],
 }
 
 pub struct Accumulators {
     pub f: f64,
     pub i: i32,
-    pub s: Option<LuaString>,
+    pub s: LuaString,
     pub c: BlockID,
     pub t: Option<TableRef>,
     pub d: LuaValue,
@@ -283,7 +283,7 @@ impl Machine {
             accumulators: Accumulators {
                 f: 0.0,
                 i: 0,
-                s: None,
+                s: LuaString::default(),
                 c: dummy_block_id,
                 t: None,
                 d: LuaValue::Nil,
@@ -298,7 +298,7 @@ impl Machine {
             argument_registers: ArgumentRegisters {
                 f: [0.0; ARG_REG_COUNT],
                 i: [0; ARG_REG_COUNT],
-                s: [(); ARG_REG_COUNT].map(|_| None),
+                s: [(); ARG_REG_COUNT].map(|_| LuaString::default()),
                 t: [(); ARG_REG_COUNT].map(|_| None),
                 d: [(); ARG_REG_COUNT].map(|_| LuaValue::Nil),
             },

reggie/src/runtime.rs

@@ -158,11 +158,11 @@ pub fn eval_loop(machine: &mut Machine) -> Result<(), EvalError> {
                 *position += 1;
             }
             Instruction::ConstS(string_id) => {
-                register!(AS) = Some(block.meta.const_strings[string_id].clone());
+                register!(AS) = block.meta.const_strings[string_id].clone();
                 *position += 1;
             }
             Instruction::WrapS => {
-                register!(AD) = LuaValue::String(register!(AS).as_ref().unwrap().clone());
+                register!(AD) = LuaValue::String(register!(AS).clone());
                 *position += 1;
             }
             Instruction::ConstC(local_block_id) => {
@@ -399,7 +399,7 @@ pub fn eval_loop(machine: &mut Machine) -> Result<(), EvalError> {
             }
             Instruction::AssocASD => {
                 let table = register!(AT).as_mut().unwrap();
-                table.assoc_str(register!(AS).clone().unwrap(), register!(AD).clone());
+                table.assoc_str(register!(AS).clone(), register!(AD).clone());
                 *position += 1;
             }
             Instruction::CastT => {
@@ -413,7 +413,7 @@ pub fn eval_loop(machine: &mut Machine) -> Result<(), EvalError> {
             }
             Instruction::TablePropertyLookupError => {
                 return Err(EvalError::from(TypeError::CannotAccessProperty {
-                    property: register!(AS).take().unwrap(),
+                    property: register!(AS).clone(),
                     of: std::mem::replace(&mut register!(AD), LuaValue::Nil),
                 }))
             }
@@ -439,7 +439,7 @@ pub fn eval_loop(machine: &mut Machine) -> Result<(), EvalError> {
                     .t
                     .as_mut()
                     .unwrap()
-                    .get_str_assoc(register!(AS).clone().unwrap());
+                    .get_str_assoc(register!(AS).clone());
                 *position += 1;
             }
             Instruction::LdaAssocAD => {
@@ -499,7 +499,7 @@ pub fn eval_loop(machine: &mut Machine) -> Result<(), EvalError> {
             }
             Instruction::TablePropertyAssignError => {
                 return Err(EvalError::from(TypeError::CannotAssignProperty {
-                    property: register!(AS).take().unwrap(),
+                    property: register!(AS).clone(),
                     of: std::mem::replace(&mut register!(AD), LuaValue::Nil),
                 }))
             }
@@ -972,7 +972,7 @@ mod test {
         code: [ConstS(StringID(0)), Ret],
         strings: ["hello"],
         post_condition: |machine: Machine| {
-            assert_eq!(register_of!(machine, AS), Some("hello".into()))
+            assert_eq!(register_of!(machine, AS), "hello".into())
         }
     }
 

reggie/src/value/string.rs

@@ -1,7 +1,9 @@
-use std::{fmt, hash::Hash, ptr::NonNull, rc::Rc};
+use std::{alloc::Layout, fmt, hash::Hash, marker::PhantomData, mem::size_of, slice};
 
 use luar_lex::Ident;
 
+const INLINE_BUFFER_SIZE: usize = std::mem::size_of::<*const ()>();
+
 #[repr(packed)]
 pub struct LuaString {
     // since the maximum alingment of LuaValue associated
@@ -11,31 +13,107 @@ pub struct LuaString {
     // potentially "borrow" std::ptr::Unique to make
     // sizeof(Option<LuaString>) == sizeof(LuaString>)
     len: u32,
-    // This shoud've been NonNull<str>, but that is a fat pointer
-    str: NonNull<()>,
+    /// Either an inline string for strings that are shorter than
+    /// INLINE_BUFFER_SIZE bytes, or a pointer to an allocation
+    /// for longer strings. Points straight to the StrBlock.
+    ptr_or_inline_data: SSOStorage,
+    _unused: PhantomData<*const ()>,
+}
+
+/// Small String Optimized storage.
+#[derive(Copy, Clone)]
+union SSOStorage {
+    inline_data: [u8; INLINE_BUFFER_SIZE],
+    heap_allocation: *mut (),
+}
+
+#[repr(C)]
+struct StrBlock {
+    /// Number of outstanding references to this allocation, minus one.
+    /// If refcount is zero, you are responsible for deallocating them.
+    /// Sharing StrBlock between threads is not safe, since refcount
+    /// is not atomic.
+    refcount: usize,
+    _unused: PhantomData<*const ()>,
+    data: str,
 }
 
 impl Default for LuaString {
     fn default() -> Self {
         Self {
             len: 0,
-            str: NonNull::dangling(),
+            _unused: PhantomData,
+            ptr_or_inline_data: SSOStorage {
+                inline_data: [0; INLINE_BUFFER_SIZE],
+            },
         }
     }
 }
 
 impl From<&str> for LuaString {
     fn from(str: &str) -> Self {
-        let allocation: Rc<str> = Rc::from(str);
-        let ptr = Rc::into_raw(allocation);
+        // SAFETY: it is safe to store len of a string in u32, since
+        //         if the value overflows u32, we will panic.
+        let len: u32 = str
+            .len()
+            .try_into()
+            .expect("size of string should not to exceed u32");
+
+        if len == 0 {
+            return Self::default();
+        }
+
+        if len <= INLINE_BUFFER_SIZE as u32 {
+            let mut inline_data = [0; INLINE_BUFFER_SIZE];
+            inline_data[..str.len()].copy_from_slice(str.as_bytes());
+            return Self {
+                len,
+                _unused: PhantomData,
+                ptr_or_inline_data: SSOStorage { inline_data },
+            };
+        }
+
+        // SAFETY: Allocation size is enough to store the StrBlock with the data
+        //         ```
+        //         #[repr(C)]
+        //         struct StrBlock {
+        //             refcount: usize,
+        //             _unused: PhantomData<*const ()>,
+        //             data: str,
+        //         }
+        //         ```
+        //         sizeof(refcount) + sizeof(PhantomData<*const ()>) + sizeof(str) == 8 + 0 + len
+        //         This is the layout of StrBlock, since it is repr(C).
+        //
+        //         Data at allocation is uninitialized, but no matter, we write
+        //         to it immediately afterwards.
+        let block_ptr = unsafe {
+            let allocation_size = len as usize + size_of::<usize>();
+            // SAFETY:
+            //       * `align` is not zero
+            //       * `align` is a power of two
+            //       * `size`, when rounded up to the nearest multiple of `align`,
+            //          cannot not overflow isize (i.e., the rounded value must be
+            //          less than or equal to `isize::MAX`).
+            let allocation =
+                std::alloc::alloc(Layout::from_size_align_unchecked(allocation_size, 1));
+            let slice = slice::from_raw_parts_mut(allocation, len as usize);
+            let block_ptr = slice as *mut [u8] as *mut StrBlock;
+            let str_block = &mut *block_ptr;
+            str_block.refcount = 0;
+            str_block
+                .data
+                .as_bytes_mut()
+                .copy_from_slice(str.as_bytes());
+            block_ptr
+        };
 
         Self {
-            len: str
-                .len()
-                .try_into()
-                .expect("size of string should not to exceed u32"),
-            // This could as well easily be NonNull::new_unchecked
-            str: NonNull::new(ptr as *mut ()).expect("pointer from Rc should never be null"),
+            len,
+            _unused: PhantomData,
+            ptr_or_inline_data: SSOStorage {
+                heap_allocation: block_ptr as *mut (),
+            },
         }
     }
 }
@@ -60,31 +138,64 @@ impl From<Ident> for LuaString {
 
 impl Drop for LuaString {
     fn drop(&mut self) {
-        if self.len != 0 {
-            // SAFETY: This is safe, because we know that the pointer is not dangling
-            //         and that the allocation is still alive
+        if self.len > INLINE_BUFFER_SIZE as u32 {
+            // SAFETY: Everything that is longer than INLINE_BUFFER_SIZE bytes in
+            //         self.ptr_or_inline_data is a valid pointer to a valid allocation,
+            //         allocated by Rc<str>.
+            //         Ptr is saved into [u8; INLINE_BUFFER_SIZE] via usize::to_ne_bytes,
+            //         and is brought back via usize::from_ne_bytes.
+            //         [8; INLINE_BUFFER_SIZE] is guaranteed to have enough space to
+            //         store a pointer. str pointer is brought back exactly as it was
+            //         saved, including it's length in the fat pointer.
+            //         Call to from_utf8_unchecked has no safety implications, since
+            //         str is not accessed in any way.
+            //         Decresing refcount is safe, since LuaString cannot be shared
+            //         between threads.
             unsafe {
-                drop(Rc::from_raw(self.str.as_ptr()));
+                let ptr = self.ptr_or_inline_data.heap_allocation;
+                // Until we have std::ptr::from_raw_parts this is a workaround for
+                // creating fat pointers to ?Sized structs
+                let slice = std::slice::from_raw_parts_mut(ptr, self.len as usize);
+                let block_ptr = slice as *mut [()] as *mut StrBlock;
+                let block = &mut *block_ptr;
+
+                if block.refcount == 0 {
+                    // No need to call Drop, since StrBlock is trivially dropable.
+                    let layout = Layout::from_size_align_unchecked(
+                        size_of::<usize>() + self.len as usize,
+                        1,
+                    );
+                    std::alloc::dealloc(block_ptr as *mut u8, layout)
+                } else {
+                    block.refcount -= 1;
+                }
             }
         }
     }
 }
 
 impl AsRef<str> for LuaString {
     fn as_ref(&self) -> &str {
-        if self.len == 0 {
-            return "";
-        }
-        let ptr = self.str.as_ptr() as *const u8;
-        // SAFETY: This is safe, because we know that the pointer is not dangling
-        //         and that the allocation is still alive
-        //         and that self.str points to a valid UTF-8 string, since
-        //         LuaString::from(&str) is only implemented for valid &str
-        //         self.len is not overflowing, since there is a check in
-        //         impl From<&str> for LuaString
+        // SAFETY: If self.len < INLINE_BUFFER_SIZE, then self.ptr_or_inline_data
+        //         contains a inline string. inline_data is a valid UTF-8 string
+        //         since the only way to construct LuaString is via a valid &str
+        //
+        //         Otherwise self.ptr_or_inline_data pointer to a valid allocation
+        //         of StrBlock.
+        //         StrBlock is valid, since we refcount outstanding references.
         unsafe {
-            let slice = std::slice::from_raw_parts(ptr, self.len as usize);
-            std::str::from_utf8_unchecked(slice)
+            if self.len <= INLINE_BUFFER_SIZE as u32 {
+                let byte_slice = &self.ptr_or_inline_data.inline_data[..self.len as usize];
+                std::str::from_utf8_unchecked(byte_slice)
+            } else {
+                let block_ptr = self.ptr_or_inline_data.heap_allocation;
+                // Until we have std::ptr::from_raw_parts this is a workaround for
+                // creating fat pointers to ?Sized structs
+                let slice = std::slice::from_raw_parts(block_ptr, self.len as usize);
+                let block_ptr = slice as *const [()] as *const StrBlock;
+                let block = &*block_ptr;
+                &block.data
+            }
         }
     }
 }
@@ -130,7 +241,29 @@ impl Hash for LuaString {
 }
 impl Clone for LuaString {
     fn clone(&self) -> Self {
-        Self::from(self.as_ref())
+        if self.len > INLINE_BUFFER_SIZE as u32 {
+            // SAFETY: Everything that is longer than INLINE_BUFFER_SIZE bytes is
+            //         stored inline. Otherwise self.ptr_or_inline_data contains a
+            //         pointer to a valid allocation of StrBlock, allocated by std::alloc.
+            //
+            //         We do not memcopy the allocation, but instead share it.
+            //         StrBlock is refcounted, so we increase the refcount.
+            //         It is safe to increase refcount, since LuaString cannot be shared
+            //         between threads.
+            unsafe {
+                let ptr = self.ptr_or_inline_data.heap_allocation;
+                let slice = std::slice::from_raw_parts_mut(ptr, self.len as usize);
+                let block_ptr = slice as *mut [()] as *mut StrBlock;
+                let block = &mut *block_ptr;
+                block.refcount += 1;
+            }
+        }
+
+        Self {
+            len: self.len,
+            _unused: PhantomData,
+            ptr_or_inline_data: self.ptr_or_inline_data,
+        }
     }
 }
 

value_size.rs

@@ -15,7 +15,7 @@ enum LuaValue {
 #[repr(packed(4))]
 pub struct LuaString {
     len: u32, 
-    block: NonNull<StrBlockHeader>
+    block: [u8; 8],
 }