src/libproc_macro/bridge/handle.rs

   1 // Copyright 2018 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11 //! Server-side handles and storage for per-handle data.
  12
  13 use std::collections::{BTreeMap, HashMap};
  14 use std::hash::Hash;
  15 use std::num::NonZeroU32;
  16 use std::ops::{Index, IndexMut};
  17 use std::sync::atomic::{AtomicUsize, Ordering};
  18
  19 pub(super) type Handle = NonZeroU32;
  20
  21 pub(super) struct OwnedStore<T: 'static> {
  22     counter: &'static AtomicUsize,
  23     data: BTreeMap<Handle, T>,
  24 }
  25
  26 impl<T> OwnedStore<T> {
  27     pub(super) fn new(counter: &'static AtomicUsize) -> Self {
  28         // Ensure the handle counter isn't 0, which would panic later,
  29         // when `NonZeroU32::new` (aka `Handle::new`) is called in `alloc`.
  30         assert_ne!(counter.load(Ordering::SeqCst), 0);
  31
  32         OwnedStore {
  33             counter,
  34             data: BTreeMap::new(),
  35         }
  36     }
  37 }
  38
  39 impl<T> OwnedStore<T> {
  40     pub(super) fn alloc(&mut self, x: T) -> Handle {
  41         let counter = self.counter.fetch_add(1, Ordering::SeqCst);
  42         let handle = Handle::new(counter as u32).expect("`proc_macro` handle counter overflowed");
  43         assert!(self.data.insert(handle, x).is_none());
  44         handle
  45     }
  46
  47     pub(super) fn take(&mut self, h: Handle) -> T {
  48         self.data
  49             .remove(&h)
  50             .expect("use-after-free in `proc_macro` handle")
  51     }
  52 }
  53
  54 impl<T> Index<Handle> for OwnedStore<T> {
  55     type Output = T;
  56     fn index(&self, h: Handle) -> &T {
  57         self.data
  58             .get(&h)
  59             .expect("use-after-free in `proc_macro` handle")
  60     }
  61 }
  62
  63 impl<T> IndexMut<Handle> for OwnedStore<T> {
  64     fn index_mut(&mut self, h: Handle) -> &mut T {
  65         self.data
  66             .get_mut(&h)
  67             .expect("use-after-free in `proc_macro` handle")
  68     }
  69 }
  70
  71 pub(super) struct InternedStore<T: 'static> {
  72     owned: OwnedStore<T>,
  73     interner: HashMap<T, Handle>,
  74 }
  75
  76 impl<T: Copy + Eq + Hash> InternedStore<T> {
  77     pub(super) fn new(counter: &'static AtomicUsize) -> Self {
  78         InternedStore {
  79             owned: OwnedStore::new(counter),
  80             interner: HashMap::new(),
  81         }
  82     }
  83
  84     pub(super) fn alloc(&mut self, x: T) -> Handle {
  85         let owned = &mut self.owned;
  86         *self.interner.entry(x).or_insert_with(|| owned.alloc(x))
  87     }
  88
  89     pub(super) fn copy(&mut self, h: Handle) -> T {
  90         self.owned[h]
  91     }
  92 }