#![allow(dead_code)]

use std::cell::UnsafeCell;
use std::collections::HashMap;
use std::iter;
use std::marker::PhantomData;
use std::mem;
use std::ops::Deref;
use std::ops::Range;
use std::sync::Mutex;
use std::thread;
use lazy_static::lazy_static;

// Jump to the examples at the end!

// The metadata we store is just a colour for each memory address.
struct MetaData(Mutex<HashMap<Address, Colour>>);

lazy_static! { static ref METADATA: MetaData = MetaData(Mutex::new(HashMap::new())); }

// Memory addresses are just a word
type Address = usize;

// The different states of memory
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum Colour {
    Unreachable,   // Memory that is allocated but can't be reached from safe code
    Unique,        // Memory that is reachable from safe code as a &mut T reference
    Shared(usize), // Memory that is reachable from safe code as a &T reference
}

// Types that know how to iterate over their addresses
trait Addressable {
    type Addresses: Iterator<Item = Address>;
    fn addresses(&self) -> Self::Addresses;
}

impl Addressable for u8 {
    type Addresses = iter::Once<Address>;
    fn addresses(&self) -> Self::Addresses {
        iter::once(self as *const u8 as Address)
    }
}

impl Addressable for u16 {
    type Addresses = Range<Address>;
    fn addresses(&self) -> Self::Addresses {
        let address = self as *const u16 as Address;
        (address..address+2)
    }
}

impl Addressable for u32 {
    type Addresses = Range<Address>;
    fn addresses(&self) -> Self::Addresses {
        let address = self as *const u32 as Address;
        (address..address+4)
    }
}

impl Addressable for u64 {
    type Addresses = Range<Address>;
    fn addresses(&self) -> Self::Addresses {
        let address = self as *const u64 as Address;
        (address..address+8)
    }
}

impl<T, U> Addressable for (T, U) where T: Addressable, U: Addressable {
    type Addresses = iter::Chain<T::Addresses, U::Addresses>;
    fn addresses(&self) -> Self::Addresses {
        self.0.addresses().chain(self.1.addresses())
    }
}

impl<T> Addressable for UnsafeCell<T> where T: Addressable {
    type Addresses = iter::Empty<Address>;
    fn addresses(&self) -> Self::Addresses {
        iter::empty()
    }
}

// How we model &T
struct Ref<'a, T: 'a + Addressable>(&'a T);

impl<'a, T: 'a + Addressable> Drop for Ref<'a, T> {
    fn drop(&mut self) {
        if thread::panicking() { return; }
        let mut metadata = METADATA.0.lock().unwrap();
        for address in self.0.addresses() {
            let n = match metadata.get(&address) {
                Some(Colour::Shared(n)) => *n,
                colour => panic!("Expected Shared, found {:?} at address 0x{:x}", colour, address),
            };
            assert!(n > 0);
            let colour = if n == 1 {
                Colour::Unreachable
            } else {
                Colour::Shared(n - 1)
            };
            metadata.insert(address, colour);
            println!("Colouring 0x{:x} as {:?}", address, colour);
        }
    }
}

impl<'a, T: 'a + Addressable> Clone for Ref<'a, T> {
    fn clone(&self) -> Self {
        let mut metadata = METADATA.0.lock().unwrap();
        for address in self.0.addresses() {
            let n = match metadata.get(&address) {
                Some(Colour::Shared(n)) => *n,
                colour => panic!("Expected Shared, found {:?} at address 0x{:x}", colour, address),
            };
            metadata.insert(address, Colour::Shared(n + 1));
            println!("Colouring 0x{:x} as Shared({})", address, n + 1);
        }
        Ref(self.0)
    }
}

impl<'a, T: 'a + Addressable> Ref<'a, T> {
    fn get(&self) -> T where T: Clone {
        self.0.clone()
    }

    fn as_ptr(self) -> *const T {
        self.0 as *const T
    }

    unsafe fn from_ptr(ptr: *const T) -> Ref<'a, T> {
        let mut metadata = METADATA.0.lock().unwrap();
        let reference = &*ptr;
        for address in reference.addresses() {
            let n = match metadata.get(&address) {
                Some(Colour::Shared(n)) => *n,
                Some(Colour::Unreachable) => 0,
                colour => panic!("UB: Expected Shared or Unreachable, found {:?} at address 0x{:x}", colour, address),
            };
            metadata.insert(address, Colour::Shared(n + 1));
            println!("Colouring 0x{:x} as Shared({})", address, n + 1);
        }
        Ref(reference)
    }

}

// How we model &mut T
struct RefMut<'a, T: 'a + Addressable>(&'a mut T);

impl<'a, T: 'a + Addressable> RefMut<'a, T> {
    fn new(x: &'a mut T) -> Self {
        let mut metadata = METADATA.0.lock().unwrap();
        for address in x.addresses() {
            let colour = metadata.get(&address).cloned().unwrap_or(Colour::Unreachable);
            assert!(colour == Colour::Unreachable,
                "Expected Unreachable, found {:?} at address 0x{:x}", colour, address);
            metadata.insert(address, Colour::Unique);
            println!("Colouring 0x{:x} as Unique", address);
        }
        RefMut(x)
    }

    fn as_ref<'b>(&'b mut self) -> Ref<'b, T> {
        let mut metadata = METADATA.0.lock().unwrap();
        for address in self.0.addresses() {
            let colour = metadata.get(&address).cloned().unwrap_or(Colour::Unreachable);
            assert!((colour == Colour::Unreachable) || (colour == Colour::Unique),
                "Expected Unreachable or Unique, found {:?} at address 0x{:x}", colour, address);
            metadata.insert(address, Colour::Shared(1));
            println!("Colouring 0x{:x} as Shared(1)", address);
        }
        Ref(self.0)
    }

    unsafe fn from_ptr(ptr: *mut T) -> RefMut<'a, T> {
        let mut metadata = METADATA.0.lock().unwrap();
        let reference = &mut *ptr;
        for address in reference.addresses() {
            let colour = metadata.get(&address).cloned().unwrap_or(Colour::Unreachable);
            assert!(colour == Colour::Unreachable,
                "UB: Expected Unreachable, found {:?} at address 0x{:x}", colour, address);
            metadata.insert(address, Colour::Unique);
            println!("Colouring 0x{:x} as Unique", address);
        }
        RefMut(reference)
    }

    fn set(&mut self, x: T) {
        let mut metadata = METADATA.0.lock().unwrap();
        for address in self.0.addresses() {
            let colour = metadata.get(&address).cloned().unwrap_or(Colour::Unreachable);
            assert!((colour == Colour::Unreachable) || (colour == Colour::Unique),
                "Expected Unreachable or Unique, found {:?} at address 0x{:x}", colour, address);
            metadata.insert(address, Colour::Unique);
            println!("Colouring 0x{:x} as Unique", address);
        }
        *self.0 = x;
    }
}

macro_rules! let_ref_mut {
    ($x:ident = $e:expr) => {
        let ref mut tmp = $e;
        let mut $x = RefMut::new(tmp);
    }
}

/*
// Another example, from Ralf Jung

fn aliasing_mut_and_shr() {
    fn inner(rc: &RefCell<u32>, aliasing: &mut u32) {
        *aliasing += 4;
        let _escape_to_raw : *const RefCell<u32> = rc;
        *aliasing += 4;
        let _shr = &*rc;
        *aliasing += 4;
        // also turning this into a frozen ref now must work
        let aliasing : &u32 = aliasing;
        let _val = *aliasing;
        let _escape_to_raw : *const RefCell<u32> = rc;
        let _val = *aliasing;
        let _shr = &*rc; // this must NOT unfreeze
        let _val = *aliasing;
    }

    let refcell = RefCell::new(23);
    let rc: &mut RefCell<u32> = &mut refcell;
    let rc: &RefCell<u32> = rc;
    let mut bmut : RefMut<u32> = rc.borrow_mut();
    inner(&rc, &mut *bmut);
    drop(bmut);
    assert_eq!(*rc.borrow(), 23+12);
}

fn main() {
    aliasing_mut_and_shr()
}
*/

fn bad(p: Ref<u8>) {
    let q = p.as_ptr() as *mut u8;
    let mut r = unsafe { RefMut::from_ptr(q) };
    r.set(5);
}

fn main() {
    let_ref_mut!(x = 37);
    let y = x.as_ref();
    println!("*y = {}", y.get());
    bad(y);
    let z = x.as_ref();
    println!("*z = {}", z.get());
    bad(z.clone());
    println!("*z = {}", z.get());
}