rustc TyCtxt model

rustc.rs

use std::cell::RefCell;
use std::collections::HashSet;

// This is the mapping of types in this model to types in rustc:
//   model                => rustc
//   TyS<'a>              => TyS<'tcx>
//   Ty<'a>               => Ty<'tcx>
//   Arena                => DroplessArena
//   Interner<'cx>        => CtxtInterner<'tcx>
//   Ctxt<'a, 'gcx, 'lcx> => TyCtxt<'a, 'gcx, 'tcx>
//
// Usually 'tcx is used for the lifetime of the global arena in rustc if there
// isn't not a local arena around. This model always uses 'gcx for the global
// arena and 'lcx for the local arena. 'cx is used for lifetime for data
// usually in either arena.

#[derive(Copy, Clone, Hash, PartialEq, Eq, Debug)]
enum TyS<'cx> {
    Usize,
    Ptr(Ty<'cx>),
}

// note: corresponds to Ty<'tcx> in rustc
type Ty<'cx> = &'cx TyS<'cx>;

// note: corresponds to DroplessArena in rustc
#[derive(Default)]
struct Arena {
    memory: RefCell<Vec<*mut usize>>,
}

impl Arena {
    fn alloc<T: Copy>(&self, val: T) -> &mut T {
        let val = Box::into_raw(Box::new(val));
        // Record the memory so we can free it when dropped
        self.memory.borrow_mut().push(val as *mut usize);
        // This is safe because we do not access `*val` again
        // and the reference cannot outlive our memory
        unsafe { &mut *val }
    }
}

impl Drop for Arena {
    fn drop(&mut self) {
        let mut vec = self.memory.borrow_mut();
        for val in vec.drain(..) {
            // The type of the box here doesn't matter as we cannot allocate things with destructors
            // We just want to free the memory
            let _: Box<usize> = unsafe { Box::from_raw(val) };
        }
    }
}

struct Interner<'cx> {
    arena: &'cx Arena,
    types: RefCell<HashSet<Ty<'cx>>>,
}

#[derive(Copy, Clone)]
struct Ctxt<'a, 'gcx: 'a + 'lcx, 'lcx: 'a> {
    global: &'a Interner<'gcx>,
    local: &'a Interner<'lcx>,
}

impl<'a, 'gcx: 'a + 'lcx, 'lcx: 'a> Ctxt<'a, 'gcx, 'lcx> {
    // This returns a global Ctxt
    fn to_global(self) -> Ctxt<'a, 'gcx, 'gcx> {
        Ctxt {
            global: self.global,
            local: self.global,
        }
    }
    
    fn with_local<F>(self, ty: Ty<'lcx>, f: F) -> Ty<'lcx>
            where F: for<'_a, '_gcx, '_lcx> FnOnce(Ctxt<'_a, '_gcx, '_lcx>, Ty<'_lcx>) -> Ty<'_lcx> {
        let local = Interner {
            arena: &Arena::default(),
            types: RefCell::<HashSet<Ty>>::default(),
        };
        
        let lcx = Ctxt {
            global: self.global,
            local: &local,
        };
        
        let result = f(lcx, ty);
        
        // Lift the local result type so it can outlive
        // the local interner and we can return it
        self.lift(result)
    }
    
    fn intern(self, ty: TyS<'lcx>) -> Ty<'lcx> {
        // Does this type exist in the local interner?
        if let Some(&ty) = self.local.types.borrow().get(&ty) {
            return ty;
        }
        
        // Does this type exist in the global interner?
        if let Some(&ty) = self.global.types.borrow().get(&ty) {
            return ty;
        }
        
        // This is a type never seen before in this Ctxt.
        // It may exist in other TyCtxtes though
        // Allocate it in our local arena
        let ty = self.local.arena.alloc(ty);
        self.local.types.borrow_mut().insert(ty);
        ty
    }
    
    // Moves a type to be allocted in the local interner
    fn lift<'b>(self, ty: Ty<'b>) -> Ty<'lcx> {
        match *ty {
            TyS::Usize => self.intern(TyS::Usize),
            TyS::Ptr(pty) => self.intern(TyS::Ptr(self.lift(pty))),
        }
    }
}

fn query<'a, 'gcx>(gcx: Ctxt<'a, 'gcx, 'gcx>, arg: Ty<'gcx>) -> Ty<'gcx> {
    if true {
        // Returns a global type
        gcx.to_global().lift(&TyS::Ptr(arg))
    } else {
        // Returns a local type
        gcx.lift(&TyS::Ptr(arg))
    }
}

fn typeck<'a, 'gcx, 'lcx>(lcx: Ctxt<'a, 'gcx, 'lcx>, arg: Ty<'lcx>) -> Ty<'lcx> {
    // ERROR: Calling query with a local type. It expects a Ty<'gcx>
    // let local_ty: Ty<'lcx> = lcx.lift(&TyS::Usize);
    // query(tcx.to_global(), local_ty);
    
    let global_ty: Ty<'gcx> = lcx.to_global().lift(&TyS::Usize);
    query(lcx.to_global(), global_ty);
    
    // Return `*arg`
    let return_ty : Ty<'lcx> = lcx.lift(&TyS::Ptr(arg));
    return_ty
}

fn main() {
    let global = Interner {
        arena: &Arena::default(),
        types: RefCell::<HashSet<Ty>>::default(),
    };
    
    let gcx = Ctxt {
        global: &global,
        local: &global,
    };
    
    let ty = gcx.lift(&TyS::Usize);
    
    // We can call typeck through a function which creates a local interner
    let global_ty1 = gcx.with_local(ty, |lcx, lty| typeck(lcx, lty));
    
    // or we can call typeck using only a global interner
    let global_ty2 = typeck(gcx, ty);
    
    assert_eq!(global_ty1, global_ty2);
}

fn _should_not_happen() {
    let arenas = &(Arena::default(), Arena::default());
    
    let global = Interner {
        arena: &arenas.0,
        types: RefCell::<HashSet<Ty>>::default(),
    };
    
    let local = Interner {
        arena: &arenas.1,
        types: RefCell::<HashSet<Ty>>::default(),
    };
    
    let gcx = Ctxt {
        global: &global,
        local: &local,
    };
    
    // ERROR: We are constructing a TyCtxt with two different arenas, but
    // passing it to a function which expects TyCtxt<'a, 'gcx, 'gcx>.
    // These functions assume that there is no local arena and all allocations
    // are global.
    query(gcx, gcx.lift(&TyS::Usize));
}