Skip to main content

uqa_sql/semantics/
ctes.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! CTE scope, reachability, and dependency ordering.
8
9use crate::plan::{CommandPlan, CtePlan, CtePlanBody, QueryPlan, RelationalPlan, SourcePlan};
10use crate::SQLError;
11use std::collections::{BTreeMap, BTreeSet};
12
13/// Decode a source reference into a CTE identifier. Qualified relations never resolve to CTEs.
14pub 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
69/// The WITH items of `ctes` that the definition of the item at `index` names: the items defined before it, or any item of a recursive list.
70pub 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
84/// The WITH items of `plan` that its primary query names outside the items' own definitions.
85pub 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
103/// The WITH items of a data-modifying statement that its own clauses name: the queries and the source relation it evaluates in the scope of `ctes`.
104pub 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
194/// Return reachable CTEs with exactly one syntactic reference in the owning query tree. Counting references outside their lexical visibility can only make this set more conservative, never cause a multiply referenced CTE to be streamed as a single-consumer input.
195pub 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}