mz_transform/
normalize_lets.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Normalize the structure of `Let` and `LetRec` operators in expressions.
11//!
12//! Normalization happens in the context of "scopes", corresponding to
13//! 1. the expression's root and 2. each instance of a `LetRec` AST node.
14//!
15//! Within each scope,
16//! 1. Each expression is normalized to have all `Let` nodes at the root
17//! of the expression, in order of identifier.
18//! 2. Each expression assigns a contiguous block of identifiers.
19//!
20//! The transform may remove some `Let` and `Get` operators, and does not
21//! introduce any new operators.
22//!
23//! The module also publishes the function `renumber_bindings` which can
24//! be used to renumber bindings in an expression starting from a provided
25//! `IdGen`, which is used to prepare distinct expressions for inlining.
26
27use mz_expr::{MirRelationExpr, visit::Visit};
28use mz_ore::assert_none;
29use mz_ore::{id_gen::IdGen, stack::RecursionLimitError};
30use mz_repr::optimize::OptimizerFeatures;
31
32use crate::{TransformCtx, catch_unwind_optimize};
33
34pub use renumbering::renumber_bindings;
35
36/// Normalize `Let` and `LetRec` structure.
37pub fn normalize_lets(
38    expr: &mut MirRelationExpr,
39    features: &OptimizerFeatures,
40) -> Result<(), crate::TransformError> {
41    catch_unwind_optimize(|| NormalizeLets::new(false).action(expr, features))
42}
43
44/// Install replace certain `Get` operators with their `Let` value.
45#[derive(Debug)]
46pub struct NormalizeLets {
47    /// If `true`, inline MFPs around a Get.
48    ///
49    /// We want this value to be true for the NormalizeLets call that comes right
50    /// before [crate::join_implementation::JoinImplementation] runs because
51    /// - JoinImplementation cannot lift MFPs through a Let.
52    /// - JoinImplementation can't extract FilterCharacteristics through a Let.
53    ///
54    /// Generally, though, we prefer to be more conservative in our inlining in
55    /// order to be able to better detect CSEs.
56    pub inline_mfp: bool,
57}
58
59impl NormalizeLets {
60    /// Construct a new [`NormalizeLets`] instance with the given `inline_mfp`.
61    pub fn new(inline_mfp: bool) -> NormalizeLets {
62        NormalizeLets { inline_mfp }
63    }
64}
65
66impl crate::Transform for NormalizeLets {
67    fn name(&self) -> &'static str {
68        "NormalizeLets"
69    }
70
71    #[mz_ore::instrument(
72        target = "optimizer",
73        level = "debug",
74        fields(path.segment = "normalize_lets")
75    )]
76    fn actually_perform_transform(
77        &self,
78        relation: &mut MirRelationExpr,
79        ctx: &mut TransformCtx,
80    ) -> Result<(), crate::TransformError> {
81        let result = self.action(relation, ctx.features);
82        mz_repr::explain::trace_plan(&*relation);
83        result
84    }
85}
86
87impl NormalizeLets {
88    /// Normalize `Let` and `LetRec` bindings in `relation`.
89    ///
90    /// Mechanically, `action` first renumbers all bindings, erroring if any shadowing is encountered.
91    /// It then promotes all `Let` and `LetRec` expressions to the roots of their expressions, fusing
92    /// `Let` bindings into containing `LetRec` bindings, but leaving stacked `LetRec` bindings unfused to each
93    /// other (for reasons of correctness). It then considers potential inlining in each `LetRec` scope.
94    /// Lastly, it refreshes the types of each `Get` operator, erroring if any scalar types have changed
95    /// but updating nullability and keys.
96    ///
97    /// We then perform a final renumbering.
98    pub fn action(
99        &self,
100        relation: &mut MirRelationExpr,
101        features: &OptimizerFeatures,
102    ) -> Result<(), crate::TransformError> {
103        // Record whether the relation was initially recursive, to confirm that we do not introduce
104        // recursion to a non-recursive expression.
105        let was_recursive = relation.is_recursive();
106
107        // Renumber all bindings to ensure that identifier order matches binding order.
108        // In particular, as we use `BTreeMap` for binding order, we want to ensure that
109        // 1. Bindings within a `LetRec` are assigned increasing identifiers, and
110        // 2. Bindings across `LetRec`s are assigned identifiers in "visibility order", corresponding to an
111        // in-order traversal.
112        // TODO: More can and perhaps should be said about "visibility order" and how let promotion is correct.
113        renumbering::renumber_bindings(relation, &mut IdGen::default())?;
114
115        // Promote all `Let` and `LetRec` AST nodes to the roots.
116        // After this, all non-`LetRec` nodes contain no further `Let` or `LetRec` nodes,
117        // placing all `LetRec` nodes around the root, if not always in a single AST node.
118        let_motion::promote_let_rec(relation);
119        let_motion::assert_no_lets(relation);
120        let_motion::assert_letrec_major(relation);
121
122        // Inlining may violate letrec-major form.
123        inlining::inline_lets(relation, self.inline_mfp)?;
124
125        // Return to letrec-major form to refresh types.
126        let_motion::promote_let_rec(relation);
127        support::refresh_types(relation, features)?;
128
129        // Renumber bindings for good measure.
130        // Ideally we could skip when `action` is a no-op, but hard to thread that through at the moment.
131        renumbering::renumber_bindings(relation, &mut IdGen::default())?;
132
133        // A final bottom-up traversal to normalize the shape of nested LetRec blocks
134        relation.try_visit_mut_post(&mut |relation| -> Result<(), RecursionLimitError> {
135            // Move a non-recursive suffix of bindings from the end of the LetRec
136            // to the LetRec body.
137            // This is unsafe when applied to expressions which contain `ArrangeBy`,
138            // as if the extracted suffixes reference arrangements they will not be
139            // able to access those arrangements from outside the `LetRec` scope.
140            // It happens to work at the moment, so we don't touch it but should fix.
141            let bindings = let_motion::harvest_nonrec_suffix(relation)?;
142            if let MirRelationExpr::LetRec {
143                ids: _,
144                values: _,
145                limits: _,
146                body,
147            } = relation
148            {
149                for (id, value) in bindings.into_iter().rev() {
150                    **body = MirRelationExpr::Let {
151                        id,
152                        value: Box::new(value),
153                        body: Box::new(body.take_dangerous()),
154                    };
155                }
156            } else {
157                for (id, value) in bindings.into_iter().rev() {
158                    *relation = MirRelationExpr::Let {
159                        id,
160                        value: Box::new(value),
161                        body: Box::new(relation.take_dangerous()),
162                    };
163                }
164            }
165
166            // Extract `Let` prefixes from `LetRec`, to reveal their non-recursive nature.
167            // This assists with hoisting e.g. arrangements out of `LetRec` blocks, a thing
168            // we don't promise to do, but it can be helpful to do. This also exposes more
169            // AST nodes to non-`LetRec` analyses, which don't always have parity with `LetRec`.
170            let bindings = let_motion::harvest_non_recursive(relation);
171            for (id, (value, max_iter)) in bindings.into_iter().rev() {
172                assert_none!(max_iter);
173                *relation = MirRelationExpr::Let {
174                    id,
175                    value: Box::new(value),
176                    body: Box::new(relation.take_dangerous()),
177                };
178            }
179
180            Ok(())
181        })?;
182
183        if !was_recursive && relation.is_recursive() {
184            Err(crate::TransformError::Internal(
185                "NormalizeLets introduced LetRec to a LetRec-free expression".to_string(),
186            ))?;
187        }
188
189        Ok(())
190    }
191}
192
193// Support methods that are unlikely to be useful to other modules.
194mod support {
195
196    use std::collections::BTreeMap;
197
198    use itertools::Itertools;
199
200    use mz_expr::{Id, LetRecLimit, LocalId, MirRelationExpr};
201    use mz_repr::optimize::OptimizerFeatures;
202
203    pub(super) fn replace_bindings_from_map(
204        map: BTreeMap<LocalId, (MirRelationExpr, Option<LetRecLimit>)>,
205        ids: &mut Vec<LocalId>,
206        values: &mut Vec<MirRelationExpr>,
207        limits: &mut Vec<Option<LetRecLimit>>,
208    ) {
209        let (new_ids, new_values, new_limits) = map_to_3vecs(map);
210        *ids = new_ids;
211        *values = new_values;
212        *limits = new_limits;
213    }
214
215    pub(super) fn map_to_3vecs(
216        map: BTreeMap<LocalId, (MirRelationExpr, Option<LetRecLimit>)>,
217    ) -> (Vec<LocalId>, Vec<MirRelationExpr>, Vec<Option<LetRecLimit>>) {
218        let (new_ids, new_values_and_limits): (Vec<_>, Vec<_>) = map.into_iter().unzip();
219        let (new_values, new_limits) = new_values_and_limits.into_iter().unzip();
220        (new_ids, new_values, new_limits)
221    }
222
223    /// Logic mapped across each use of a `LocalId`.
224    pub(super) fn for_local_id<F>(expr: &MirRelationExpr, mut logic: F)
225    where
226        F: FnMut(LocalId),
227    {
228        expr.visit_pre(|expr| {
229            if let MirRelationExpr::Get {
230                id: Id::Local(i), ..
231            } = expr
232            {
233                logic(*i);
234            }
235        });
236    }
237
238    /// Populates `counts` with the number of uses of each local identifier in `expr`.
239    pub(super) fn count_local_id_uses(
240        expr: &MirRelationExpr,
241        counts: &mut std::collections::BTreeMap<LocalId, usize>,
242    ) {
243        for_local_id(expr, |i| *counts.entry(i).or_insert(0) += 1)
244    }
245
246    /// Visit `LetRec` stages and determine and update type information for `Get` nodes.
247    ///
248    /// This method errors if the scalar type information has changed (number of columns, or types).
249    /// It only refreshes the nullability and unique key information. As this information can regress,
250    /// we do not error if the type weakens, even though that may be something we want to look into.
251    ///
252    /// The method relies on the `analysis::{UniqueKeys, SqlRelationType}` analyses to improve its type
253    /// information for `LetRec` stages.
254    pub(super) fn refresh_types(
255        expr: &mut MirRelationExpr,
256        features: &OptimizerFeatures,
257    ) -> Result<(), crate::TransformError> {
258        // Assemble type information once for the whole expression.
259        use crate::analysis::{DerivedBuilder, SqlRelationType, UniqueKeys};
260        let mut builder = DerivedBuilder::new(features);
261        builder.require(SqlRelationType);
262        builder.require(UniqueKeys);
263        let derived = builder.visit(expr);
264        let derived_view = derived.as_view();
265
266        // Collect id -> type mappings.
267        let mut types = BTreeMap::new();
268        let mut todo = vec![(&*expr, derived_view)];
269        while let Some((expr, view)) = todo.pop() {
270            let ids = match expr {
271                MirRelationExpr::Let { id, .. } => std::slice::from_ref(id),
272                MirRelationExpr::LetRec { ids, .. } => ids,
273                _ => &[],
274            };
275            if !ids.is_empty() {
276                // The `skip(1)` skips the `body` child, and is followed by binding children.
277                for (id, view) in ids.iter().rev().zip_eq(view.children_rev().skip(1)) {
278                    let cols = view
279                        .value::<SqlRelationType>()
280                        .expect("SqlRelationType required")
281                        .clone()
282                        .expect("Expression not well typed");
283                    let keys = view
284                        .value::<UniqueKeys>()
285                        .expect("UniqueKeys required")
286                        .clone();
287                    types.insert(*id, mz_repr::SqlRelationType::new(cols).with_keys(keys));
288                }
289            }
290            todo.extend(expr.children().rev().zip_eq(view.children_rev()));
291        }
292
293        // Install the new types in each `Get`.
294        let mut todo = vec![&mut *expr];
295        while let Some(expr) = todo.pop() {
296            if let MirRelationExpr::Get {
297                id: Id::Local(i),
298                typ,
299                ..
300            } = expr
301            {
302                if let Some(new_type) = types.get(i) {
303                    // Assert that the column length has not changed.
304                    if !new_type.column_types.len() == typ.column_types.len() {
305                        Err(crate::TransformError::Internal(format!(
306                            "column lengths do not match: {:?} v {:?}",
307                            new_type.column_types, typ.column_types
308                        )))?;
309                    }
310                    // Assert that the column types have not changed.
311                    if !new_type
312                        .column_types
313                        .iter()
314                        .zip_eq(typ.column_types.iter())
315                        .all(|(t1, t2)| t1.scalar_type.base_eq(&t2.scalar_type))
316                    {
317                        Err(crate::TransformError::Internal(format!(
318                            "scalar types do not match: {:?} v {:?}",
319                            new_type.column_types, typ.column_types
320                        )))?;
321                    }
322
323                    typ.clone_from(new_type);
324                } else {
325                    panic!("Type not found for: {:?}", i);
326                }
327            }
328            todo.extend(expr.children_mut());
329        }
330        Ok(())
331    }
332}
333
334mod let_motion {
335
336    use std::collections::{BTreeMap, BTreeSet};
337
338    use itertools::Itertools;
339    use mz_expr::{LetRecLimit, LocalId, MirRelationExpr};
340    use mz_ore::stack::RecursionLimitError;
341
342    use crate::normalize_lets::support::replace_bindings_from_map;
343
344    /// Promotes all `Let` and `LetRec` nodes to the roots of their expressions.
345    ///
346    /// We cannot (without further reasoning) fuse stacked `LetRec` stages, and instead we just promote
347    /// `LetRec` to the roots of their expressions (e.g. as children of another `LetRec` stage).
348    pub(crate) fn promote_let_rec(expr: &mut MirRelationExpr) {
349        // First, promote all `LetRec` nodes above all other nodes.
350        let mut worklist = vec![&mut *expr];
351        while let Some(mut expr) = worklist.pop() {
352            hoist_bindings(expr);
353            while let MirRelationExpr::LetRec { values, body, .. } = expr {
354                worklist.extend(values.iter_mut().rev());
355                expr = body;
356            }
357        }
358
359        // Harvest any potential `Let` nodes, via a post-order traversal.
360        post_order_harvest_lets(expr);
361    }
362
363    /// A stand in for the types of bindings we might encounter.
364    ///
365    /// As we dissolve various `Let` and `LetRec` expressions, a `Binding` will carry
366    /// the relevant information as we hoist it to the root of the expression.
367    enum Binding {
368        // Binding resulting from a `Let` expression.
369        Let(LocalId, MirRelationExpr),
370        // Bindings resulting from a `LetRec` expression.
371        LetRec(Vec<(LocalId, MirRelationExpr, Option<LetRecLimit>)>),
372    }
373
374    /// Hoist all exposed bindings to the root of the expression.
375    ///
376    /// A binding is "exposed" if the path from the root does not cross a LetRec binding.
377    /// After the call, the expression should be a linear sequence of bindings, where each
378    /// `Let` binding is of a let-free expression. There may be `LetRec` expressions in the
379    /// sequence, and their bindings will have hoisted bindings to their root, but not out
380    /// of the binding.
381    fn hoist_bindings(expr: &mut MirRelationExpr) {
382        // Bindings we have extracted but not fully processed.
383        let mut worklist = Vec::new();
384        // Bindings we have extracted and then fully processed.
385        let mut finished = Vec::new();
386
387        extract_bindings(expr, &mut worklist);
388        while let Some(mut bind) = worklist.pop() {
389            match &mut bind {
390                Binding::Let(_id, value) => {
391                    extract_bindings(value, &mut worklist);
392                }
393                Binding::LetRec(_binds) => {
394                    // nothing to do here; we cannot hoist letrec bindings and refine
395                    // them in an outer loop.
396                }
397            }
398            finished.push(bind);
399        }
400
401        // The worklist is empty and finished should contain only LetRec bindings and Let
402        // bindings with let-free expressions bound. We need to re-assemble them now in
403        // the correct order. The identifiers are "sequential", so we should be able to
404        // sort by them, with some care.
405
406        // We only extract non-empty letrec bindings, so it is safe to peek at the first.
407        finished.sort_by_key(|b| match b {
408            Binding::Let(id, _) => *id,
409            Binding::LetRec(binds) => binds[0].0,
410        });
411
412        // To match historical behavior we fuse let bindings into adjacent letrec bindings.
413        // We could alternately make each a singleton letrec binding (just, non-recursive).
414        // We don't yet have a strong opinion on which is most helpful and least harmful.
415        // In the absence of any letrec bindings, we form one to house the let bindings.
416        let mut ids = Vec::new();
417        let mut values = Vec::new();
418        let mut limits = Vec::new();
419        let mut compact = Vec::new();
420        for bind in finished {
421            match bind {
422                Binding::Let(id, value) => {
423                    ids.push(id);
424                    values.push(value);
425                    limits.push(None);
426                }
427                Binding::LetRec(binds) => {
428                    for (id, value, limit) in binds {
429                        ids.push(id);
430                        values.push(value);
431                        limits.push(limit);
432                    }
433                    compact.push((ids, values, limits));
434                    ids = Vec::new();
435                    values = Vec::new();
436                    limits = Vec::new();
437                }
438            }
439        }
440
441        // Remaining bindings can either be fused to the prior letrec, or put in their own.
442        if let Some((last_ids, last_vals, last_lims)) = compact.last_mut() {
443            last_ids.extend(ids);
444            last_vals.extend(values);
445            last_lims.extend(limits);
446        } else if !ids.is_empty() {
447            compact.push((ids, values, limits));
448        }
449
450        while let Some((ids, values, limits)) = compact.pop() {
451            *expr = MirRelationExpr::LetRec {
452                ids,
453                values,
454                limits,
455                body: Box::new(expr.take_dangerous()),
456            };
457        }
458    }
459
460    /// Extracts exposed bindings into `bindings`.
461    ///
462    /// After this call `expr` will contain no let or letrec bindings, though the bindings
463    /// it introduces to `bindings` may themselves contain such bindings (and they should
464    /// be further processed if the goal is to maximally extract let bindings).
465    fn extract_bindings(expr: &mut MirRelationExpr, bindings: &mut Vec<Binding>) {
466        let mut todo = vec![expr];
467        while let Some(expr) = todo.pop() {
468            match expr {
469                MirRelationExpr::Let { id, value, body } => {
470                    bindings.push(Binding::Let(*id, value.take_dangerous()));
471                    *expr = body.take_dangerous();
472                    todo.push(expr);
473                }
474                MirRelationExpr::LetRec {
475                    ids,
476                    values,
477                    limits,
478                    body,
479                } => {
480                    use itertools::Itertools;
481                    let binds: Vec<_> = ids
482                        .drain(..)
483                        .zip_eq(values.drain(..))
484                        .zip_eq(limits.drain(..))
485                        .map(|((i, v), l)| (i, v, l))
486                        .collect();
487                    if !binds.is_empty() {
488                        bindings.push(Binding::LetRec(binds));
489                    }
490                    *expr = body.take_dangerous();
491                    todo.push(expr);
492                }
493                _ => {
494                    todo.extend(expr.children_mut());
495                }
496            }
497        }
498    }
499
500    /// Performs a post-order traversal of the `LetRec` nodes at the root of an expression.
501    ///
502    /// The traversal is only of the `LetRec` nodes, for which fear of stack exhaustion is nominal.
503    fn post_order_harvest_lets(expr: &mut MirRelationExpr) {
504        if let MirRelationExpr::LetRec {
505            ids,
506            values,
507            limits,
508            body,
509        } = expr
510        {
511            // Only recursively descend through `LetRec` stages.
512            for value in values.iter_mut() {
513                post_order_harvest_lets(value);
514            }
515
516            let mut bindings = BTreeMap::new();
517            for ((id, mut value), max_iter) in ids
518                .drain(..)
519                .zip_eq(values.drain(..))
520                .zip_eq(limits.drain(..))
521            {
522                bindings.extend(harvest_non_recursive(&mut value));
523                bindings.insert(id, (value, max_iter));
524            }
525            bindings.extend(harvest_non_recursive(body));
526            replace_bindings_from_map(bindings, ids, values, limits);
527        }
528    }
529
530    /// Harvest any safe-to-lift non-recursive bindings from a `LetRec`
531    /// expression.
532    ///
533    /// At the moment, we reason that a binding can be lifted without changing
534    /// the output if both:
535    /// 1. It references no other non-lifted binding bound in `expr`,
536    /// 2. It is referenced by no prior non-lifted binding in `expr`.
537    ///
538    /// The rationale is that (1) ensures that the binding's value does not
539    /// change across iterations, and that (2) ensures that all observations of
540    /// the binding are after it assumes its first value, rather than when it
541    /// could be empty.
542    pub(crate) fn harvest_non_recursive(
543        expr: &mut MirRelationExpr,
544    ) -> BTreeMap<LocalId, (MirRelationExpr, Option<LetRecLimit>)> {
545        if let MirRelationExpr::LetRec {
546            ids,
547            values,
548            limits,
549            body,
550        } = expr
551        {
552            // Bindings to lift.
553            let mut lifted = BTreeMap::<LocalId, (MirRelationExpr, Option<LetRecLimit>)>::new();
554            // Bindings to retain.
555            let mut retained = BTreeMap::<LocalId, (MirRelationExpr, Option<LetRecLimit>)>::new();
556
557            // All remaining LocalIds bound by the enclosing LetRec.
558            let mut id_set = ids.iter().cloned().collect::<BTreeSet<LocalId>>();
559            // All LocalIds referenced up to (including) the current binding.
560            let mut cannot = BTreeSet::<LocalId>::new();
561            // The reference count of the current bindings.
562            let mut refcnt = BTreeMap::<LocalId, usize>::new();
563
564            for ((id, value), max_iter) in ids
565                .drain(..)
566                .zip_eq(values.drain(..))
567                .zip_eq(limits.drain(..))
568            {
569                refcnt.clear();
570                super::support::count_local_id_uses(&value, &mut refcnt);
571
572                // LocalIds that have already been referenced cannot be lifted.
573                cannot.extend(refcnt.keys().cloned());
574
575                // - The first conjunct excludes bindings that have already been
576                //   referenced.
577                // - The second conjunct excludes bindings that reference a
578                //   LocalId that either defined later or is a known retained.
579                if !cannot.contains(&id) && !refcnt.keys().any(|i| id_set.contains(i)) {
580                    lifted.insert(id, (value, None)); // Non-recursive bindings don't need a limit
581                    id_set.remove(&id);
582                } else {
583                    retained.insert(id, (value, max_iter));
584                }
585            }
586
587            replace_bindings_from_map(retained, ids, values, limits);
588            if values.is_empty() {
589                *expr = body.take_dangerous();
590            }
591
592            lifted
593        } else {
594            BTreeMap::new()
595        }
596    }
597
598    /// Harvest any safe-to-lower non-recursive suffix of binding from a
599    /// `LetRec` expression.
600    pub(crate) fn harvest_nonrec_suffix(
601        expr: &mut MirRelationExpr,
602    ) -> Result<BTreeMap<LocalId, MirRelationExpr>, RecursionLimitError> {
603        if let MirRelationExpr::LetRec {
604            ids,
605            values,
606            limits,
607            body,
608        } = expr
609        {
610            // Bindings to lower.
611            let mut lowered = BTreeMap::<LocalId, MirRelationExpr>::new();
612
613            let rec_ids = MirRelationExpr::recursive_ids(ids, values);
614
615            while ids.last().map(|id| !rec_ids.contains(id)).unwrap_or(false) {
616                let id = ids.pop().expect("non-empty ids");
617                let value = values.pop().expect("non-empty values");
618                let _limit = limits.pop().expect("non-empty limits");
619
620                lowered.insert(id, value); // Non-recursive bindings don't need a limit
621            }
622
623            if values.is_empty() {
624                *expr = body.take_dangerous();
625            }
626
627            Ok(lowered)
628        } else {
629            Ok(BTreeMap::new())
630        }
631    }
632
633    pub(crate) fn assert_no_lets(expr: &MirRelationExpr) {
634        expr.visit_pre(|expr| {
635            assert!(!matches!(expr, MirRelationExpr::Let { .. }));
636        });
637    }
638
639    /// Asserts that `expr` in "LetRec-major" form.
640    ///
641    /// This means `expr` is either `LetRec`-free, or a `LetRec` whose values and body are `LetRec`-major.
642    pub(crate) fn assert_letrec_major(expr: &MirRelationExpr) {
643        let mut todo = vec![expr];
644        while let Some(expr) = todo.pop() {
645            match expr {
646                MirRelationExpr::LetRec {
647                    ids: _,
648                    values,
649                    limits: _,
650                    body,
651                } => {
652                    todo.extend(values.iter());
653                    todo.push(body);
654                }
655                _ => {
656                    expr.visit_pre(|expr| {
657                        assert!(!matches!(expr, MirRelationExpr::LetRec { .. }));
658                    });
659                }
660            }
661        }
662    }
663}
664
665mod inlining {
666
667    use std::collections::BTreeMap;
668
669    use itertools::Itertools;
670    use mz_expr::{Id, LetRecLimit, LocalId, MirRelationExpr};
671
672    use crate::normalize_lets::support::replace_bindings_from_map;
673
674    pub(super) fn inline_lets(
675        expr: &mut MirRelationExpr,
676        inline_mfp: bool,
677    ) -> Result<(), crate::TransformError> {
678        let mut worklist = vec![&mut *expr];
679        while let Some(expr) = worklist.pop() {
680            inline_lets_core(expr, inline_mfp)?;
681            // We descend only into `LetRec` nodes, because `promote_let_rec` ensured that all
682            // `LetRec` nodes are clustered near the root. This means that we can get to all the
683            // `LetRec` nodes by just descending into `LetRec` nodes, as there can't be any other
684            // nodes between them.
685            if let MirRelationExpr::LetRec {
686                ids: _,
687                values,
688                limits: _,
689                body,
690            } = expr
691            {
692                worklist.extend(values);
693                worklist.push(body);
694            }
695        }
696        Ok(())
697    }
698
699    /// Considers inlining actions to perform for a sequence of bindings and a
700    /// following body.
701    ///
702    /// A let binding may be inlined only in subsequent bindings or in the body;
703    /// other bindings should not "immediately" observe the binding, as that
704    /// would be a change to the semantics of `LetRec`. For example, it would
705    /// not be correct to replace `C` with `A` in the definition of `B` here:
706    /// ```ignore
707    /// let A = ...;
708    /// let B = A - C;
709    /// let C = A;
710    /// ```
711    /// The explanation is that `B` should always be the difference between the
712    /// current and previous `A`, and that the substitution of `C` would instead
713    /// make it always zero, changing its definition.
714    ///
715    /// Here a let binding is proposed for inlining if any of the following is true:
716    ///  1. It has a single reference across all bindings and the body.
717    ///  2. It is a "sufficiently simple" `Get`, determined in part by the
718    ///     `inline_mfp` argument.
719    ///
720    /// We don't need extra checks for `limits`, because
721    ///  - `limits` is only relevant when a binding is directly used through a back edge (because
722    ///    that is where the rendering puts the `limits` check);
723    ///  - when a binding is directly used through a back edge, it can't be inlined anyway.
724    ///  - Also note that if a `LetRec` completely disappears at the end of `inline_lets_core`, then
725    ///    there was no recursion in it.
726    ///
727    /// The case of `Constant` binding is handled here (as opposed to
728    /// `FoldConstants`) in a somewhat limited manner (see database-issues#5346). Although a
729    /// bit weird, constants should also not be inlined into prior bindings as
730    /// this does change the behavior from one where the collection is initially
731    /// empty to one where it is always the constant.
732    ///
733    /// Having inlined bindings, many of them may now be dead (with no
734    /// transitive references from `body`). These can now be removed. They may
735    /// not be exactly those bindings that were inlineable, as we may not always
736    /// be able to apply inlining due to ordering (we cannot inline a binding
737    /// into one that is not strictly later).
738    pub(super) fn inline_lets_core(
739        expr: &mut MirRelationExpr,
740        inline_mfp: bool,
741    ) -> Result<(), crate::TransformError> {
742        if let MirRelationExpr::LetRec {
743            ids,
744            values,
745            limits,
746            body,
747        } = expr
748        {
749            // Count the number of uses of each local id across all expressions.
750            let mut counts = BTreeMap::new();
751            for value in values.iter() {
752                super::support::count_local_id_uses(value, &mut counts);
753            }
754            super::support::count_local_id_uses(body, &mut counts);
755
756            // Each binding can reach one of three positions on its inlineability:
757            //  1. The binding is used once and is available to be directly taken.
758            //  2. The binding is simple enough that it can just be cloned.
759            //  3. The binding is not available for inlining.
760            let mut inline_offers = BTreeMap::new();
761
762            // Each binding may require the expiration of prior inlining offers.
763            // This occurs when an inlined body references the prior iterate of a binding,
764            // and inlining it would change the meaning to be the current iterate.
765            // Roughly, all inlining offers expire just after the binding of the least
766            // identifier they contain that is greater than the bound identifier itself.
767            let mut expire_offers = BTreeMap::new();
768            let mut expired_offers = Vec::new();
769
770            // For each binding, inline `Get`s and then determine if *it* should be inlined.
771            // It is important that we do the substitution in-order and before reasoning
772            // about the inlineability of each binding, to ensure that our conclusion about
773            // the inlineability of a binding stays put. Specifically,
774            //   1. by going in order no substitution will increase the `Get`-count of an
775            //      identifier beyond one, as all in values with strictly greater identifiers.
776            //   2. by performing the substitution before reasoning, the structure of the value
777            //      as it would be substituted is fixed.
778            for ((id, mut expr), max_iter) in ids
779                .drain(..)
780                .zip_eq(values.drain(..))
781                .zip_eq(limits.drain(..))
782            {
783                // Substitute any appropriate prior let bindings.
784                inline_lets_helper(&mut expr, &mut inline_offers)?;
785
786                // Determine the first `id'` at which any inlining offer must expire.
787                // An inlining offer expires because it references an `id'` that is not yet bound,
788                // indicating a reference to the *prior* iterate of that identifier. Inlining the
789                // expression once `id'` becomes bound would advance the reference to be the
790                // *current* iterate of the identifier.
791                MirRelationExpr::collect_expirations(id, &expr, &mut expire_offers);
792
793                // Gets for `id` only occur in later expressions, so this should still be correct.
794                let num_gets = counts.get(&id).map(|x| *x).unwrap_or(0);
795                // Counts of zero or one lead to substitution; otherwise certain simple structures
796                // are cloned in to `Get` operators, and all others emitted as `Let` bindings.
797                if num_gets == 0 {
798                } else if num_gets == 1 {
799                    inline_offers.insert(id, InlineOffer::Take(Some(expr), max_iter));
800                } else {
801                    let clone_binding = {
802                        let stripped_value = if inline_mfp {
803                            mz_expr::MapFilterProject::extract_non_errors_from_expr(&expr).1
804                        } else {
805                            &expr
806                        };
807                        match stripped_value {
808                            MirRelationExpr::Get { .. } | MirRelationExpr::Constant { .. } => true,
809                            _ => false,
810                        }
811                    };
812
813                    if clone_binding {
814                        inline_offers.insert(id, InlineOffer::Clone(expr, max_iter));
815                    } else {
816                        inline_offers.insert(id, InlineOffer::Unavailable(expr, max_iter));
817                    }
818                }
819
820                // We must now discard any offers that reference `id`, as it is no longer correct
821                // to inline such an offer as it would have access to this iteration's binding of
822                // `id` rather than the prior iteration's binding of `id`.
823                expired_offers.extend(MirRelationExpr::do_expirations(
824                    id,
825                    &mut expire_offers,
826                    &mut inline_offers,
827                ));
828            }
829            // Complete the inlining in `body`.
830            inline_lets_helper(body, &mut inline_offers)?;
831
832            // Re-introduce expired offers for the subsequent logic that expects to see them all.
833            for (id, offer) in expired_offers.into_iter() {
834                inline_offers.insert(id, offer);
835            }
836
837            // We may now be able to discard some of `inline_offer` based on the remaining pattern of `Get` expressions.
838            // Starting from `body` and working backwards, we can activate bindings that are still required because we
839            // observe `Get` expressions referencing them. Any bindings not so identified can be dropped (including any
840            // that may be part of a cycle not reachable from `body`).
841            let mut let_bindings = BTreeMap::new();
842            let mut todo = Vec::new();
843            super::support::for_local_id(body, |id| todo.push(id));
844            while let Some(id) = todo.pop() {
845                if let Some(offer) = inline_offers.remove(&id) {
846                    let (value, max_iter) = match offer {
847                        InlineOffer::Take(value, max_iter) => (
848                            value.ok_or_else(|| {
849                                crate::TransformError::Internal(
850                                    "Needed value already taken".to_string(),
851                                )
852                            })?,
853                            max_iter,
854                        ),
855                        InlineOffer::Clone(value, max_iter) => (value, max_iter),
856                        InlineOffer::Unavailable(value, max_iter) => (value, max_iter),
857                    };
858                    super::support::for_local_id(&value, |id| todo.push(id));
859                    let_bindings.insert(id, (value, max_iter));
860                }
861            }
862
863            // If bindings remain we update the `LetRec`, otherwise we remove it.
864            if !let_bindings.is_empty() {
865                replace_bindings_from_map(let_bindings, ids, values, limits);
866            } else {
867                *expr = body.take_dangerous();
868            }
869        }
870        Ok(())
871    }
872
873    /// Possible states of let binding inlineability.
874    enum InlineOffer {
875        /// There is a unique reference to this value and given the option it should take this expression.
876        Take(Option<MirRelationExpr>, Option<LetRecLimit>),
877        /// Any reference to this value should clone this expression.
878        Clone(MirRelationExpr, Option<LetRecLimit>),
879        /// Any reference to this value should do no inlining of it.
880        Unavailable(MirRelationExpr, Option<LetRecLimit>),
881    }
882
883    /// Substitute `Get{id}` expressions for any proposed expressions.
884    ///
885    /// The proposed expressions can be proposed either to be taken or cloned.
886    fn inline_lets_helper(
887        expr: &mut MirRelationExpr,
888        inline_offer: &mut BTreeMap<LocalId, InlineOffer>,
889    ) -> Result<(), crate::TransformError> {
890        let mut worklist = vec![expr];
891        while let Some(expr) = worklist.pop() {
892            if let MirRelationExpr::Get {
893                id: Id::Local(id), ..
894            } = expr
895            {
896                if let Some(offer) = inline_offer.get_mut(id) {
897                    // It is important that we *not* continue to iterate
898                    // on the contents of `offer`, which has already been
899                    // maximally inlined. If we did, we could mis-inline
900                    // bindings into bodies that precede them, which would
901                    // change the semantics of the expression.
902                    match offer {
903                        InlineOffer::Take(value, _max_iter) => {
904                            *expr = value.take().ok_or_else(|| {
905                                crate::TransformError::Internal(format!(
906                                    "Value already taken for {:?}",
907                                    id
908                                ))
909                            })?;
910                        }
911                        InlineOffer::Clone(value, _max_iter) => {
912                            *expr = value.clone();
913                        }
914                        InlineOffer::Unavailable(_, _) => {
915                            // Do nothing.
916                        }
917                    }
918                } else {
919                    // Presumably a reference to an outer scope.
920                }
921            } else {
922                worklist.extend(expr.children_mut().rev());
923            }
924        }
925        Ok(())
926    }
927}
928
929mod renumbering {
930
931    use std::collections::BTreeMap;
932
933    use itertools::Itertools;
934    use mz_expr::{Id, LocalId, MirRelationExpr};
935    use mz_ore::id_gen::IdGen;
936
937    /// Re-assign an identifier to each `Let`.
938    ///
939    /// Under the assumption that `id_gen` produces identifiers in order, this process
940    /// maintains in-orderness of `LetRec` identifiers.
941    pub fn renumber_bindings(
942        relation: &mut MirRelationExpr,
943        id_gen: &mut IdGen,
944    ) -> Result<(), crate::TransformError> {
945        let mut renaming = BTreeMap::new();
946        determine(&*relation, &mut renaming, id_gen)?;
947        implement(relation, &renaming)?;
948        Ok(())
949    }
950
951    /// Performs an in-order traversal of the AST, assigning identifiers as it goes.
952    fn determine(
953        relation: &MirRelationExpr,
954        remap: &mut BTreeMap<LocalId, LocalId>,
955        id_gen: &mut IdGen,
956    ) -> Result<(), crate::TransformError> {
957        // The stack contains pending work as `Result<LocalId, &MirRelationExpr>`, where
958        // 1. 'Ok(id)` means the identifier `id` is ready for renumbering,
959        // 2. `Err(expr)` means that the expression `expr` needs to be further processed.
960        let mut stack: Vec<Result<LocalId, _>> = vec![Err(relation)];
961        while let Some(action) = stack.pop() {
962            match action {
963                Ok(id) => {
964                    if remap.contains_key(&id) {
965                        Err(crate::TransformError::Internal(format!(
966                            "Shadowing of let binding for {:?}",
967                            id
968                        )))?;
969                    } else {
970                        remap.insert(id, LocalId::new(id_gen.allocate_id()));
971                    }
972                }
973                Err(expr) => match expr {
974                    MirRelationExpr::Let { id, value, body } => {
975                        stack.push(Err(body));
976                        stack.push(Ok(*id));
977                        stack.push(Err(value));
978                    }
979                    MirRelationExpr::LetRec {
980                        ids,
981                        values,
982                        limits: _,
983                        body,
984                    } => {
985                        stack.push(Err(body));
986                        for (id, value) in ids.iter().rev().zip_eq(values.iter().rev()) {
987                            stack.push(Ok(*id));
988                            stack.push(Err(value));
989                        }
990                    }
991                    _ => {
992                        stack.extend(expr.children().rev().map(Err));
993                    }
994                },
995            }
996        }
997        Ok(())
998    }
999
1000    fn implement(
1001        relation: &mut MirRelationExpr,
1002        remap: &BTreeMap<LocalId, LocalId>,
1003    ) -> Result<(), crate::TransformError> {
1004        let mut worklist = vec![relation];
1005        while let Some(expr) = worklist.pop() {
1006            match expr {
1007                MirRelationExpr::Let { id, .. } => {
1008                    *id = *remap
1009                        .get(id)
1010                        .ok_or(crate::TransformError::IdentifierMissing(*id))?;
1011                }
1012                MirRelationExpr::LetRec { ids, .. } => {
1013                    for id in ids.iter_mut() {
1014                        *id = *remap
1015                            .get(id)
1016                            .ok_or(crate::TransformError::IdentifierMissing(*id))?;
1017                    }
1018                }
1019                MirRelationExpr::Get {
1020                    id: Id::Local(id), ..
1021                } => {
1022                    *id = *remap
1023                        .get(id)
1024                        .ok_or(crate::TransformError::IdentifierMissing(*id))?;
1025                }
1026                _ => {
1027                    // Remapped identifiers not used in these patterns.
1028                }
1029            }
1030            // The order is not critical, but behave as a stack for clarity.
1031            worklist.extend(expr.children_mut().rev());
1032        }
1033        Ok(())
1034    }
1035}