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
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
127/// 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.
128pub 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}