1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.

//! Re-assign type information and identifiers to each `Get` to ensure
//! uniqueness of identifiers.

use std::collections::HashMap;

use crate::TransformArgs;
use mz_expr::visit::VisitChildren;
use mz_expr::{Id, LocalId, MirRelationExpr, RECURSION_LIMIT};
use mz_ore::id_gen::IdGen;
use mz_ore::stack::{CheckedRecursion, RecursionGuard};
use mz_repr::RelationType;

/// Refreshes identifiers and types for local let bindings.
///
/// The analysis is capable of handling shadowing of identifiers, which
/// *shouldn't* happen, but if it does and we wanted to kick and scream,
/// this is one place we could do that. Instead, we'll just come up with
/// guaranteed unique names for each let binding.
#[derive(Debug)]
pub struct UpdateLet {
    recursion_guard: RecursionGuard,
}

impl Default for UpdateLet {
    fn default() -> UpdateLet {
        UpdateLet {
            recursion_guard: RecursionGuard::with_limit(RECURSION_LIMIT),
        }
    }
}

impl CheckedRecursion for UpdateLet {
    fn recursion_guard(&self) -> &RecursionGuard {
        &self.recursion_guard
    }
}

impl crate::Transform for UpdateLet {
    #[tracing::instrument(
        target = "optimizer"
        level = "trace",
        skip_all,
        fields(path.segment = "update_let")
    )]
    fn transform(
        &self,
        relation: &mut MirRelationExpr,
        args: TransformArgs,
    ) -> Result<(), crate::TransformError> {
        let result = self.transform_without_trace(relation, args);
        mz_repr::explain_new::trace_plan(&*relation);
        result
    }
}

impl UpdateLet {
    /// Performs the `UpdateLet` transformation without tracing the result.
    pub fn transform_without_trace(
        &self,
        relation: &mut MirRelationExpr,
        _args: TransformArgs,
    ) -> Result<(), crate::TransformError> {
        let mut id_gen = IdGen::default(); // Get a fresh IdGen.
        self.action(relation, &mut HashMap::new(), &mut id_gen)
    }

    /// Re-assign type information and identifier to each `Get`.
    pub fn action(
        &self,
        relation: &mut MirRelationExpr,
        remap: &mut HashMap<LocalId, (LocalId, RelationType)>,
        id_gen: &mut IdGen,
    ) -> Result<(), crate::TransformError> {
        self.checked_recur(|_| {
            match relation {
                MirRelationExpr::Let { id, value, body } => {
                    self.action(value, remap, id_gen)?;
                    // If a local id, assign a new identifier and refresh the type.
                    let new_id = LocalId::new(id_gen.allocate_id());
                    let prev = remap.insert(id.clone(), (new_id, value.typ()));
                    self.action(body, remap, id_gen)?;
                    remap.remove(id);
                    if let Some(prev_stuff) = prev {
                        remap.insert(id.clone(), prev_stuff);
                    }
                    *id = new_id;
                    Ok(())
                }
                MirRelationExpr::Get { id, typ } => {
                    if let Id::Local(local_id) = id {
                        if let Some((new_id, new_type)) = remap.get(local_id) {
                            *local_id = new_id.clone();
                            *typ = new_type.clone()
                        }
                    }
                    Ok(())
                }
                _ => relation.try_visit_mut_children(|e| self.action(e, remap, id_gen)),
            }
        })
    }
}