1use crate::plan::{CommandPlan, CtePlan, CtePlanBody, QueryPlan, RelationalPlan, SourcePlan};
10use crate::SQLError;
11use std::collections::{BTreeMap, BTreeSet};
12
13pub fn cte_reference_name(reference: &str) -> Option<String> {
15 let (schema, name) = crate::RelationIdentity::parse_reference(reference).ok()?;
16 schema.is_none().then_some(name)
17}
18
19pub fn reachable_plan_cte_names(plan: &QueryPlan) -> BTreeSet<String> {
20 let targets = plan
21 .ctes
22 .iter()
23 .map(|cte| cte.name.clone())
24 .collect::<BTreeSet<_>>();
25 if targets.is_empty() {
26 return BTreeSet::new();
27 }
28
29 let mut reachable = plan
30 .ctes
31 .iter()
32 .filter(|cte| cte.body.modifies_data())
33 .map(|cte| cte.name.clone())
34 .collect::<BTreeSet<_>>();
35 collect_target_cte_references_from_root(&plan.root, &targets, &BTreeSet::new(), &mut reachable);
36
37 let mut expanded = BTreeSet::new();
38 loop {
39 let pending = plan
40 .ctes
41 .iter()
42 .enumerate()
43 .filter(|(_, cte)| reachable.contains(&cte.name) && !expanded.contains(&cte.name))
44 .collect::<Vec<_>>();
45 if pending.is_empty() {
46 break;
47 }
48 for (index, cte) in pending {
49 expanded.insert(cte.name.clone());
50 let visible_dependencies = if cte.recursive {
51 targets.clone()
52 } else {
53 plan.ctes[..index]
54 .iter()
55 .map(|dependency| dependency.name.clone())
56 .collect::<BTreeSet<_>>()
57 };
58 collect_target_cte_references_from_body(
59 &cte.body,
60 &visible_dependencies,
61 &BTreeSet::new(),
62 &mut reachable,
63 );
64 }
65 }
66 reachable
67}
68
69pub fn cte_definition_references(ctes: &[CtePlan], index: usize) -> BTreeSet<String> {
71 let cte = &ctes[index];
72 let visible = if cte.recursive {
73 ctes.iter()
74 } else {
75 ctes[..index].iter()
76 }
77 .map(|cte| cte.name.clone())
78 .collect::<BTreeSet<_>>();
79 let mut references = BTreeSet::new();
80 collect_target_cte_references_from_body(&cte.body, &visible, &BTreeSet::new(), &mut references);
81 references
82}
83
84pub fn primary_query_cte_references(plan: &QueryPlan) -> BTreeSet<String> {
86 let targets = plan
87 .ctes
88 .iter()
89 .map(|cte| cte.name.clone())
90 .collect::<BTreeSet<_>>();
91 let mut references = BTreeSet::new();
92 if !targets.is_empty() {
93 collect_target_cte_references_from_root(
94 &plan.root,
95 &targets,
96 &BTreeSet::new(),
97 &mut references,
98 );
99 }
100 references
101}
102
103pub fn primary_command_cte_references(
105 ctes: &[CtePlan],
106 queries: &[&QueryPlan],
107 source: Option<&SourcePlan>,
108) -> BTreeSet<String> {
109 let targets = ctes
110 .iter()
111 .map(|cte| cte.name.clone())
112 .collect::<BTreeSet<_>>();
113 let mut references = BTreeSet::new();
114 if targets.is_empty() {
115 return references;
116 }
117 for query in queries {
118 collect_target_cte_references_from_nested_query(
119 query,
120 &targets,
121 &BTreeSet::new(),
122 &mut references,
123 );
124 }
125 if let Some(source) = source {
126 collect_target_cte_references_from_source(
127 source,
128 &targets,
129 &BTreeSet::new(),
130 &mut references,
131 );
132 }
133 references
134}
135
136pub fn cte_references_own_name(cte: &CtePlan) -> bool {
137 let targets = BTreeSet::from([cte.name.clone()]);
138 let mut references = BTreeSet::new();
139 collect_target_cte_references_from_body(&cte.body, &targets, &BTreeSet::new(), &mut references);
140 references.contains(&cte.name)
141}
142
143pub fn ordered_plan_ctes(plan: &QueryPlan) -> Result<Vec<&CtePlan>, SQLError> {
144 ordered_cte_plans(&plan.ctes)
145}
146
147pub fn ordered_cte_plans(ctes: &[CtePlan]) -> Result<Vec<&CtePlan>, SQLError> {
148 order_cte_plans(ctes.iter().collect())
149}
150
151pub fn order_cte_plans(plans: Vec<&CtePlan>) -> Result<Vec<&CtePlan>, SQLError> {
152 if !plans.iter().any(|cte| cte.recursive) {
153 return Ok(plans);
154 }
155 let targets = plans
156 .iter()
157 .map(|cte| cte.name.clone())
158 .collect::<BTreeSet<_>>();
159 let dependencies = plans
160 .iter()
161 .map(|cte| {
162 let mut references = BTreeSet::new();
163 collect_target_cte_references_from_body(
164 &cte.body,
165 &targets,
166 &BTreeSet::new(),
167 &mut references,
168 );
169 references.remove(&cte.name);
170 references
171 })
172 .collect::<Vec<_>>();
173 let mut emitted = BTreeSet::new();
174 let mut ordered = Vec::with_capacity(plans.len());
175 let mut remaining = (0..plans.len()).collect::<BTreeSet<_>>();
176 while !remaining.is_empty() {
177 let ready = remaining
178 .iter()
179 .copied()
180 .find(|index| dependencies[*index].is_subset(&emitted));
181 let Some(index) = ready else {
182 return Err(SQLError::Routine {
183 sqlstate: "0A000".into(),
184 message: "mutual recursion between WITH items is not implemented".into(),
185 });
186 };
187 remaining.remove(&index);
188 emitted.insert(plans[index].name.clone());
189 ordered.push(plans[index]);
190 }
191 Ok(ordered)
192}
193
194pub fn single_reference_plan_cte_names(plan: &QueryPlan) -> BTreeSet<String> {
196 let targets = plan
197 .ctes
198 .iter()
199 .map(|cte| cte.name.clone())
200 .collect::<BTreeSet<_>>();
201 let mut counts = targets
202 .iter()
203 .map(|name| (name.clone(), 0usize))
204 .collect::<BTreeMap<_, _>>();
205 count_plan_cte_references(plan, &targets, &mut counts);
206 counts
207 .into_iter()
208 .filter_map(|(name, count)| (count == 1).then_some(name))
209 .collect()
210}
211
212fn count_plan_cte_references(
213 plan: &QueryPlan,
214 targets: &BTreeSet<String>,
215 counts: &mut BTreeMap<String, usize>,
216) {
217 for cte in &plan.ctes {
218 count_cte_body_references(&cte.body, targets, counts);
219 }
220 count_relational_cte_references(&plan.root, targets, counts);
221}
222
223fn count_cte_body_references(
224 body: &CtePlanBody,
225 targets: &BTreeSet<String>,
226 counts: &mut BTreeMap<String, usize>,
227) {
228 match body {
229 CtePlanBody::Query(query) => count_plan_cte_references(query, targets, counts),
230 CtePlanBody::Command(command) => {
231 for cte in command.ctes() {
232 count_cte_body_references(&cte.body, targets, counts);
233 }
234 for query in command.query_inputs() {
235 count_plan_cte_references(query, targets, counts);
236 }
237 if let Some(source) = command.source_input() {
238 count_source_cte_references(source, targets, counts);
239 }
240 }
241 }
242}
243
244fn count_relational_cte_references(
245 plan: &RelationalPlan,
246 targets: &BTreeSet<String>,
247 counts: &mut BTreeMap<String, usize>,
248) {
249 match plan {
250 RelationalPlan::QueryBlock(block) => {
251 if let Some(source) = &block.from {
252 count_source_cte_references(source, targets, counts);
253 }
254 for subquery in &block.subqueries {
255 count_plan_cte_references(subquery, targets, counts);
256 }
257 }
258 RelationalPlan::SetOp {
259 left,
260 right,
261 subqueries,
262 ..
263 } => {
264 count_plan_cte_references(left, targets, counts);
265 count_plan_cte_references(right, targets, counts);
266 for subquery in subqueries {
267 count_plan_cte_references(subquery, targets, counts);
268 }
269 }
270 RelationalPlan::Values { subqueries, .. } => {
271 for subquery in subqueries {
272 count_plan_cte_references(subquery, targets, counts);
273 }
274 }
275 }
276}
277
278fn count_source_cte_references(
279 source: &SourcePlan,
280 targets: &BTreeSet<String>,
281 counts: &mut BTreeMap<String, usize>,
282) {
283 match source {
284 SourcePlan::Table { name, .. } => {
285 if let Some(name) = cte_reference_name(name).filter(|name| targets.contains(name)) {
286 *counts.entry(name).or_default() += 1;
287 }
288 }
289 SourcePlan::Join { left, right, .. } => {
290 count_source_cte_references(left, targets, counts);
291 count_source_cte_references(right, targets, counts);
292 }
293 SourcePlan::Subquery { body, .. } => {
294 count_plan_cte_references(body, targets, counts);
295 }
296 SourcePlan::Values { .. }
297 | SourcePlan::Function { .. }
298 | SourcePlan::FunctionGroup { .. } => {}
299 }
300}
301
302fn collect_target_cte_references_from_root(
303 root: &RelationalPlan,
304 targets: &BTreeSet<String>,
305 shadowed: &BTreeSet<String>,
306 references: &mut BTreeSet<String>,
307) {
308 match root {
309 RelationalPlan::QueryBlock(block) => {
310 if let Some(source) = &block.from {
311 collect_target_cte_references_from_source(source, targets, shadowed, references);
312 }
313 for subquery in &block.subqueries {
314 collect_target_cte_references_from_nested_query(
315 subquery, targets, shadowed, references,
316 );
317 }
318 }
319 RelationalPlan::SetOp {
320 left,
321 right,
322 subqueries,
323 ..
324 } => {
325 collect_target_cte_references_from_nested_query(left, targets, shadowed, references);
326 collect_target_cte_references_from_nested_query(right, targets, shadowed, references);
327 for subquery in subqueries {
328 collect_target_cte_references_from_nested_query(
329 subquery, targets, shadowed, references,
330 );
331 }
332 }
333 RelationalPlan::Values { subqueries, .. } => {
334 for subquery in subqueries {
335 collect_target_cte_references_from_nested_query(
336 subquery, targets, shadowed, references,
337 );
338 }
339 }
340 }
341}
342
343fn collect_target_cte_references_from_source(
344 source: &SourcePlan,
345 targets: &BTreeSet<String>,
346 shadowed: &BTreeSet<String>,
347 references: &mut BTreeSet<String>,
348) {
349 match source {
350 SourcePlan::Table { name, .. } => {
351 if let Some(name) = cte_reference_name(name)
352 .filter(|name| targets.contains(name) && !shadowed.contains(name))
353 {
354 references.insert(name);
355 }
356 }
357 SourcePlan::Join { left, right, .. } => {
358 collect_target_cte_references_from_source(left, targets, shadowed, references);
359 collect_target_cte_references_from_source(right, targets, shadowed, references);
360 }
361 SourcePlan::Subquery { body, .. } => {
362 collect_target_cte_references_from_nested_query(body, targets, shadowed, references);
363 }
364 SourcePlan::Values { .. }
365 | SourcePlan::Function { .. }
366 | SourcePlan::FunctionGroup { .. } => {}
367 }
368}
369
370fn collect_target_cte_references_from_command_root(
371 command: &CommandPlan,
372 targets: &BTreeSet<String>,
373 shadowed: &BTreeSet<String>,
374 references: &mut BTreeSet<String>,
375) {
376 for query in command.query_inputs() {
377 collect_target_cte_references_from_nested_query(query, targets, shadowed, references);
378 }
379 if let Some(source) = command.source_input() {
380 collect_target_cte_references_from_source(source, targets, shadowed, references);
381 }
382}
383
384fn collect_target_cte_references_from_body(
385 body: &CtePlanBody,
386 targets: &BTreeSet<String>,
387 shadowed: &BTreeSet<String>,
388 references: &mut BTreeSet<String>,
389) {
390 let CtePlanBody::Command(command) = body else {
391 if let CtePlanBody::Query(query) = body {
392 collect_target_cte_references_from_nested_query(query, targets, shadowed, references);
393 }
394 return;
395 };
396 let locals = command
397 .ctes()
398 .iter()
399 .map(|cte| cte.name.clone())
400 .collect::<BTreeSet<_>>();
401 let mut reachable = command
402 .ctes()
403 .iter()
404 .filter(|cte| cte.body.modifies_data())
405 .map(|cte| cte.name.clone())
406 .collect::<BTreeSet<_>>();
407 collect_target_cte_references_from_command_root(
408 command,
409 &locals,
410 &BTreeSet::new(),
411 &mut reachable,
412 );
413 let mut expanded = BTreeSet::new();
414 loop {
415 let pending = command
416 .ctes()
417 .iter()
418 .enumerate()
419 .filter(|(_, cte)| reachable.contains(&cte.name) && !expanded.contains(&cte.name))
420 .collect::<Vec<_>>();
421 if pending.is_empty() {
422 break;
423 }
424 for (index, cte) in pending {
425 expanded.insert(cte.name.clone());
426 let visible = if cte.recursive {
427 locals.clone()
428 } else {
429 command.ctes()[..index]
430 .iter()
431 .map(|cte| cte.name.clone())
432 .collect()
433 };
434 collect_target_cte_references_from_body(
435 &cte.body,
436 &visible,
437 &BTreeSet::new(),
438 &mut reachable,
439 );
440 }
441 }
442 let mut root_shadowed = shadowed.clone();
443 root_shadowed.extend(locals.iter().cloned());
444 collect_target_cte_references_from_command_root(command, targets, &root_shadowed, references);
445 let mut preceding = shadowed.clone();
446 for cte in command.ctes() {
447 if reachable.contains(&cte.name) {
448 let definition = if cte.recursive {
449 shadowed.union(&locals).cloned().collect()
450 } else {
451 preceding.clone()
452 };
453 collect_target_cte_references_from_body(&cte.body, targets, &definition, references);
454 }
455 preceding.insert(cte.name.clone());
456 }
457}
458
459fn collect_target_cte_references_from_nested_query(
460 plan: &QueryPlan,
461 targets: &BTreeSet<String>,
462 shadowed: &BTreeSet<String>,
463 references: &mut BTreeSet<String>,
464) {
465 let local_reachable = reachable_plan_cte_names(plan);
466 let mut root_shadowed = shadowed.clone();
467 root_shadowed.extend(
468 plan.ctes
469 .iter()
470 .map(|cte| cte.name.clone())
471 .filter(|name| targets.contains(name)),
472 );
473 collect_target_cte_references_from_root(&plan.root, targets, &root_shadowed, references);
474 let recursive_scope = plan.ctes.iter().any(|cte| cte.recursive).then(|| {
475 plan.ctes
476 .iter()
477 .map(|cte| cte.name.clone())
478 .collect::<BTreeSet<_>>()
479 });
480 let mut preceding = BTreeSet::new();
481 for cte in &plan.ctes {
482 if local_reachable.contains(&cte.name) {
483 let mut definition_shadowed = shadowed.clone();
484 if let Some(recursive_scope) = recursive_scope.as_ref() {
485 definition_shadowed.extend(
486 recursive_scope
487 .iter()
488 .filter(|name| targets.contains(*name))
489 .cloned(),
490 );
491 } else {
492 definition_shadowed.extend(
493 preceding
494 .iter()
495 .filter(|name| targets.contains(*name))
496 .cloned(),
497 );
498 }
499 collect_target_cte_references_from_body(
500 &cte.body,
501 targets,
502 &definition_shadowed,
503 references,
504 );
505 }
506 preceding.insert(cte.name.clone());
507 }
508}