rustc_builtin_macros/
global_allocator.rs

1use rustc_ast::expand::allocator::{
2    ALLOCATOR_METHODS, AllocatorMethod, AllocatorMethodInput, AllocatorTy, global_fn_name,
3};
4use rustc_ast::{
5    self as ast, AttrVec, Expr, Fn, FnHeader, FnSig, Generics, ItemKind, Mutability, Param, Safety,
6    Stmt, StmtKind, Ty, TyKind,
7};
8use rustc_expand::base::{Annotatable, ExtCtxt};
9use rustc_span::{Ident, Span, Symbol, kw, sym};
10use thin_vec::{ThinVec, thin_vec};
11
12use crate::errors;
13use crate::util::check_builtin_macro_attribute;
14
15pub(crate) fn expand(
16    ecx: &mut ExtCtxt<'_>,
17    _span: Span,
18    meta_item: &ast::MetaItem,
19    item: Annotatable,
20) -> Vec<Annotatable> {
21    check_builtin_macro_attribute(ecx, meta_item, sym::global_allocator);
22
23    let orig_item = item.clone();
24
25    // Allow using `#[global_allocator]` on an item statement
26    // FIXME - if we get deref patterns, use them to reduce duplication here
27    let (item, ident, is_stmt, ty_span) = if let Annotatable::Item(item) = &item
28        && let ItemKind::Static(box ast::StaticItem { ident, ty, .. }) = &item.kind
29    {
30        (item, *ident, false, ecx.with_def_site_ctxt(ty.span))
31    } else if let Annotatable::Stmt(stmt) = &item
32        && let StmtKind::Item(item) = &stmt.kind
33        && let ItemKind::Static(box ast::StaticItem { ident, ty, .. }) = &item.kind
34    {
35        (item, *ident, true, ecx.with_def_site_ctxt(ty.span))
36    } else {
37        ecx.dcx().emit_err(errors::AllocMustStatics { span: item.span() });
38        return vec![orig_item];
39    };
40
41    // Generate a bunch of new items using the AllocFnFactory
42    let span = ecx.with_def_site_ctxt(item.span);
43    let f = AllocFnFactory { span, ty_span, global: ident, cx: ecx };
44
45    // Generate item statements for the allocator methods.
46    let stmts = ALLOCATOR_METHODS.iter().map(|method| f.allocator_fn(method)).collect();
47
48    // Generate anonymous constant serving as container for the allocator methods.
49    let const_ty = ecx.ty(ty_span, TyKind::Tup(ThinVec::new()));
50    let const_body = ast::ConstItemRhs::Body(ecx.expr_block(ecx.block(span, stmts)));
51    let const_item = ecx.item_const(span, Ident::new(kw::Underscore, span), const_ty, const_body);
52    let const_item = if is_stmt {
53        Annotatable::Stmt(Box::new(ecx.stmt_item(span, const_item)))
54    } else {
55        Annotatable::Item(const_item)
56    };
57
58    // Return the original item and the new methods.
59    vec![orig_item, const_item]
60}
61
62struct AllocFnFactory<'a, 'b> {
63    span: Span,
64    ty_span: Span,
65    global: Ident,
66    cx: &'a ExtCtxt<'b>,
67}
68
69impl AllocFnFactory<'_, '_> {
70    fn allocator_fn(&self, method: &AllocatorMethod) -> Stmt {
71        let mut abi_args = ThinVec::new();
72        let args = method.inputs.iter().map(|input| self.arg_ty(input, &mut abi_args)).collect();
73        let result = self.call_allocator(method.name, args);
74        let output_ty = self.ret_ty(&method.output);
75        let decl = self.cx.fn_decl(abi_args, ast::FnRetTy::Ty(output_ty));
76        let header = FnHeader { safety: Safety::Unsafe(self.span), ..FnHeader::default() };
77        let sig = FnSig { decl, header, span: self.span };
78        let body = Some(self.cx.block_expr(result));
79        let kind = ItemKind::Fn(Box::new(Fn {
80            defaultness: ast::Defaultness::Final,
81            sig,
82            ident: Ident::from_str_and_span(&global_fn_name(method.name), self.span),
83            generics: Generics::default(),
84            contract: None,
85            body,
86            define_opaque: None,
87            eii_impls: ThinVec::new(),
88        }));
89        let item = self.cx.item(self.span, self.attrs(method), kind);
90        self.cx.stmt_item(self.ty_span, item)
91    }
92
93    fn call_allocator(&self, method: Symbol, mut args: ThinVec<Box<Expr>>) -> Box<Expr> {
94        let method = self.cx.std_path(&[sym::alloc, sym::GlobalAlloc, method]);
95        let method = self.cx.expr_path(self.cx.path(self.ty_span, method));
96        let allocator = self.cx.path_ident(self.ty_span, self.global);
97        let allocator = self.cx.expr_path(allocator);
98        let allocator = self.cx.expr_addr_of(self.ty_span, allocator);
99        args.insert(0, allocator);
100
101        self.cx.expr_call(self.ty_span, method, args)
102    }
103
104    fn attrs(&self, method: &AllocatorMethod) -> AttrVec {
105        let alloc_attr = match method.name {
106            sym::alloc => sym::rustc_allocator,
107            sym::dealloc => sym::rustc_deallocator,
108            sym::realloc => sym::rustc_reallocator,
109            sym::alloc_zeroed => sym::rustc_allocator_zeroed,
110            _ => unreachable!("Unknown allocator method!"),
111        };
112        thin_vec![
113            self.cx.attr_word(sym::rustc_std_internal_symbol, self.span),
114            self.cx.attr_word(alloc_attr, self.span)
115        ]
116    }
117
118    fn arg_ty(&self, input: &AllocatorMethodInput, args: &mut ThinVec<Param>) -> Box<Expr> {
119        match input.ty {
120            AllocatorTy::Layout => {
121                // If an allocator method is ever introduced having multiple
122                // Layout arguments, these argument names need to be
123                // disambiguated somehow. Currently the generated code would
124                // fail to compile with "identifier is bound more than once in
125                // this parameter list".
126                let size = Ident::from_str_and_span("size", self.span);
127                let align = Ident::from_str_and_span("align", self.span);
128
129                let usize = self.cx.path_ident(self.span, Ident::new(sym::usize, self.span));
130                let ty_usize = self.cx.ty_path(usize);
131                args.push(self.cx.param(self.span, size, ty_usize.clone()));
132                args.push(self.cx.param(self.span, align, ty_usize));
133
134                let layout_new =
135                    self.cx.std_path(&[sym::alloc, sym::Layout, sym::from_size_align_unchecked]);
136                let layout_new = self.cx.expr_path(self.cx.path(self.span, layout_new));
137                let size = self.cx.expr_ident(self.span, size);
138                let align = self.cx.expr_ident(self.span, align);
139                let layout = self.cx.expr_call(self.span, layout_new, thin_vec![size, align]);
140                layout
141            }
142
143            AllocatorTy::Ptr => {
144                let ident = Ident::from_str_and_span(input.name, self.span);
145                args.push(self.cx.param(self.span, ident, self.ptr_u8()));
146                self.cx.expr_ident(self.span, ident)
147            }
148
149            AllocatorTy::Usize => {
150                let ident = Ident::from_str_and_span(input.name, self.span);
151                args.push(self.cx.param(self.span, ident, self.usize()));
152                self.cx.expr_ident(self.span, ident)
153            }
154
155            AllocatorTy::Never | AllocatorTy::ResultPtr | AllocatorTy::Unit => {
156                panic!("can't convert AllocatorTy to an argument")
157            }
158        }
159    }
160
161    fn ret_ty(&self, ty: &AllocatorTy) -> Box<Ty> {
162        match *ty {
163            AllocatorTy::ResultPtr => self.ptr_u8(),
164
165            AllocatorTy::Unit => self.cx.ty(self.span, TyKind::Tup(ThinVec::new())),
166
167            AllocatorTy::Layout | AllocatorTy::Never | AllocatorTy::Usize | AllocatorTy::Ptr => {
168                panic!("can't convert `AllocatorTy` to an output")
169            }
170        }
171    }
172
173    fn usize(&self) -> Box<Ty> {
174        let usize = self.cx.path_ident(self.span, Ident::new(sym::usize, self.span));
175        self.cx.ty_path(usize)
176    }
177
178    fn ptr_u8(&self) -> Box<Ty> {
179        let u8 = self.cx.path_ident(self.span, Ident::new(sym::u8, self.span));
180        let ty_u8 = self.cx.ty_path(u8);
181        self.cx.ty_ptr(self.span, ty_u8, Mutability::Mut)
182    }
183}
rustc_builtin_macros/global_allocator.rs

rustc_builtin_macros/
global_allocator.rs
