rust heap-memory dynamic-memory-allocation

How can you allocate a raw mutable pointer in stable Rust?

I was trying to build a naive implementation of a custom String-like struct with small string optimization. Now that unions are allowed in stable Rust, I came up with the following code:

struct Large {
    capacity: usize,
    buffer: *mut u8,
}

struct Small([u8; 16]);

union Container {
    large: Large,
    small: Small,
}

struct MyString {
    len: usize,
    container: Container,
}

I can't seem to find a way how to allocate that *mut u8. Is it possible to do in stable Rust? It looks like using alloc::heap would work, but it is only available in nightly.

Solution

As of Rust 1.28, std::alloc::alloc is stable.

Here is an example which shows in general how it can be used.

use std::{
    alloc::{self, Layout},
    cmp, mem, ptr, slice, str,
};

// This really should **not** be copied
#[derive(Copy, Clone)]
struct Large {
    capacity: usize,
    buffer: *mut u8,
}

// This really should **not** be copied
#[derive(Copy, Clone, Default)]
struct Small([u8; 16]);

union Container {
    large: Large,
    small: Small,
}

struct MyString {
    len: usize,
    container: Container,
}

impl MyString {
    fn new() -> Self {
        MyString {
            len: 0,
            container: Container {
                small: Small::default(),
            },
        }
    }

    fn as_buf(&self) -> &[u8] {
        unsafe {
            if self.len <= 16 {
                &self.container.small.0[..self.len]
            } else {
                slice::from_raw_parts(self.container.large.buffer, self.len)
            }
        }
    }

    pub fn as_str(&self) -> &str {
        unsafe { str::from_utf8_unchecked(self.as_buf()) }
    }

    // Not actually UTF-8 safe!
    fn push(&mut self, c: u8) {
        unsafe {
            use cmp::Ordering::*;

            match self.len.cmp(&16) {
                Less => {
                    self.container.small.0[self.len] = c;
                }
                Equal => {
                    let capacity = 17;
                    let layout = Layout::from_size_align(capacity, mem::align_of::<u8>())
                        .expect("Bad layout");

                    let buffer = alloc::alloc(layout);

                    {
                        let buf = self.as_buf();
                        ptr::copy_nonoverlapping(buf.as_ptr(), buffer, buf.len());
                    }

                    self.container.large = Large { capacity, buffer };

                    *self.container.large.buffer.offset(self.len as isize) = c;
                }
                Greater => {
                    let Large {
                        mut capacity,
                        buffer,
                    } = self.container.large;
                    capacity += 1;

                    let layout = Layout::from_size_align(capacity, mem::align_of::<u8>())
                        .expect("Bad layout");

                    let buffer = alloc::realloc(buffer, layout, capacity);

                    self.container.large = Large { capacity, buffer };

                    *self.container.large.buffer.offset(self.len as isize) = c;
                }
            }

            self.len += 1;
        }
    }
}

impl Drop for MyString {
    fn drop(&mut self) {
        unsafe {
            if self.len > 16 {
                let Large { capacity, buffer } = self.container.large;
                let layout =
                    Layout::from_size_align(capacity, mem::align_of::<u8>()).expect("Bad layout");
                alloc::dealloc(buffer, layout);
            }
        }
    }
}

fn main() {
    let mut s = MyString::new();

    for _ in 0..32 {
        s.push(b'a');
        println!("{}", s.as_str());
    }
}

I believe this code to be correct with respect to allocations, but not for anything else. Like all unsafe code, verify it yourself. It's also completely inefficient as it reallocates for every additional character.

If you'd like to allocate a collection of u8 instead of a single u8, you can create a Vec and then convert it into the constituent pieces, such as by calling as_mut_ptr:

use std::mem;

fn main() {
    let mut foo = vec![0; 1024]; // or Vec::<u8>::with_capacity(1024);

    let ptr = foo.as_mut_ptr();
    let cap = foo.capacity();
    let len = foo.len();
    mem::forget(foo); // Avoid calling the destructor!

    let foo_again = unsafe { Vec::from_raw_parts(ptr, len, cap) }; // Rebuild it to drop it
    // Do *NOT* use `ptr` / `cap` / `len` anymore
}

Re allocating is a bit of a pain though; you'd have to convert back to a Vec and do the whole dance forwards and backwards

~~That being said, your Large struct seems to be missing a length, which would be distinct from capacity. You could just use a Vec instead of writing it out.~~ I see now it's up a bit in the hierarchy.

I wonder if having a full String wouldn't be a lot easier, even if it were a bit less efficient in that the length is double-counted...

union Container {
    large: String,
    small: Small,
}