dprint_core/formatting/
thread_state.rs

1use 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
32/// This file is a dumpster fireand the API is actually really
33/// unsafe if used incorrectly. The way it ended up this way was
34/// because I wanted to try out the performance improvements of
35/// a bump allocator without having to deal with lifetimes everywhere.
36/// Anyway, this landed and I haven't had time to go through and
37/// make the API safe.
38pub 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      // Leak the Rc that gets stored in the bump allocator, then add another
58      // rc in a vector that we store here, which will cause a decrement of the
59      // rc when this is dropped.
60      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}