#![allow(dead_code)]

use std::cell::RefCell;
use std::cell::RefMut;
use std::collections::HashMap;
use std::marker::PhantomData;
use std::mem;
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 {
    Undefined,   // Memory that has not been allocated
    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,      // Memory that is reachable from safe code as a &T reference
}


// Types that can change colour.
// For the whole language, every type would implement this,\
// for now we just do tuples of integers.
trait Colourable {
    fn colour(address: Address, from: Colour, to: Colour);
}

impl Colourable for u8 {
    fn colour(address: Address, from: Colour, to: Colour) {
        let old = METADATA.0.lock().unwrap().insert(address, to).unwrap_or(Colour::Undefined);
        if !thread::panicking() && (from != old) {
            panic!("UB: Found colour {:?}, expecting {:?} at address 0x{:x}.", old, from, address);
        }
    }
}

impl Colourable for u16 {
    fn colour(address: Address, from: Colour, to: Colour) {
        u8::colour(address, from, to);
        u8::colour(address + 1, from, to);
    }
}

impl Colourable for u32 {
    fn colour(address: Address, from: Colour, to: Colour) {
        u16::colour(address, from, to);
        u16::colour(address + 2, from, to);
    }
}

impl Colourable for u64 {
        fn colour(address: Address, from: Colour, to: Colour) {
        u32::colour(address, from, to);
        u32::colour(address + 4, from, to);
    }
}

impl Colourable for usize {
    fn colour(address: Address, from: Colour, to: Colour) {
        u64::colour(address, from, to);
    }
}

impl<T, U> Colourable for (T, U) where T: Colourable, U: Colourable {
    fn colour(address: Address, from: Colour, to: Colour) {
        T::colour(address, from, to);
        U::colour(address + mem::size_of::<T>(), from, to);
    }
}

impl<T> Colourable for RefCell<T> where T: Colourable {
    fn colour(address: Address, from: Colour, to: Colour) {
        // This assumes that the address of the RefCell is not
        // the same as the address of the pointer in the RefCell.
        u8::colour(address, from, to);
        // Really we should be able to this without casting,
        // but the RefCell's fields are private.
        let refcell = unsafe { &*(address as *const RefCell<T>) };
        let address = refcell.as_ptr() as Address;
        if from == Colour::Undefined {
            T::colour(address, Colour::Undefined, Colour::Unreachable);
        } else if from == Colour::Undefined {
            T::colour(address, Colour::Unreachable, Colour::Undefined);
        }
    }
}

// When we recolour memory, we leave a Recolour marker on the stack.
// When its drop code runs, it undoes the recolouring.
struct Recolour<T: Colourable> {
    old: Colour,
    new: Colour,
    address: Address,
    marker: PhantomData<T>,
}

impl<T: Colourable> Drop for Recolour<T> {
    fn drop(&mut self) {
        T::colour(self.address, self.new, self.old);
    }
}

// A trick to make sure the Recolour marker lasts as long as the borrow does
impl<T: Colourable> Recolour<T> {
    fn set_address<'a>(&'a mut self, reference: &'a T) -> &'a T {
        self.address = reference as *const T as Address;
        T::colour(self.address, self.old, self.new);
        reference
    }
    fn set_address_mut<'a>(&'a mut self, reference: &'a mut T) -> &'a mut T {
        self.address = reference as *mut T as Address;
        T::colour(self.address, self.old, self.new);
        reference
    }
}

// A macro for tracking borrows and reborrows.
// This is my attempt to code the borrow model directly in Rust.
macro_rules! track {
    // Borrow a T as a &mut T
    (let $x:ident : &mut $t:ty = &mut $e:expr) => {
        let mut tmp : $t = $e;
        let mut recolour = Recolour {
            old: Colour::Undefined,
            new: Colour::Unique,
            address: 0,
            marker: PhantomData,
        };
        let $x : &mut $t = recolour.set_address_mut(&mut tmp);
    };
    // Reborrow a &mut T as a &T
    (let $x:ident : & $t:ty = $e:expr) => {
        let tmp: &mut $t  = $e;
        let mut recolour = Recolour {
            old: Colour::Unique,
            new: Colour::Shared,
            address: 0,
            marker: PhantomData,
        };
        let $x : & $t = recolour.set_address(tmp);
    };
    // Cast a *mut T to a &mut T
    (let $x:ident : &mut $t:ty = unsafe { &mut * ($e:expr) }) => {
        let tmp: *mut $t = $e;
        let mut recolour = Recolour {
            old: Colour::Unreachable,
            new: Colour::Unique,
            address: 0,
            marker: PhantomData,
        };
        let $x : &mut $t = recolour.set_address_mut(unsafe { &mut *tmp });
    };
    // Cast a *const T to a &T
    (let $x:ident : & $t:ty = unsafe { & * ($e:expr) }) => {
        let tmp: *const $t = $e;
        let mut recolour = Recolour {
            old: Colour::Shared,
            new: Colour::Shared,
            address: 0,
            marker: PhantomData,
        };
        let $x : & $t = recolour.set_address(unsafe { &*tmp });
    };
    // Cast a &mut T to a *mut T
    (let $x:ident : *mut $t:ty = $e:expr) => {
        let tmp: &mut $t = $e;
        let mut recolour = Recolour {
            old: Colour::Unique,
            new: Colour::Unreachable,
            address: 0,
            marker: PhantomData,
        };
        let $x : *mut $t = recolour.set_address_mut(tmp);
    };
    // Cast a &T to a *const T
    (let $x:ident : *const $t:ty = $e:expr) => {
        let tmp: & $t = $e;
        let mut recolour = Recolour {
            old: Colour::Shared,
            new: Colour::Shared,
            address: 0,
            marker: PhantomData,
        };
        let $x : *const $t = recolour.set_address(tmp);
    };
    // Borrow mutably from a RefCell
    (let mut $x:ident : RefMut<$t:ty> = $y:ident .borrow_mut()) => {
        let tmp: &RefCell<$t> = $y;
        let mut $x : RefMut<$t> = tmp.borrow_mut();
        <$t as Colourable>::colour(&*$x as *const $t as Address, Colour::Unreachable, Colour::Unique);
    };
    // Drop a mutable borrow from a RefCell
    (drop ($x:ident as RefMut<$t:ty>)) => {
        let tmp: RefMut<$t> = $x;
        <$t as Colourable>::colour(&*tmp as *const $t as Address, Colour::Unique, Colour::Unreachable);
        drop (tmp);
    };

    // Casting between raw pointers doesn't require tracking
}

// An example, showing how we can track metadata for how memory is being used

fn bad(p: &usize) {
    track!(let q : *const usize = p);
    track!(let r : &mut usize = unsafe { &mut *(q as *mut usize) });
    *r = 5;
}

fn simple() {
    track!(let x : &mut usize = &mut 37);
    {
        track!(let y : &usize = x);
        bad(y);
    }
    *x += 1;
    print!("*x = {}", *x);
}

// Another example, from Ralf Jung

fn aliasing_mut_and_shr() {
    fn inner(rc: &RefCell<u32>, aliasing: &mut u32) {
        *aliasing += 4;
        track!(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
        track!(let aliasing : &u32 = aliasing);
        let _val = *aliasing;
        track!(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);
    track!(let rc: &mut RefCell<u32> = &mut refcell);
    track!(let rc: &RefCell<u32> = rc);
    track!(let mut bmut : RefMut<u32> = rc.borrow_mut());
    inner(&rc, &mut *bmut);
    track!(drop (bmut as RefMut<u32>));
    assert_eq!(*rc.borrow(), 23+12);
}

fn main() {
    aliasing_mut_and_shr()
}