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_references_own_name(cte: &CtePlan) -> bool {
70 let targets = BTreeSet::from([cte.name.clone()]);
71 let mut references = BTreeSet::new();
72 collect_target_cte_references_from_body(&cte.body, &targets, &BTreeSet::new(), &mut references);
73 references.contains(&cte.name)
74}
75
76pub fn ordered_plan_ctes(plan: &QueryPlan) -> Result<Vec<&CtePlan>, SQLError> {
77 ordered_cte_plans(&plan.ctes)
78}
79
80pub fn ordered_cte_plans(ctes: &[CtePlan]) -> Result<Vec<&CtePlan>, SQLError> {
81 order_cte_plans(ctes.iter().collect())
82}
83
84pub fn order_cte_plans(plans: Vec<&CtePlan>) -> Result<Vec<&CtePlan>, SQLError> {
85 if !plans.iter().any(|cte| cte.recursive) {
86 return Ok(plans);
87 }
88 let targets = plans
89 .iter()
90 .map(|cte| cte.name.clone())
91 .collect::<BTreeSet<_>>();
92 let dependencies = plans
93 .iter()
94 .map(|cte| {
95 let mut references = BTreeSet::new();
96 collect_target_cte_references_from_body(
97 &cte.body,
98 &targets,
99 &BTreeSet::new(),
100 &mut references,
101 );
102 references.remove(&cte.name);
103 references
104 })
105 .collect::<Vec<_>>();
106 let mut emitted = BTreeSet::new();
107 let mut ordered = Vec::with_capacity(plans.len());
108 let mut remaining = (0..plans.len()).collect::<BTreeSet<_>>();
109 while !remaining.is_empty() {
110 let ready = remaining
111 .iter()
112 .copied()
113 .find(|index| dependencies[*index].is_subset(&emitted));
114 let Some(index) = ready else {
115 return Err(SQLError::Routine {
116 sqlstate: "0A000".into(),
117 message: "mutual recursion between WITH items is not implemented".into(),
118 });
119 };
120 remaining.remove(&index);
121 emitted.insert(plans[index].name.clone());
122 ordered.push(plans[index]);
123 }
124 Ok(ordered)
125}
126
127pub fn single_reference_plan_cte_names(plan: &QueryPlan) -> BTreeSet<String> {
129 let targets = plan
130 .ctes
131 .iter()
132 .map(|cte| cte.name.clone())
133 .collect::<BTreeSet<_>>();
134 let mut counts = targets
135 .iter()
136 .map(|name| (name.clone(), 0usize))
137 .collect::<BTreeMap<_, _>>();
138 count_plan_cte_references(plan, &targets, &mut counts);
139 counts
140 .into_iter()
141 .filter_map(|(name, count)| (count == 1).then_some(name))
142 .collect()
143}
144
145fn count_plan_cte_references(
146 plan: &QueryPlan,
147 targets: &BTreeSet<String>,
148 counts: &mut BTreeMap<String, usize>,
149) {
150 for cte in &plan.ctes {
151 count_cte_body_references(&cte.body, targets, counts);
152 }
153 count_relational_cte_references(&plan.root, targets, counts);
154}
155
156fn count_cte_body_references(
157 body: &CtePlanBody,
158 targets: &BTreeSet<String>,
159 counts: &mut BTreeMap<String, usize>,
160) {
161 match body {
162 CtePlanBody::Query(query) => count_plan_cte_references(query, targets, counts),
163 CtePlanBody::Command(command) => {
164 for cte in command.ctes() {
165 count_cte_body_references(&cte.body, targets, counts);
166 }
167 for query in command.query_inputs() {
168 count_plan_cte_references(query, targets, counts);
169 }
170 if let Some(source) = command.source_input() {
171 count_source_cte_references(source, targets, counts);
172 }
173 }
174 }
175}
176
177fn count_relational_cte_references(
178 plan: &RelationalPlan,
179 targets: &BTreeSet<String>,
180 counts: &mut BTreeMap<String, usize>,
181) {
182 match plan {
183 RelationalPlan::QueryBlock(block) => {
184 if let Some(source) = &block.from {
185 count_source_cte_references(source, targets, counts);
186 }
187 for subquery in &block.subqueries {
188 count_plan_cte_references(subquery, targets, counts);
189 }
190 }
191 RelationalPlan::SetOp {
192 left,
193 right,
194 subqueries,
195 ..
196 } => {
197 count_plan_cte_references(left, targets, counts);
198 count_plan_cte_references(right, targets, counts);
199 for subquery in subqueries {
200 count_plan_cte_references(subquery, targets, counts);
201 }
202 }
203 RelationalPlan::Values { subqueries, .. } => {
204 for subquery in subqueries {
205 count_plan_cte_references(subquery, targets, counts);
206 }
207 }
208 }
209}
210
211fn count_source_cte_references(
212 source: &SourcePlan,
213 targets: &BTreeSet<String>,
214 counts: &mut BTreeMap<String, usize>,
215) {
216 match source {
217 SourcePlan::Table { name, .. } => {
218 if let Some(name) = cte_reference_name(name).filter(|name| targets.contains(name)) {
219 *counts.entry(name).or_default() += 1;
220 }
221 }
222 SourcePlan::Join { left, right, .. } => {
223 count_source_cte_references(left, targets, counts);
224 count_source_cte_references(right, targets, counts);
225 }
226 SourcePlan::Subquery { body, .. } => {
227 count_plan_cte_references(body, targets, counts);
228 }
229 SourcePlan::Values { .. }
230 | SourcePlan::Function { .. }
231 | SourcePlan::FunctionGroup { .. } => {}
232 }
233}
234
235fn collect_target_cte_references_from_root(
236 root: &RelationalPlan,
237 targets: &BTreeSet<String>,
238 shadowed: &BTreeSet<String>,
239 references: &mut BTreeSet<String>,
240) {
241 match root {
242 RelationalPlan::QueryBlock(block) => {
243 if let Some(source) = &block.from {
244 collect_target_cte_references_from_source(source, targets, shadowed, references);
245 }
246 for subquery in &block.subqueries {
247 collect_target_cte_references_from_nested_query(
248 subquery, targets, shadowed, references,
249 );
250 }
251 }
252 RelationalPlan::SetOp {
253 left,
254 right,
255 subqueries,
256 ..
257 } => {
258 collect_target_cte_references_from_nested_query(left, targets, shadowed, references);
259 collect_target_cte_references_from_nested_query(right, targets, shadowed, references);
260 for subquery in subqueries {
261 collect_target_cte_references_from_nested_query(
262 subquery, targets, shadowed, references,
263 );
264 }
265 }
266 RelationalPlan::Values { subqueries, .. } => {
267 for subquery in subqueries {
268 collect_target_cte_references_from_nested_query(
269 subquery, targets, shadowed, references,
270 );
271 }
272 }
273 }
274}
275
276fn collect_target_cte_references_from_source(
277 source: &SourcePlan,
278 targets: &BTreeSet<String>,
279 shadowed: &BTreeSet<String>,
280 references: &mut BTreeSet<String>,
281) {
282 match source {
283 SourcePlan::Table { name, .. } => {
284 if let Some(name) = cte_reference_name(name)
285 .filter(|name| targets.contains(name) && !shadowed.contains(name))
286 {
287 references.insert(name);
288 }
289 }
290 SourcePlan::Join { left, right, .. } => {
291 collect_target_cte_references_from_source(left, targets, shadowed, references);
292 collect_target_cte_references_from_source(right, targets, shadowed, references);
293 }
294 SourcePlan::Subquery { body, .. } => {
295 collect_target_cte_references_from_nested_query(body, targets, shadowed, references);
296 }
297 SourcePlan::Values { .. }
298 | SourcePlan::Function { .. }
299 | SourcePlan::FunctionGroup { .. } => {}
300 }
301}
302
303fn collect_target_cte_references_from_command_root(
304 command: &CommandPlan,
305 targets: &BTreeSet<String>,
306 shadowed: &BTreeSet<String>,
307 references: &mut BTreeSet<String>,
308) {
309 for query in command.query_inputs() {
310 collect_target_cte_references_from_nested_query(query, targets, shadowed, references);
311 }
312 if let Some(source) = command.source_input() {
313 collect_target_cte_references_from_source(source, targets, shadowed, references);
314 }
315}
316
317fn collect_target_cte_references_from_body(
318 body: &CtePlanBody,
319 targets: &BTreeSet<String>,
320 shadowed: &BTreeSet<String>,
321 references: &mut BTreeSet<String>,
322) {
323 let CtePlanBody::Command(command) = body else {
324 if let CtePlanBody::Query(query) = body {
325 collect_target_cte_references_from_nested_query(query, targets, shadowed, references);
326 }
327 return;
328 };
329 let locals = command
330 .ctes()
331 .iter()
332 .map(|cte| cte.name.clone())
333 .collect::<BTreeSet<_>>();
334 let mut reachable = command
335 .ctes()
336 .iter()
337 .filter(|cte| cte.body.modifies_data())
338 .map(|cte| cte.name.clone())
339 .collect::<BTreeSet<_>>();
340 collect_target_cte_references_from_command_root(
341 command,
342 &locals,
343 &BTreeSet::new(),
344 &mut reachable,
345 );
346 let mut expanded = BTreeSet::new();
347 loop {
348 let pending = command
349 .ctes()
350 .iter()
351 .enumerate()
352 .filter(|(_, cte)| reachable.contains(&cte.name) && !expanded.contains(&cte.name))
353 .collect::<Vec<_>>();
354 if pending.is_empty() {
355 break;
356 }
357 for (index, cte) in pending {
358 expanded.insert(cte.name.clone());
359 let visible = if cte.recursive {
360 locals.clone()
361 } else {
362 command.ctes()[..index]
363 .iter()
364 .map(|cte| cte.name.clone())
365 .collect()
366 };
367 collect_target_cte_references_from_body(
368 &cte.body,
369 &visible,
370 &BTreeSet::new(),
371 &mut reachable,
372 );
373 }
374 }
375 let mut root_shadowed = shadowed.clone();
376 root_shadowed.extend(locals.iter().cloned());
377 collect_target_cte_references_from_command_root(command, targets, &root_shadowed, references);
378 let mut preceding = shadowed.clone();
379 for cte in command.ctes() {
380 if reachable.contains(&cte.name) {
381 let definition = if cte.recursive {
382 shadowed.union(&locals).cloned().collect()
383 } else {
384 preceding.clone()
385 };
386 collect_target_cte_references_from_body(&cte.body, targets, &definition, references);
387 }
388 preceding.insert(cte.name.clone());
389 }
390}
391
392fn collect_target_cte_references_from_nested_query(
393 plan: &QueryPlan,
394 targets: &BTreeSet<String>,
395 shadowed: &BTreeSet<String>,
396 references: &mut BTreeSet<String>,
397) {
398 let local_reachable = reachable_plan_cte_names(plan);
399 let mut root_shadowed = shadowed.clone();
400 root_shadowed.extend(
401 plan.ctes
402 .iter()
403 .map(|cte| cte.name.clone())
404 .filter(|name| targets.contains(name)),
405 );
406 collect_target_cte_references_from_root(&plan.root, targets, &root_shadowed, references);
407 let recursive_scope = plan.ctes.iter().any(|cte| cte.recursive).then(|| {
408 plan.ctes
409 .iter()
410 .map(|cte| cte.name.clone())
411 .collect::<BTreeSet<_>>()
412 });
413 let mut preceding = BTreeSet::new();
414 for cte in &plan.ctes {
415 if local_reachable.contains(&cte.name) {
416 let mut definition_shadowed = shadowed.clone();
417 if let Some(recursive_scope) = recursive_scope.as_ref() {
418 definition_shadowed.extend(
419 recursive_scope
420 .iter()
421 .filter(|name| targets.contains(*name))
422 .cloned(),
423 );
424 } else {
425 definition_shadowed.extend(
426 preceding
427 .iter()
428 .filter(|name| targets.contains(*name))
429 .cloned(),
430 );
431 }
432 collect_target_cte_references_from_body(
433 &cte.body,
434 targets,
435 &definition_shadowed,
436 references,
437 );
438 }
439 preceding.insert(cte.name.clone());
440 }
441}