Search code examples
rustlifetimeborrow-checkerunsafe

How to create a self referential struct?


Context: I'm working on a very fast algorithm to solve the Advent of Code 2023 Day 8 part 2 challenge without the least-common-multiple approach (which is meant to be impossible, hence the optimizations).

One of my solutions works with HashMaps internally. As keys and values, they contain strings or references to strings that are present in the input data (think of it as just a large string).

So basically, I'm taking apart the input data and spread it across a bunch of derived data structures.

I can now either copy these bits using .to_owned(), or just use references to sections within the original input, which brings runtime down from 2m12s to 1m38s.

My trouble is that I cannot use the fast reference approach without tying the lifetime of my data-structures to the lifetime of the reference to the input data, which means that the caller has to deal with that.

What I would prefer is to consume the input data, so the caller can forget about it, and still work with references internally. But that just doesn't work out.

Let some dummy code highlight the issue:

use std::fs;

/**
 * Works fine, but needs to duplicate data, which makes it a bit slower
 */
pub fn read_file_and_analyze_little_slow(file_name: &str) -> (String, Vec<String>) {
    let content = fs::read_to_string(file_name).expect("Could not read the file you told me to analyze");
    let lines: Vec<String> = content.split("\n").map(|x| { x.to_owned() }).collect();

    (content, lines)
}

/**
 * I wish I could do this, but the read content of the file is dropped
 * at the end of the file (or so the borrow checker thinks), invalidating
 * the references
 */
pub fn read_file_and_analyze_fast(file_name: &str) -> (String, Vec<&str>) {
    let content = fs::read_to_string(file_name).expect("Could not read the file you told me to analyze");
    let lines: Vec<&str> = content.split("\n").collect();

    // Would need to make content live on somehow, so the references remain valid
    (content, lines) // Error: cannot move out of `content` because it is borrowed
}

/**
 * This works super well, but we've just transferred the issue to the calling function,
 * because now our result is tied to the lifetime of input string, meaning the parent
 * function can't pass it as _its_ result, unless it also works with lifetimes
 */
pub fn read_file_and_analyze_unwieldy<'a>(already_read_file: &'a str) -> Vec<&'a str> {
    already_read_file.split("\n").collect()
}

There are of course more variations, like have the input passed in as String, which would align better with the question title but afaik not make an actual difference to the problem.

I'm looking for solutions to make read_file_and_analyze_fast's API work. I suspect it's probably possible but may require unsafe, which I haven't worked with before.

This is a very interesting problem, so some general guidance to deal with issues like this one would also be appreciated


Solution

  • This is one of the cases where String::leak might be appropriate. This consumes the String (without releasing the heap allocation it owns) and gives you back a reference to a slice of the entire string, but with 'static lifetime.

    pub fn read_file_and_analyze_fast(file_name: &str) -> Vec<&'static str> {
        let content = std::fs::read_to_string(file_name)
            .expect("Could not read the file you told me to analyze")
            .leak();
    
        let lines: Vec<&str> = content.split("\n").collect();
    
        lines
    }
    

    Note that I changed the return type to Vec<&'static str> to make it clear that the string slices contained in the Vec have 'static lifetime. If you don't change this, they will be limited to the same lifetime as the input slice reference file_name (which could also be 'static, but that doesn't really matter here).

    Beware that a leaked string can never be safely reclaimed, so it would be a bad idea to use this with any string that you don't expect to be used the entire time the program is running. Especially be careful about using it in loops!


    Another approach, since the input file is going to be the same every time, is to use include_str!, which is a macro that expands to a &'static str obtained by reading the contents of the provided file name at compile time. (This is what I did when I was working on Advent of Code.)


    If you are hell-bent on a self-referential data structure, this is safe with strings since the string data is stored in a separate allocation that does not move when the String itself does. To ensure this is sound, the String needs to live at least as long as the data structure that references its contents, and the String cannot be changed for that duration.

    To make this work in anything resembling an ergonomic manner, where we don't need a separate type for every different kind of container, we can create a trait with a lifetime-generic associated type ("lifetime GAT").

    pub trait StringSliceContainer {
        type Container<'a>;
    }
    

    Now let's implement for Vec<&str>.

    impl StringSliceContainer for Vec<&str> {
        type Container<'a> = Vec<&'a str>;
    }
    

    Simple enough so far. Note that the type StringSliceContainer is implemented on doesn't have to match the type of Container. We could have created our own struct VecStr; and implemented on that instead, still with type Container<'a> = Vec<&'a str>;.

    There's a few things we need to keep in mind when we create our self-owning type. The idea is that it will own both a String and a T::Container for some T implementing StringSliceContainer. There's two crucial things that will inform how this type is defined, and one is a consequence of the other.

    First, the only valid lifetime we can use, without introducing a named lifetime, is 'static. This means we have to define our type as something like this:

    pub struct SelfOwnedStringSlices<T: StringSliceContainer> {
        _string: String,
        container: T::Container<'static>,
    }
    

    However, this seems wrong... and it kind of is. T::Container<'static> isn't right, is it? Well, it's the best we can do here. The important thing is that we never actually use this field as a T::Container<'static>! We must always shorten its lifetime parameter first, to something that is no larger than the lifetime of our owned string.

    Second, because of this, we cannot rely on the compiler-generated drop glue, because it's going to drop a T::Container<'static>, which won't be valid. Therefore, we must tweak the definition of this struct to wrap the container in ManuallyDrop. This ensures that the compiler-generated drop glue won't actually drop this field. Instead, we'll manage dropping it ourselves after we've adjusted the lifetime appropriately.

    So now we have this:

    pub struct SelfOwnedStringSlices<T: StringSliceContainer> {
        _string: String,
        container: ManuallyDrop<T::Container<'static>>,
    }
    

    Ok. Now how do we make one of these things? Well, all we need is a string, and some way to create a T::Container<'static> from a &'a str. Except we don't actually want the consumer of this type to produce a T::Container<'static> because then they can't reference anything in the string! So they need to give us a way to produce a T::Container<'a> from a &'a str. This is easy enough to express with a closure.

    impl<T: StringSliceContainer> SelfOwnedStringSlices<T> {
        pub fn new<F>(string: String, container_factory: F) -> Self
        where
            F: for<'a> FnOnce(&'a str) -> T::Container<'a>,
        {
            todo!()
        }
    }
    

    The body of this function should call the closure, and then create a new Self using the returned container and the string. But we get a T::Container<'a>, not a T::Container<'static>! Well, thankfully, this is one of the cases where transmute can help us. The layout of a type does not depend on any lifetime parameters, so this kind of transmutation is safe. However, as mentioned before, we must be careful to never use the container with this lifetime. We always need to transmute the lifetime back to one that cannot exceed the lifetime of the string that we own.

    impl<T: StringSliceContainer> SelfOwnedStringSlices<T> {
        pub fn new<F>(string: String, container_factory: F) -> Self
        where
            F: for<'a> FnOnce(&'a str) -> T::Container<'a>,
        {
            let container = container_factory(&string);
    
            // SAFETY: We are only changing the lifetime parameter to be 'static,
            // and we will never actually use this container with the 'static
            // lifetime.
            let container =
                unsafe { std::mem::transmute::<T::Container<'_>, T::Container<'static>>(container) };
    
            Self {
                container: ManuallyDrop::new(container),
                _string: string,
            }
        }
    }
    

    How do we actually use the container, now? Thankfully, that's basically the same transmutation in the other direction (replacing 'static with the lifetime of the &self reference), only on a reference to the container.

    impl<T: StringSliceContainer> SelfOwnedStringSlices<T> {
        pub fn get<'a>(&'a self) -> &'a T::Container<'a> {
            // SAFETY: We are only changing the lifetime parameter, and we are
            // shortening it to the lifetime of the self reference, which ensures
            // that the caller cannot retain any (non-static) reference stored in
            // the container for longer than self exists.
            unsafe {
                std::mem::transmute::<&'a T::Container<'static>, &'a T::Container<'a>>(&self.container)
            }
        }
    }
    

    Not too bad.

    Note that we unfortunately can't implement Deref because when defining Deref::Target we'd need to use T::Container<'something> but in that context we don't have a reference from which we can derive a lifetime.

    The final piece of the puzzle is how we handle dropping the container. Interestingly, we need to desugar the lifetime in Drop::drop so we have a named lifetime we can transmute the container's lifetime parameter to. Once we have that named lifetime, the transmutation here is about the same thing happening in the get method above, just with an exclusive reference and retaining the ManuallyDrop layer.

    impl<T: StringSliceContainer> Drop for SelfOwnedStringSlices<T> {
        fn drop<'a>(&'a mut self) {
            // SAFETY: We change the lifetime of the container so it's safe to drop,
            // then we drop it.
            unsafe {
                let container = std::mem::transmute::<
                    &'a mut ManuallyDrop<T::Container<'static>>,
                    &'a mut ManuallyDrop<T::Container<'a>>,
                >(&mut self.container);
    
                ManuallyDrop::drop(container);
            }
        }
    }
    

    And we're done!

    Using this type is now this simple:

    let string = "this is a test".to_owned();
    
    let vec = SelfOwnedStringSlices::<Vec<&str>>::new(
        string,
        |s| s.split_ascii_whitespace().collect()
    );
    
    println!("{:?}", vec.get());
    

    (Playground)

    Note there's quite a lot of unsafe gymnastics being performed here. To my knowledge, everything done here is sound. However, the solutions using String::leak and include_str! involve you writing no unsafe code, so they are trivially provably sound. They're also a lot shorter!

    Put another way, you have one solution that is clever and fancy (which are both synonyms for "complicated" in software engineering) and two alternative solutions that are simpler. In this context (Advent of Code) all three solutions should have the same performance characteristics.

    In my mind, the simpler solutions would win out. However, programming challenges like Advent of Code are often the perfect place to learn and play around with new things, so using the more complicated approach to learn about unsafe Rust and how to expose a safe interface on top of unsafe code would also be a perfectly valid decision.