1use brokk_bifrost_core::hash::{HashMap, HashSet};
19use tree_sitter::Node;
20
21pub trait LocalBindingBudget {
29 type Error;
30 fn enter_node(&mut self) -> Result<(), Self::Error>;
32 fn before_insert(&mut self) -> Result<(), Self::Error>;
34 fn charge_name(&mut self, name: &str) -> Result<(), Self::Error>;
36}
37
38#[derive(Debug, Default)]
41pub struct UnboundedLocalBindingBudget;
42
43impl LocalBindingBudget for UnboundedLocalBindingBudget {
44 type Error = std::convert::Infallible;
45
46 fn enter_node(&mut self) -> Result<(), Self::Error> {
47 Ok(())
48 }
49
50 fn before_insert(&mut self) -> Result<(), Self::Error> {
51 Ok(())
52 }
53
54 fn charge_name(&mut self, _name: &str) -> Result<(), Self::Error> {
55 Ok(())
56 }
57}
58
59#[derive(Clone, Default)]
66pub struct LocalBindingTimeline {
67 entry_bindings: HashSet<Box<str>>,
68 activations: HashMap<Box<str>, usize>,
69}
70
71impl LocalBindingTimeline {
72 pub fn is_active_at(&self, name: &str, source_byte: usize) -> bool {
73 self.entry_bindings.contains(name)
74 || self
75 .activations
76 .get(name)
77 .is_some_and(|activation| *activation <= source_byte)
78 }
79
80 pub fn activations(&self) -> impl Iterator<Item = (&str, usize)> {
83 self.activations
84 .iter()
85 .map(|(name, start)| (name.as_ref(), *start))
86 }
87
88 pub fn active_names_at(&self, source_byte: usize) -> Vec<&str> {
89 let mut names = self
90 .entry_bindings
91 .iter()
92 .map(Box::as_ref)
93 .chain(
94 self.activations
95 .iter()
96 .filter(|(_, activation)| **activation <= source_byte)
97 .map(|(name, _)| name.as_ref()),
98 )
99 .collect::<Vec<_>>();
100 names.sort_unstable();
101 names
102 }
103}
104
105pub struct LocalBindingCollection {
106 pub timeline: LocalBindingTimeline,
107 pub has_parameter_defaults: bool,
108}
109
110fn node_text<'source>(source: &'source str, node: Node<'_>) -> Option<&'source str> {
111 node.utf8_text(source.as_bytes()).ok()
112}
113
114fn named_children(node: Node<'_>) -> Vec<Node<'_>> {
115 (0..node.named_child_count())
116 .filter_map(|index| node.named_child(index))
117 .collect()
118}
119
120fn children_by_field_name<'tree>(node: Node<'tree>, field: &str) -> Vec<Node<'tree>> {
121 node.child_by_field_name(field).into_iter().collect()
122}
123
124struct LocalBindingCollector<'source, 'request, B: LocalBindingBudget> {
125 source: &'source str,
126 timeline: LocalBindingTimeline,
127 has_parameter_defaults: bool,
128 budget: &'request mut B,
129}
130
131impl<'source, 'request, B: LocalBindingBudget> LocalBindingCollector<'source, 'request, B> {
132 fn new(source: &'source str, budget: &'request mut B) -> Self {
133 Self {
134 source,
135 timeline: LocalBindingTimeline::default(),
136 has_parameter_defaults: false,
137 budget,
138 }
139 }
140
141 fn visit(&mut self) -> Result<(), B::Error> {
142 self.budget.enter_node()
143 }
144
145 fn insert_entry_name(&mut self, name: &str) -> Result<(), B::Error> {
146 self.budget.before_insert()?;
147 if self.timeline.entry_bindings.contains(name) {
148 return Ok(());
149 }
150 if self.timeline.activations.remove(name).is_none() {
151 self.budget.charge_name(name)?;
152 }
153 self.timeline.entry_bindings.insert(name.into());
154 Ok(())
155 }
156
157 fn insert_activation(&mut self, name: &str, source_byte: usize) -> Result<(), B::Error> {
158 self.budget.before_insert()?;
159 if self.timeline.entry_bindings.contains(name) {
160 return Ok(());
161 }
162 if let Some(activation) = self.timeline.activations.get_mut(name) {
163 *activation = (*activation).min(source_byte);
164 return Ok(());
165 }
166 self.budget.charge_name(name)?;
167 self.timeline.activations.insert(name.into(), source_byte);
168 Ok(())
169 }
170
171 fn insert_entry_identifier(&mut self, node: Node<'_>) -> Result<(), B::Error> {
172 if node.kind() == "identifier"
173 && let Some(name) = node_text(self.source, node)
174 {
175 self.insert_entry_name(name)?;
176 }
177 Ok(())
178 }
179
180 fn insert_activation_identifier(&mut self, node: Node<'_>) -> Result<(), B::Error> {
181 if node.kind() == "identifier"
182 && let Some(name) = node_text(self.source, node)
183 {
184 self.insert_activation(name, node.start_byte())?;
185 }
186 Ok(())
187 }
188
189 fn collect_parameters(&mut self, node: Node<'_>) -> Result<(), B::Error> {
190 let mut stack = vec![node];
191 while let Some(current) = stack.pop() {
192 self.visit()?;
193 match current.kind() {
194 "identifier" => self.insert_entry_identifier(current)?,
195 "optional_parameter"
196 | "keyword_parameter"
197 | "splat_parameter"
198 | "hash_splat_parameter"
199 | "block_parameter" => {
200 self.has_parameter_defaults |= current.kind() == "optional_parameter"
201 || (current.kind() == "keyword_parameter"
202 && current.child_by_field_name("value").is_some());
203 if let Some(name) = current.child_by_field_name("name") {
204 self.insert_entry_identifier(name)?;
205 }
206 }
207 "method_parameters"
208 | "lambda_parameters"
209 | "block_parameters"
210 | "destructured_parameter" => {
211 stack.extend(named_children(current).into_iter().rev());
212 }
213 "forward_parameter" | "hash_splat_nil" => {}
214 _ => {}
215 }
216 }
217 Ok(())
218 }
219
220 fn collect_assignment(&mut self, node: Node<'_>) -> Result<(), B::Error> {
221 let mut stack = vec![node];
222 while let Some(current) = stack.pop() {
223 self.visit()?;
224 match current.kind() {
225 "identifier" => self.insert_activation_identifier(current)?,
226 "left_assignment_list"
227 | "right_assignment_list"
228 | "destructured_left_assignment"
229 | "rest_assignment"
230 | "exception_variable" => {
231 stack.extend(named_children(current).into_iter().rev());
232 }
233 _ => {}
234 }
235 }
236 Ok(())
237 }
238
239 fn collect_pattern(&mut self, node: Node<'_>) -> Result<(), B::Error> {
240 let mut stack = vec![node];
241 while let Some(current) = stack.pop() {
242 self.visit()?;
243 match current.kind() {
244 "identifier" => self.insert_activation_identifier(current)?,
245 "as_pattern" => {
246 if let Some(name) = current.child_by_field_name("name") {
247 self.insert_activation_identifier(name)?;
248 }
249 stack.extend(children_by_field_name(current, "value"));
250 }
251 "keyword_pattern" => {
252 if let Some(value) = current.child_by_field_name("value") {
253 stack.push(value);
254 } else if let Some(key) = current.child_by_field_name("key")
255 && let Some(name) = node_text(self.source, key)
256 {
257 self.insert_activation(
258 name.strip_suffix(':').unwrap_or(name),
259 key.start_byte(),
260 )?;
261 }
262 }
263 "splat_parameter" | "hash_splat_parameter" => {
264 if let Some(name) = current.child_by_field_name("name") {
265 self.insert_activation_identifier(name)?;
266 }
267 }
268 "variable_reference_pattern" | "expression_reference_pattern" => {}
269 "array_pattern" | "find_pattern" | "hash_pattern" => {
270 let class_id = current.child_by_field_name("class").map(|class| class.id());
271 let children = named_children(current)
272 .into_iter()
273 .filter(|child| Some(child.id()) != class_id)
274 .collect::<Vec<_>>();
275 stack.extend(children.into_iter().rev());
276 }
277 "alternative_pattern" | "parenthesized_pattern" => {
278 stack.extend(named_children(current).into_iter().rev());
279 }
280 _ => {}
281 }
282 }
283 Ok(())
284 }
285
286 fn finish(self) -> LocalBindingCollection {
287 LocalBindingCollection {
288 timeline: self.timeline,
289 has_parameter_defaults: self.has_parameter_defaults,
290 }
291 }
292}
293
294pub fn callable_parameters<'tree>(callable: Node<'tree>, body: Node<'tree>) -> Option<Node<'tree>> {
297 callable.child_by_field_name("parameters").or_else(|| {
298 (callable.kind() == "lambda")
299 .then(|| body.child_by_field_name("parameters"))
300 .flatten()
301 })
302}
303
304pub fn collect_local_bindings<B: LocalBindingBudget>(
312 source: &str,
313 callable: Node<'_>,
314 body: Node<'_>,
315 inherited: Option<(&LocalBindingTimeline, usize)>,
316 budget: &mut B,
317) -> Result<LocalBindingCollection, B::Error> {
318 let mut collector = LocalBindingCollector::new(source, budget);
319 let parameters = callable_parameters(callable, body);
320 if let Some(parameters) = parameters {
321 collector.collect_parameters(parameters)?;
322 }
323 if matches!(callable.kind(), "lambda" | "block" | "do_block") && parameters.is_none() {
324 for name in ["_1", "_2", "_3", "_4", "_5", "_6", "_7", "_8", "_9", "it"] {
325 collector.insert_entry_name(name)?;
326 }
327 }
328 if let Some((inherited, source_byte)) = inherited {
329 for name in inherited.active_names_at(source_byte) {
330 collector.insert_entry_name(name)?;
331 }
332 }
333
334 let mut stack = vec![body];
335 if let Some(parameters) = parameters.filter(|parameters| {
336 parameters.start_byte() < body.start_byte() || parameters.end_byte() > body.end_byte()
337 }) {
338 stack.push(parameters);
339 }
340 while let Some(node) = stack.pop() {
341 collector.visit()?;
342 match node.kind() {
343 "assignment" | "operator_assignment" => {
344 if let Some(left) = node.child_by_field_name("left") {
345 collector.collect_assignment(left)?;
346 }
347 }
348 "for" => {
349 if let Some(pattern) = node.child_by_field_name("pattern") {
350 collector.collect_assignment(pattern)?;
351 }
352 }
353 "rescue" => {
354 if let Some(variable) = node.child_by_field_name("variable") {
355 collector.collect_assignment(variable)?;
356 }
357 }
358 "match_pattern" | "test_pattern" | "in_clause" => {
359 if let Some(pattern) = node.child_by_field_name("pattern") {
360 collector.collect_pattern(pattern)?;
361 }
362 }
363 _ => {}
364 }
365 for child in named_children(node).into_iter().rev() {
366 if child.id() != body.id()
367 && matches!(
368 child.kind(),
369 "method"
370 | "singleton_method"
371 | "lambda"
372 | "block"
373 | "do_block"
374 | "class"
375 | "module"
376 | "singleton_class"
377 )
378 {
379 continue;
380 }
381 stack.push(child);
382 }
383 }
384 Ok(collector.finish())
385}