dprint_core/formatting/
thread_state.rs1use std::borrow::Cow;
2use std::cell::UnsafeCell;
3use std::rc::Rc;
4
5use super::collections::GraphNode;
6use super::collections::NodeStackNode;
7use super::Condition;
8use super::ConditionResolver;
9use super::PrintNodeCell;
10use super::SavePoint;
11use super::StringContainer;
12use super::UnsafePrintLifetime;
13use super::WriteItem;
14
15#[derive(Default)]
16pub struct Counts {
17 line_number_anchor_id_count: u32,
18 line_number_id_count: u32,
19 column_number_id_count: u32,
20 is_start_of_line_id: u32,
21 indent_level_id_count: u32,
22 line_start_column_number_id_count: u32,
23 line_start_indent_level_id_count: u32,
24 condition_id_count: u32,
25 condition_reevaluation_id_count: u32,
26 #[cfg(feature = "tracing")]
27 print_node_id_count: u32,
28 #[cfg(feature = "tracing")]
29 graph_node_id_count: u32,
30}
31
32pub struct BumpAllocator {
39 condition_resolvers: Vec<ConditionResolver>,
40 bump: bumpalo::Bump,
41}
42
43impl BumpAllocator {
44 fn new() -> Self {
45 Self {
46 condition_resolvers: Default::default(),
47 bump: bumpalo::Bump::new(),
48 }
49 }
50
51 pub fn inner(&self) -> &bumpalo::Bump {
52 &self.bump
53 }
54
55 pub fn alloc_condition(&mut self, condition: Condition) -> UnsafePrintLifetime<Condition> {
56 unsafe {
57 let condition_resolver = condition.condition.clone();
61 let rc_raw = Rc::into_raw(condition_resolver);
62 Rc::decrement_strong_count(rc_raw);
63 self.condition_resolvers.push(Rc::from_raw(rc_raw));
64 }
65 let condition = self.bump.alloc(condition);
66 unsafe { std::mem::transmute::<&Condition, UnsafePrintLifetime<Condition>>(condition) }
67 }
68
69 pub fn alloc_string(&self, item: Cow<'static, str>) -> UnsafePrintLifetime<StringContainer> {
70 let string = match item {
71 Cow::Borrowed(item) => item,
72 Cow::Owned(item) => {
73 let string = self.bump.alloc(bumpalo::collections::String::from_str_in(&item, &self.bump));
74 unsafe { std::mem::transmute::<&bumpalo::collections::String, UnsafePrintLifetime<bumpalo::collections::String>>(string) }
75 }
76 };
77 let string = StringContainer::new(string);
78 let string = self.bump.alloc(string);
79 unsafe { std::mem::transmute::<&StringContainer, UnsafePrintLifetime<StringContainer>>(string) }
80 }
81
82 pub fn alloc_write_item_graph_node<'a>(&'a self, node: GraphNode<'a, WriteItem<'a>>) -> &'a GraphNode<'a, WriteItem<'a>> {
83 self.bump.alloc(node)
84 }
85
86 pub fn alloc_node_stack_node<'a>(&'a self, node: NodeStackNode<'a>) -> &'a NodeStackNode<'a> {
87 self.bump.alloc(node)
88 }
89
90 pub fn alloc_save_point<'a>(&'a self, save_point: SavePoint<'a>) -> &'a SavePoint<'a> {
91 self.bump.alloc(save_point)
92 }
93
94 pub fn alloc_print_node_cell(&self, cell: PrintNodeCell) -> UnsafePrintLifetime<PrintNodeCell> {
95 let result = self.bump.alloc(cell);
96 unsafe { std::mem::transmute::<&PrintNodeCell, UnsafePrintLifetime<PrintNodeCell>>(result) }
97 }
98
99 pub fn reset(&mut self) {
100 self.bump.reset();
101 self.condition_resolvers.clear();
102 }
103}
104
105thread_local! {
106 static BUMP_ALLOCATOR: UnsafeCell<BumpAllocator> = UnsafeCell::new(BumpAllocator::new());
107 static COUNTS: UnsafeCell<Counts> = UnsafeCell::new(Default::default());
108}
109
110pub fn with_bump_allocator<TReturn>(action: impl FnOnce(&mut BumpAllocator) -> TReturn) -> TReturn {
111 BUMP_ALLOCATOR.with(|bump_cell| unsafe {
112 let bump = bump_cell.get();
113 action(&mut *bump)
114 })
115}
116
117pub fn take_counts() -> Counts {
118 COUNTS.with(|cell| unsafe { std::mem::take(&mut (*cell.get())) })
119}
120
121pub fn set_counts(counts: Counts) {
122 COUNTS.with(|cell| unsafe {
123 *cell.get() = counts;
124 })
125}
126
127pub fn next_line_number_anchor_id() -> u32 {
128 COUNTS.with(|cell| unsafe {
129 let counts = &mut *cell.get();
130 let value = counts.line_number_anchor_id_count;
131 counts.line_number_anchor_id_count += 1;
132 value
133 })
134}
135
136pub fn next_line_number_id() -> u32 {
137 COUNTS.with(|cell| unsafe {
138 let counts = &mut *cell.get();
139 let value = counts.line_number_id_count;
140 counts.line_number_id_count += 1;
141 value
142 })
143}
144
145pub fn next_column_number_id() -> u32 {
146 COUNTS.with(|cell| unsafe {
147 let counts = &mut *cell.get();
148 let value = counts.column_number_id_count;
149 counts.column_number_id_count += 1;
150 value
151 })
152}
153
154pub fn next_is_start_of_line_id() -> u32 {
155 COUNTS.with(|cell| unsafe {
156 let counts = &mut *cell.get();
157 let value = counts.is_start_of_line_id;
158 counts.is_start_of_line_id += 1;
159 value
160 })
161}
162
163pub fn next_indent_level_id() -> u32 {
164 COUNTS.with(|cell| unsafe {
165 let counts = &mut *cell.get();
166 let value = counts.indent_level_id_count;
167 counts.indent_level_id_count += 1;
168 value
169 })
170}
171
172pub fn next_line_start_column_number_id() -> u32 {
173 COUNTS.with(|cell| unsafe {
174 let counts = &mut *cell.get();
175 let value = counts.line_start_column_number_id_count;
176 counts.line_start_column_number_id_count += 1;
177 value
178 })
179}
180
181pub fn next_line_start_indent_level_id() -> u32 {
182 COUNTS.with(|cell| unsafe {
183 let counts = &mut *cell.get();
184 let value = counts.line_start_indent_level_id_count;
185 counts.line_start_indent_level_id_count += 1;
186 value
187 })
188}
189
190pub fn next_condition_id() -> u32 {
191 COUNTS.with(|cell| unsafe {
192 let counts = &mut *cell.get();
193 let value = counts.condition_id_count;
194 counts.condition_id_count += 1;
195 value
196 })
197}
198
199pub fn next_condition_reevaluation_id() -> u32 {
200 COUNTS.with(|cell| unsafe {
201 let counts = &mut *cell.get();
202 let value = counts.condition_reevaluation_id_count;
203 counts.condition_reevaluation_id_count += 1;
204 value
205 })
206}
207
208#[cfg(feature = "tracing")]
209pub fn next_print_node_id() -> u32 {
210 COUNTS.with(|cell| unsafe {
211 let counts = &mut *cell.get();
212 let value = counts.print_node_id_count;
213 counts.print_node_id_count += 1;
214 value
215 })
216}
217
218#[cfg(feature = "tracing")]
219pub fn next_graph_node_id() -> u32 {
220 COUNTS.with(|cell| unsafe {
221 let counts = &mut *cell.get();
222 let value = counts.graph_node_id_count;
223 counts.graph_node_id_count += 1;
224 value
225 })
226}