1use rustc_ast::expand::allocator::{
2 ALLOCATOR_METHODS, AllocatorMethod, AllocatorMethodInput, AllocatorTy, global_fn_name,
3};
4use rustc_ast::ptr::P;
5use rustc_ast::{
6 self as ast, AttrVec, Expr, Fn, FnHeader, FnSig, Generics, ItemKind, Mutability, Param, Safety,
7 Stmt, StmtKind, Ty, TyKind,
8};
9use rustc_expand::base::{Annotatable, ExtCtxt};
10use rustc_span::{Ident, Span, Symbol, kw, sym};
11use thin_vec::{ThinVec, thin_vec};
12
13use crate::errors;
14use crate::util::check_builtin_macro_attribute;
15
16pub(crate) fn expand(
17 ecx: &mut ExtCtxt<'_>,
18 _span: Span,
19 meta_item: &ast::MetaItem,
20 item: Annotatable,
21) -> Vec<Annotatable> {
22 check_builtin_macro_attribute(ecx, meta_item, sym::global_allocator);
23
24 let orig_item = item.clone();
25
26 let (item, ident, is_stmt, ty_span) = if let Annotatable::Item(item) = &item
29 && let ItemKind::Static(box ast::StaticItem { ident, ty, .. }) = &item.kind
30 {
31 (item, *ident, false, ecx.with_def_site_ctxt(ty.span))
32 } else if let Annotatable::Stmt(stmt) = &item
33 && let StmtKind::Item(item) = &stmt.kind
34 && let ItemKind::Static(box ast::StaticItem { ident, ty, .. }) = &item.kind
35 {
36 (item, *ident, true, ecx.with_def_site_ctxt(ty.span))
37 } else {
38 ecx.dcx().emit_err(errors::AllocMustStatics { span: item.span() });
39 return vec![orig_item];
40 };
41
42 let span = ecx.with_def_site_ctxt(item.span);
44 let f = AllocFnFactory { span, ty_span, global: ident, cx: ecx };
45
46 let stmts = ALLOCATOR_METHODS.iter().map(|method| f.allocator_fn(method)).collect();
48
49 let const_ty = ecx.ty(ty_span, TyKind::Tup(ThinVec::new()));
51 let const_body = ecx.expr_block(ecx.block(span, stmts));
52 let const_item = ecx.item_const(span, Ident::new(kw::Underscore, span), const_ty, const_body);
53 let const_item = if is_stmt {
54 Annotatable::Stmt(P(ecx.stmt_item(span, const_item)))
55 } else {
56 Annotatable::Item(const_item)
57 };
58
59 vec![orig_item, const_item]
61}
62
63struct AllocFnFactory<'a, 'b> {
64 span: Span,
65 ty_span: Span,
66 global: Ident,
67 cx: &'a ExtCtxt<'b>,
68}
69
70impl AllocFnFactory<'_, '_> {
71 fn allocator_fn(&self, method: &AllocatorMethod) -> Stmt {
72 let mut abi_args = ThinVec::new();
73 let args = method.inputs.iter().map(|input| self.arg_ty(input, &mut abi_args)).collect();
74 let result = self.call_allocator(method.name, args);
75 let output_ty = self.ret_ty(&method.output);
76 let decl = self.cx.fn_decl(abi_args, ast::FnRetTy::Ty(output_ty));
77 let header = FnHeader { safety: Safety::Unsafe(self.span), ..FnHeader::default() };
78 let sig = FnSig { decl, header, span: self.span };
79 let body = Some(self.cx.block_expr(result));
80 let kind = ItemKind::Fn(Box::new(Fn {
81 defaultness: ast::Defaultness::Final,
82 sig,
83 ident: Ident::from_str_and_span(&global_fn_name(method.name), self.span),
84 generics: Generics::default(),
85 contract: None,
86 body,
87 define_opaque: None,
88 }));
89 let item = self.cx.item(self.span, self.attrs(), kind);
90 self.cx.stmt_item(self.ty_span, item)
91 }
92
93 fn call_allocator(&self, method: Symbol, mut args: ThinVec<P<Expr>>) -> P<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) -> AttrVec {
105 thin_vec![self.cx.attr_word(sym::rustc_std_internal_symbol, self.span)]
106 }
107
108 fn arg_ty(&self, input: &AllocatorMethodInput, args: &mut ThinVec<Param>) -> P<Expr> {
109 match input.ty {
110 AllocatorTy::Layout => {
111 let size = Ident::from_str_and_span("size", self.span);
117 let align = Ident::from_str_and_span("align", self.span);
118
119 let usize = self.cx.path_ident(self.span, Ident::new(sym::usize, self.span));
120 let ty_usize = self.cx.ty_path(usize);
121 args.push(self.cx.param(self.span, size, ty_usize.clone()));
122 args.push(self.cx.param(self.span, align, ty_usize));
123
124 let layout_new =
125 self.cx.std_path(&[sym::alloc, sym::Layout, sym::from_size_align_unchecked]);
126 let layout_new = self.cx.expr_path(self.cx.path(self.span, layout_new));
127 let size = self.cx.expr_ident(self.span, size);
128 let align = self.cx.expr_ident(self.span, align);
129 let layout = self.cx.expr_call(self.span, layout_new, thin_vec![size, align]);
130 layout
131 }
132
133 AllocatorTy::Ptr => {
134 let ident = Ident::from_str_and_span(input.name, self.span);
135 args.push(self.cx.param(self.span, ident, self.ptr_u8()));
136 self.cx.expr_ident(self.span, ident)
137 }
138
139 AllocatorTy::Usize => {
140 let ident = Ident::from_str_and_span(input.name, self.span);
141 args.push(self.cx.param(self.span, ident, self.usize()));
142 self.cx.expr_ident(self.span, ident)
143 }
144
145 AllocatorTy::ResultPtr | AllocatorTy::Unit => {
146 panic!("can't convert AllocatorTy to an argument")
147 }
148 }
149 }
150
151 fn ret_ty(&self, ty: &AllocatorTy) -> P<Ty> {
152 match *ty {
153 AllocatorTy::ResultPtr => self.ptr_u8(),
154
155 AllocatorTy::Unit => self.cx.ty(self.span, TyKind::Tup(ThinVec::new())),
156
157 AllocatorTy::Layout | AllocatorTy::Usize | AllocatorTy::Ptr => {
158 panic!("can't convert `AllocatorTy` to an output")
159 }
160 }
161 }
162
163 fn usize(&self) -> P<Ty> {
164 let usize = self.cx.path_ident(self.span, Ident::new(sym::usize, self.span));
165 self.cx.ty_path(usize)
166 }
167
168 fn ptr_u8(&self) -> P<Ty> {
169 let u8 = self.cx.path_ident(self.span, Ident::new(sym::u8, self.span));
170 let ty_u8 = self.cx.ty_path(u8);
171 self.cx.ty_ptr(self.span, ty_u8, Mutability::Mut)
172 }
173}