sendable_rc.rs

#![feature(unsafe_destructor)]

use rooted::{RcSend, Root};
use std::fmt;
use std::rc::Rc;

mod rooted {
    use std::collections::trie_map::{mod, TrieMap};
    use std::fmt;
    use std::kinds::marker;
    use std::mem;
    use std::num::Int;
    use std::rc::Rc;

    pub struct Root;

    pub struct Rooted<'a, T: 'a> {
        // Keys are byte reversed addresses of RcMuts
        // Note: we only need explicit drop because the addresses aren't represented
        // as addresses.
        // Note 2: since address can be retrieved from key, only value is needed.
        roots: TrieMap<uint>,
        #[allow(dead_code)] root: &'a mut Root,
    }

    impl Root {
        pub fn new<'a, T>(&'a mut self) -> Rooted<'a, T> {
            Rooted { roots: TrieMap::new(), root: self }
        }

    }

    impl<'a, T> Rooted<'a, T> where T: Sync {
        pub fn root<'b>(&'b mut self, ptr: Rc<T>) -> RcSend<'a, T> {
            unsafe {
                // This assumes Rc is uint-sized; we will be transmuting it later
                // in the function, so the compiler will enforce this.
                let addr = mem::transmute_copy::<_, uint>(&ptr);
                // Tends to speed up tries for closely clustered values, YMMV
                // (might want to use another data structure like a splay tree with
                // better cache locality as well)
                let trie_key = addr.swap_bytes();
                match self.roots.entry(trie_key) {
                    trie_map::Occupied(mut entry) => {
                        // Already rooted, so we know the refcount is at least one;
                        // go ahead and drop the Rc we were just given.
                        drop(ptr);
                        // Note that even if the above panics, everything is fine,
                        // because Drop for Rooted<T> always drops exactly once per
                        // node in the map (it doesn't depend on the keys).
                        let count = *entry.get();
                        entry.set(count + 1);
                    },
                    trie_map::Vacant(entry) => {
                        // We unfortunately cannot 100% guarantee that this will neither
                        // leak nor cause memory unsafety (if the set panics, which it
                        // shouldn't).  So we opt for the lesser evil and accept the
                        // possibility of a leak.  If we had slightly stronger guarantees
                        // (that entry.set()) always sets the entry atomically) we could
                        // guarantee this though.
                        //
                        // We transmute instead of mem::forget just to make the compiler
                        // catch a mistake in size.  Do not remove this line, without it
                        // we will double drop.
                        let _ = mem::transmute::<_, uint>(ptr);
                        // Refcounts start at 0, not 1.  We never want to hold onto T if
                        // the entry is zero.
                        entry.set(0);
                    }
                }

                // Note that this transmute is safe because we know that we
                // will not *use* it until the reference is safely a key in
                // the TrieMap and it is wrapped in a RcSend.
                // (This also relies on Rc not to move the pointer around... I think
                // that is safe to rely upon though).
                let ptr_ = &*(addr as *const ());
                RcSend {
                    inner_: ptr_,
                    marker: marker::NoCopy,
                }
            }
        }

        // Note that the return type of this is there in case you try to send it to the
        // wrong Rootable.  RcSend does not itself actually implement Drop, so you cannot
        // cause memory unsafety by not handling the result, and the Rootables will all
        // eventually free their references.  Still, it's best practice to handle it since
        // it probably indicates an error in program logic if you tried to free in the wrong
        // thread (you might do something like cycle through available Rootables to find the
        // appropriate one, for instance).
        pub fn unroot(&mut self, ptr: RcSend<'a, T>) -> Result<(), RcSend<'a, T>> {
            unsafe {
                let addr = ptr.inner_ as *const _ as uint;
                let trie_key = addr.swap_bytes();
                match self.roots.entry(trie_key) {
                    trie_map::Occupied(mut entry) => {
                        // Already rooted, so we know the refcount is at least one;
                        // go ahead and drop the Rc we were just given.
                        // Note that even if the above panics, everything is fine,
                        // because Drop for Rooted<T> always drops exactly once per
                        // node in the map (it doesn't depend on the keys).
                        let count = *entry.get();
                        if count == 0 {
                            // We want to drop the entry for real.
                            //
                            // We can leak here if the take() fails (but it shouldn't).
                            let _ = entry.take();
                            // Now we can safely transmute the Rc and will neither
                            // leak nor cause memory unsafety.
                            let rc = mem::transmute::<_, Rc<T>>(ptr.inner_);
                            // Explicit drop for clarity.
                            drop(rc);
                        } else {
                            // No need to drop.  If this fails, we might leak though
                            // (but it shouldn't).
                            entry.set(count - 1);
                        }
                        Ok(())
                    },
                    trie_map::Vacant(_) => Err(ptr),
                }
            }
        }
    }

    #[unsafe_destructor]
    impl<'a, T> Drop for Rooted<'a, T> where T: Sync {
        fn drop<'b>(&'b mut self) {
            for trie_key in self.roots.keys() {
                // Get the original address back
                let addr = trie_key.swap_bytes();
                unsafe {
                    // Transmute to Rc.
                    let rc = mem::transmute::<_, Rc<T>>(addr);
                    // Explicit drop for clarity.  If this fails, we may leak.
                    drop(rc);
                }
            }
            // We rely on drop not being called more than once for safety... if this is suspected to be
            // a problem, we can explicitly call clear() here, but I suspect it is overkill.
            // self.roots.clear();
        }
    }

    // Do not implement Clone on this, or anything like Clone.
    // Mutating or copying it is not legal.
    #[must_use]
    pub struct RcSend<'a, T: 'a> where T: Sync {
        inner_: &'a (),
        // Shouldn't have to put marker::NoSend, but it should not be allowed to outlast 'a...
        marker: marker::NoCopy,
    }

    impl<'a, T: 'a> fmt::Show for RcSend<'a, T> where T: Sync {
        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
            write!(f, "RcSend")
        }
    }

    impl<'a, T> Deref<T> for RcSend<'a, T> where T: Sync {
        fn deref<'b>(&'b self) -> &'b T {
            unsafe {
                // Safe because we should always construct this from a Rc.
                // Would be great if RcBox / RcInner etc. were exposed
                // so we could do this less hackily.
                let rc = &self.inner_;
                // We don't want to transmute to Rc directly in case the Rc
                // deref fails (I know it shouldn't fail, but we're being
                // extra careful).  Hopefully this gets optimized out.
                let rc = rc as *const _ as *const Rc<T>;
                // This lifetime is known to be longer than this block,
                // because the Rc itself has lifetime 'a (thanks to the
                // Rooted in which it lives).
                let reference = mem::copy_lifetime(self, (*rc).deref());
                // The reference is only allowed to live as long as 'b, though.
                reference
            }
        }
    }
}

// Again, something that can go away if my RFC is accepted.
pub fn local_channel<'a, T>() -> (Sender<RcSend<'a, T>>, Receiver<RcSend<'a, T>>)
where T: Sync + 'static {
    unsafe {
        return ::std::mem::transmute(channel::<RcSend<'static, T>>());
    }
}

// Same as above
pub fn local_send<'a, 'b, T>(tx: &'b Sender<RcSend<'a, T>>, msg: RcSend<'a, T>) where T: Sync + 'static {
    unsafe {
        let tx = ::std::mem::transmute::<_, &'b Sender<RcSend<'static, T>>>(tx);
        let msg = ::std::mem::transmute::<_, RcSend<'static, T>>(msg);
        tx.send(msg);
    }
}

// Same as above
pub fn local_recv<'a, 'b, T>(rx: &'b Receiver<RcSend<'a, T>>) -> RcSend<'a, T> where T: Sync + 'static {
    unsafe {
        let rx = ::std::mem::transmute::<_, &'b Receiver<RcSend<'static, T>>>(rx);
        ::std::mem::transmute(rx.recv())
    }
}

pub fn spawn_test<'a, T>(//root: RcSend<'a, T>,
                         rx: Receiver<RcSend<'a, T>>,
                         tx_: Sender<RcSend<'a, T>>) where T: 'static + Sync + ::std::fmt::Show {
    // This relies on something like rayon taking a closure bounded by Sync + 'a or (if my RFC goes
    // through) Send + 'a.  For this quick example, we use unsafe code to transmute the 'a to 'static.
    unsafe {
        //let root = ::std::mem::transmute::<_, RcSend<'static, T>>(root);
        let rx = ::std::mem::transmute::<_, Receiver<RcSend<'static, T>>>(rx);
        let tx_ = ::std::mem::transmute::<_, Sender<RcSend<'static, T>>>(tx_);
        spawn(proc() {
            let foo = rx.recv();
            println!("{}", *foo);
            tx_.send(foo);
        })
    }
}

struct Example { marker: ::std::kinds::marker::NoCopy }

impl fmt::Show for Example {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Example")
    }
}

impl Drop for Example {
    fn drop(&mut self) {
        println!("Dropped");
    }
}

fn main() {
    let mut root = Root;
    let mut rooted = root.new();
    {
        let foo = Rc::new(Example { marker: ::std::kinds::marker::NoCopy });
        let foo_ = rooted.root(foo.clone());
        let (tx, rx) = local_channel();
        let (tx_, rx_) = local_channel();
        spawn_test(rx, tx_);
        local_send(&tx, foo_);
        drop(foo);
        println!("Should not be dropped here.");
        rooted.unroot(local_recv(&rx_)).unwrap();
        println!("Should be dropped here.");
    }
}