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}