use std::collections::HashMap;
use std::marker::PhantomData;
use std::mem;
use std::sync::{Arc, Weak};
use std::sync::atomic::{AtomicUsize, Ordering};

struct WeakCell<T>(AtomicUsize, PhantomData<Weak<T>>);

impl<T> WeakCell<T> {
    /// Creates a new `WeakCell`.
    pub fn new(t: Weak<T>) -> WeakCell<T> {
        WeakCell(AtomicUsize::new(unsafe { mem::transmute(t) }), PhantomData)
    }

    fn take(&self) -> Weak<T> {
        loop {
            match self.0.swap(0, Ordering::Acquire) {
                0 => {},
                n => return unsafe { mem::transmute(n) },
            }
        }
    }

    fn put(&self, t: Weak<T>) {
        debug_assert_eq!(self.0.load(Ordering::SeqCst), 0);
        self.0
            .store(unsafe { mem::transmute(t) }, Ordering::Release);
    }

    /// Stores a new value in the `WeakCell`, returning the previous
    /// value.
    pub fn set(&self, t: Weak<T>) -> Weak<T> {
        let old = self.take();
        self.put(t);
        old
    }

    /// Returns a copy of the value stored by the `WeakCell`.
    pub fn get(&self) -> Weak<T> {
        let t = self.take();
        // NB: correctness here depends on Weak's clone impl not panicking
        let out = t.clone();
        self.put(t);
        out
    }
}

impl<T> Default for WeakCell<T> {
    fn default() -> Self { WeakCell(AtomicUsize::new(0), PhantomData) }
}

impl<T> Drop for WeakCell<T> {
    fn drop(&mut self) { self.take(); }
}

#[derive(Default)]
struct A {
    map: HashMap<u32, Vec<B>>,
}

struct B {
    weak: WeakCell<A>,
}

impl A {
    pub fn new(map: HashMap<u32, Vec<B>>) -> Arc<A> {
        let mut a = Arc::new(A { map });
        let copy = a.clone();
        for (_, bs) in &mut Arc::get_mut(&mut a).unwrap().map {
            for b in bs {
                b.weak.set(Arc::downgrade(&copy));
            }
        }
        a
    }
}

impl B {
    pub fn new(a: &Arc<A>) -> B { B { weak: WeakCell::new(Arc::downgrade(a)), } }
}

fn main() {
    let dummy = Arc::new(A::default());

    let (b1, b2, b3) = (B::new(&dummy), B::new(&dummy), B::new(&dummy));

    let mut map = HashMap::new();
    map.insert(5, vec![b1, b2]);
    map.insert(10, vec![b3]);

    let _a = A::new(map);

    //  Do something!
}