// 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.

//! Lowering [`DataflowDescription`]s from MIR ([`MirRelationExpr`]) to LIR ([`Plan`]).

use std::collections::BTreeMap;

use mz_expr::JoinImplementation::{DeltaQuery, Differential, IndexedFilter, Unimplemented};
use mz_expr::{
    permutation_for_arrangement, Id, JoinInputMapper, MapFilterProject, MirRelationExpr,
    MirScalarExpr, OptimizedMirRelationExpr,
};
use mz_ore::{assert_none, soft_assert_eq_or_log, soft_panic_or_log};
use mz_repr::optimize::OptimizerFeatures;
use mz_repr::GlobalId;
use timely::progress::Timestamp;

use crate::dataflows::{BuildDesc, DataflowDescription, IndexImport};
use crate::plan::join::{DeltaJoinPlan, JoinPlan, LinearJoinPlan};
use crate::plan::reduce::{KeyValPlan, ReducePlan};
use crate::plan::threshold::ThresholdPlan;
use crate::plan::top_k::TopKPlan;
use crate::plan::{AvailableCollections, GetPlan, LirId, Plan};

pub(super) struct Context {
    /// Known bindings to (possibly arranged) collections.
    arrangements: BTreeMap<Id, AvailableCollections>,
    /// Tracks the next available `LirId`.
    next_lir_id: LirId,
    /// Information to print along with error messages.
    debug_info: LirDebugInfo,
    /// Whether to enable fusion of MFPs in reductions.
    enable_reduce_mfp_fusion: bool,
}

impl Context {
    pub fn new(debug_name: String, features: &OptimizerFeatures) -> Self {
        Self {
            arrangements: Default::default(),
            next_lir_id: 0,
            debug_info: LirDebugInfo {
                debug_name,
                id: GlobalId::Transient(0),
            },
            enable_reduce_mfp_fusion: features.enable_reduce_mfp_fusion,
        }
    }

    fn allocate_lir_id(&mut self) -> LirId {
        let id = self.next_lir_id;
        self.next_lir_id += 1;
        id
    }

    pub fn lower<T: Timestamp>(
        mut self,
        desc: DataflowDescription<OptimizedMirRelationExpr>,
    ) -> Result<DataflowDescription<Plan<T>>, String> {
        // Sources might provide arranged forms of their data, in the future.
        // Indexes provide arranged forms of their data.
        for IndexImport {
            desc: index_desc,
            typ,
            ..
        } in desc.index_imports.values()
        {
            let key = index_desc.key.clone();
            // TODO[btv] - We should be told the permutation by
            // `index_desc`, and it should have been generated
            // at the same point the thinning logic was.
            //
            // We should for sure do that soon, but it requires
            // a bit of a refactor, so for now we just
            // _assume_ that they were both generated by `permutation_for_arrangement`,
            // and recover it here.
            let (permutation, thinning) = permutation_for_arrangement(&key, typ.arity());
            let index_keys = self
                .arrangements
                .entry(Id::Global(index_desc.on_id))
                .or_insert_with(AvailableCollections::default);
            index_keys.arranged.push((key, permutation, thinning));
            index_keys.types = Some(typ.column_types.clone());
        }
        for id in desc.source_imports.keys() {
            self.arrangements
                .entry(Id::Global(*id))
                .or_insert_with(AvailableCollections::new_raw);
        }

        // Build each object in order, registering the arrangements it forms.
        let mut objects_to_build = Vec::with_capacity(desc.objects_to_build.len());
        for build in desc.objects_to_build {
            self.debug_info.id = build.id;
            let (plan, keys) = self.lower_mir_expr(&build.plan)?;

            self.arrangements.insert(Id::Global(build.id), keys);
            objects_to_build.push(BuildDesc { id: build.id, plan });
        }

        Ok(DataflowDescription {
            source_imports: desc.source_imports,
            index_imports: desc.index_imports,
            objects_to_build,
            index_exports: desc.index_exports,
            sink_exports: desc.sink_exports,
            as_of: desc.as_of,
            until: desc.until,
            initial_storage_as_of: desc.initial_storage_as_of,
            refresh_schedule: desc.refresh_schedule,
            debug_name: desc.debug_name,
        })
    }

    /// This method converts a MirRelationExpr into a plan that can be directly rendered.
    ///
    /// The rough structure is that we repeatedly extract map/filter/project operators
    /// from each expression we see, bundle them up as a `MapFilterProject` object, and
    /// then produce a plan for the combination of that with the next operator.
    ///
    /// The method accesses `self.arrangements`, which it will locally add to and remove from for
    /// `Let` bindings (by the end of the call it should contain the same bindings as when it
    /// started).
    ///
    /// The result of the method is both a `Plan`, but also a list of arrangements that
    /// are certain to be produced, which can be relied on by the next steps in the plan.
    /// Each of the arrangement keys is associated with an MFP that must be applied if that
    /// arrangement is used, to back out the permutation associated with that arrangement.
    ///
    /// An empty list of arrangement keys indicates that only a `Collection` stream can
    /// be assumed to exist.
    fn lower_mir_expr<T: Timestamp>(
        &mut self,
        expr: &MirRelationExpr,
    ) -> Result<(Plan<T>, AvailableCollections), String> {
        // This function is recursive and can overflow its stack, so grow it if
        // needed. The growth here is unbounded. Our general solution for this problem
        // is to use [`ore::stack::RecursionGuard`] to additionally limit the stack
        // depth. That however requires upstream error handling. This function is
        // currently called by the Coordinator after calls to `catalog_transact`,
        // and thus are not allowed to fail. Until that allows errors, we choose
        // to allow the unbounded growth here. We are though somewhat protected by
        // higher levels enforcing their own limits on stack depth (in the parser,
        // transformer/desugarer, and planner).
        mz_ore::stack::maybe_grow(|| self.lower_mir_expr_stack_safe(expr))
    }

    fn lower_mir_expr_stack_safe<T>(
        &mut self,
        expr: &MirRelationExpr,
    ) -> Result<(Plan<T>, AvailableCollections), String>
    where
        T: Timestamp,
    {
        // Extract a maximally large MapFilterProject from `expr`.
        // We will then try and push this in to the resulting expression.
        //
        // Importantly, `mfp` may contain temporal operators and not be a "safe" MFP.
        // While we would eventually like all plan stages to be able to absorb such
        // general operators, not all of them can.
        let (mut mfp, expr) = MapFilterProject::extract_from_expression(expr);
        // We attempt to plan what we have remaining, in the context of `mfp`.
        // We may not be able to do this, and must wrap some operators with a `Mfp` stage.
        let (mut plan, mut keys) = match expr {
            // These operators should have been extracted from the expression.
            MirRelationExpr::Map { .. } => {
                panic!("This operator should have been extracted");
            }
            MirRelationExpr::Filter { .. } => {
                panic!("This operator should have been extracted");
            }
            MirRelationExpr::Project { .. } => {
                panic!("This operator should have been extracted");
            }
            // These operators may not have been extracted, and need to result in a `Plan`.
            MirRelationExpr::Constant { rows, typ: _ } => {
                let plan = Plan::Constant {
                    rows: rows.clone().map(|rows| {
                        rows.into_iter()
                            .map(|(row, diff)| (row, T::minimum(), diff))
                            .collect()
                    }),
                    lir_id: self.allocate_lir_id(),
                };
                // The plan, not arranged in any way.
                (plan, AvailableCollections::new_raw())
            }
            MirRelationExpr::Get { id, typ: _, .. } => {
                // This stage can absorb arbitrary MFP operators.
                let mut mfp = mfp.take();
                // If `mfp` is the identity, we can surface all imported arrangements.
                // Otherwise, we apply `mfp` and promise no arrangements.
                let mut in_keys = self
                    .arrangements
                    .get(id)
                    .cloned()
                    .unwrap_or_else(AvailableCollections::new_raw);

                // Seek out an arrangement key that might be constrained to a literal.
                // Note: this code has very little use nowadays, as its job was mostly taken over
                // by `LiteralConstraints` (see in the below longer comment).
                let key_val = in_keys
                    .arranged
                    .iter()
                    .filter_map(|key| {
                        mfp.literal_constraints(&key.0)
                            .map(|val| (key.clone(), val))
                    })
                    .max_by_key(|(key, _val)| key.0.len());

                // Determine the plan of action for the `Get` stage.
                let plan = if let Some(((key, permutation, thinning), val)) = &key_val {
                    // This code path used to handle looking up literals from indexes, but it's
                    // mostly deprecated, as this is nowadays performed by the `LiteralConstraints`
                    // MIR transform instead. However, it's still called in a couple of tricky
                    // special cases:
                    // - `LiteralConstraints` handles only Gets of global ids, so this code still
                    //   gets to handle Filters on top of Gets of local ids.
                    // - Lowering does a `MapFilterProject::extract_from_expression`, while
                    //   `LiteralConstraints` does
                    //   `MapFilterProject::extract_non_errors_from_expr_mut`.
                    // - It might happen that new literal constraint optimization opportunities
                    //   appear somewhere near the end of the MIR optimizer after
                    //   `LiteralConstraints` has already run.
                    // (Also note that a similar literal constraint handling machinery is also
                    // present when handling the leftover MFP after this big match.)
                    mfp.permute(permutation.clone(), thinning.len() + key.len());
                    in_keys.arranged = vec![(key.clone(), permutation.clone(), thinning.clone())];
                    GetPlan::Arrangement(key.clone(), Some(val.clone()), mfp)
                } else if !mfp.is_identity() {
                    // We need to ensure a collection exists, which means we must form it.
                    if let Some((key, permutation, thinning)) =
                        in_keys.arbitrary_arrangement().cloned()
                    {
                        mfp.permute(permutation.clone(), thinning.len() + key.len());
                        in_keys.arranged = vec![(key.clone(), permutation, thinning)];
                        GetPlan::Arrangement(key, None, mfp)
                    } else {
                        GetPlan::Collection(mfp)
                    }
                } else {
                    // By default, just pass input arrangements through.
                    GetPlan::PassArrangements
                };

                let out_keys = if let GetPlan::PassArrangements = plan {
                    in_keys.clone()
                } else {
                    AvailableCollections::new_raw()
                };

                // Return the plan, and any keys if an identity `mfp`.
                (
                    Plan::Get {
                        id: id.clone(),
                        keys: in_keys,
                        plan,
                        lir_id: self.allocate_lir_id(),
                    },
                    out_keys,
                )
            }
            MirRelationExpr::Let { id, value, body } => {
                // It would be unfortunate to have a non-trivial `mfp` here, as we hope
                // that they would be pushed down. I am not sure if we should take the
                // initiative to push down the `mfp` ourselves.

                // Plan the value using only the initial arrangements, but
                // introduce any resulting arrangements bound to `id`.
                let (value, v_keys) = self.lower_mir_expr(value)?;
                let pre_existing = self.arrangements.insert(Id::Local(*id), v_keys);
                assert_none!(pre_existing);
                // Plan the body using initial and `value` arrangements,
                // and then remove reference to the value arrangements.
                let (body, b_keys) = self.lower_mir_expr(body)?;
                self.arrangements.remove(&Id::Local(*id));
                // Return the plan, and any `body` arrangements.
                (
                    Plan::Let {
                        id: id.clone(),
                        value: Box::new(value),
                        body: Box::new(body),
                        lir_id: self.allocate_lir_id(),
                    },
                    b_keys,
                )
            }
            MirRelationExpr::LetRec {
                ids,
                values,
                limits,
                body,
            } => {
                assert_eq!(ids.len(), values.len());
                assert_eq!(ids.len(), limits.len());
                // Plan the values using only the available arrangements, but
                // introduce any resulting arrangements bound to each `id`.
                // Arrangements made available cannot be used by prior bindings,
                // as we cannot circulate an arrangement through a `Variable` yet.
                let mut lir_values = Vec::with_capacity(values.len());
                for (id, value) in ids.iter().zip(values) {
                    let (mut lir_value, mut v_keys) = self.lower_mir_expr(value)?;
                    // If `v_keys` does not contain an unarranged collection, we must form it.
                    if !v_keys.raw {
                        // Choose an "arbitrary" arrangement; TODO: prefer a specific one.
                        let (input_key, permutation, thinning) =
                            v_keys.arbitrary_arrangement().unwrap();
                        let mut input_mfp = MapFilterProject::new(value.arity());
                        input_mfp.permute(permutation.clone(), thinning.len() + input_key.len());
                        let input_key = Some(input_key.clone());

                        let forms = AvailableCollections::new_raw();

                        // We just want to insert an `ArrangeBy` to form an unarranged collection,
                        // but there is a complication: We shouldn't break the invariant (created by
                        // `NormalizeLets`, and relied upon by the rendering) that there isn't
                        // anything between two `LetRec`s. So if `lir_value` is itself a `LetRec`,
                        // then we insert the `ArrangeBy` on the `body` of the inner `LetRec`,
                        // instead of on top of the inner `LetRec`.
                        lir_value = match lir_value {
                            Plan::LetRec {
                                ids,
                                values,
                                limits,
                                body,
                                lir_id,
                            } => Plan::LetRec {
                                ids,
                                values,
                                limits,
                                body: Box::new(Plan::ArrangeBy {
                                    input: body,
                                    forms,
                                    input_key,
                                    input_mfp,
                                    lir_id: self.allocate_lir_id(),
                                }),
                                lir_id,
                            },
                            lir_value => Plan::ArrangeBy {
                                input: Box::new(lir_value),
                                forms,
                                input_key,
                                input_mfp,
                                lir_id: self.allocate_lir_id(),
                            },
                        };
                        v_keys.raw = true;
                    }
                    let pre_existing = self.arrangements.insert(Id::Local(*id), v_keys);
                    assert_none!(pre_existing);
                    lir_values.push(lir_value);
                }
                // As we exit the iterative scope, we must leave all arrangements behind,
                // as they reference a timestamp coordinate that must be stripped off.
                for id in ids.iter() {
                    self.arrangements
                        .insert(Id::Local(*id), AvailableCollections::new_raw());
                }
                // Plan the body using initial and `value` arrangements,
                // and then remove reference to the value arrangements.
                let (body, b_keys) = self.lower_mir_expr(body)?;
                for id in ids.iter() {
                    self.arrangements.remove(&Id::Local(*id));
                }
                // Return the plan, and any `body` arrangements.
                (
                    Plan::LetRec {
                        ids: ids.clone(),
                        values: lir_values,
                        limits: limits.clone(),
                        body: Box::new(body),
                        lir_id: self.allocate_lir_id(),
                    },
                    b_keys,
                )
            }
            MirRelationExpr::FlatMap { input, func, exprs } => {
                let (input, keys) = self.lower_mir_expr(input)?;
                // This stage can absorb arbitrary MFP instances.
                let mfp = mfp.take();
                let mut exprs = exprs.clone();
                let input_key = if let Some((k, permutation, _)) = keys.arbitrary_arrangement() {
                    // We don't permute the MFP here, because it runs _after_ the table function,
                    // whose output is in a fixed order.
                    //
                    // We _do_, however, need to permute the `expr`s that provide input to the
                    // `func`.
                    for expr in &mut exprs {
                        expr.permute_map(permutation);
                    }

                    Some(k.clone())
                } else {
                    None
                };
                // Return the plan, and no arrangements.
                (
                    Plan::FlatMap {
                        input: Box::new(input),
                        func: func.clone(),
                        exprs: exprs.clone(),
                        mfp_after: mfp,
                        input_key,
                        lir_id: self.allocate_lir_id(),
                    },
                    AvailableCollections::new_raw(),
                )
            }
            MirRelationExpr::Join {
                inputs,
                equivalences,
                implementation,
            } => {
                let input_mapper = JoinInputMapper::new(inputs);

                // Plan each of the join inputs independently.
                // The `plans` get surfaced upwards, and the `input_keys` should
                // be used as part of join planning / to validate the existing
                // plans / to aid in indexed seeding of update streams.
                let mut plans = Vec::new();
                let mut input_keys = Vec::new();
                let mut input_arities = Vec::new();
                for input in inputs.iter() {
                    let (plan, keys) = self.lower_mir_expr(input)?;
                    input_arities.push(input.arity());
                    plans.push(plan);
                    input_keys.push(keys);
                }

                // Extract temporal predicates as joins cannot currently absorb them.
                let (plan, missing) = match implementation {
                    IndexedFilter(_coll_id, _idx_id, key, _val) => {
                        // Start with the constant input. (This used to be important before #14059
                        // was fixed.)
                        let start: usize = 1;
                        let order = vec![(0usize, key.clone(), None)];
                        // All columns of the constant input will be part of the arrangement key.
                        let source_arrangement = (
                            (0..key.len())
                                .map(MirScalarExpr::Column)
                                .collect::<Vec<_>>(),
                            (0..key.len()).map(|i| (i, i)).collect::<BTreeMap<_, _>>(),
                            Vec::<usize>::new(),
                        );
                        let (ljp, missing) = LinearJoinPlan::create_from(
                            start,
                            Some(&source_arrangement),
                            equivalences,
                            &order,
                            input_mapper,
                            &mut mfp,
                            &input_keys,
                        );
                        (JoinPlan::Linear(ljp), missing)
                    }
                    Differential((start, start_arr, _start_characteristic), order) => {
                        let source_arrangement = start_arr.as_ref().and_then(|key| {
                            input_keys[*start]
                                .arranged
                                .iter()
                                .find(|(k, _, _)| k == key)
                                .clone()
                        });
                        let (ljp, missing) = LinearJoinPlan::create_from(
                            *start,
                            source_arrangement,
                            equivalences,
                            order,
                            input_mapper,
                            &mut mfp,
                            &input_keys,
                        );
                        (JoinPlan::Linear(ljp), missing)
                    }
                    DeltaQuery(orders) => {
                        let (djp, missing) = DeltaJoinPlan::create_from(
                            equivalences,
                            orders,
                            input_mapper,
                            &mut mfp,
                            &input_keys,
                        );
                        (JoinPlan::Delta(djp), missing)
                    }
                    // Other plans are errors, and should be reported as such.
                    Unimplemented => return Err("unimplemented join".to_string()),
                };
                // The renderer will expect certain arrangements to exist; if any of those are not available, the join planning functions above should have returned them in
                // `missing`. We thus need to plan them here so they'll exist.
                let is_delta = matches!(plan, JoinPlan::Delta(_));
                for (((input_plan, input_keys), missing), arity) in plans
                    .iter_mut()
                    .zip(input_keys.iter())
                    .zip(missing.into_iter())
                    .zip(input_arities.iter().cloned())
                {
                    if missing != Default::default() {
                        if is_delta {
                            // join_implementation.rs produced a sub-optimal plan here;
                            // we shouldn't plan delta joins at all if not all of the required
                            // arrangements are available. Soft panic in CI and log an error in
                            // production to increase the chances that we will catch all situations
                            // that violate this constraint.
                            soft_panic_or_log!("Arrangements depended on by delta join alarmingly absent: {:?}
Dataflow info: {}
This is not expected to cause incorrect results, but could indicate a performance issue in Materialize.", missing, self.debug_info);
                        } else {
                            soft_panic_or_log!("Arrangements depended on by a non-delta join are absent: {:?}
Dataflow info: {}
This is not expected to cause incorrect results, but could indicate a performance issue in Materialize.", missing, self.debug_info);
                            // Nowadays MIR transforms take care to insert MIR ArrangeBys for each
                            // Join input. (Earlier, they were missing in the following cases:
                            //  - They were const-folded away for constant inputs. This is not
                            //    happening since
                            //    https://github.com/MaterializeInc/materialize/pull/16351
                            //  - They were not being inserted for the constant input of
                            //    `IndexedFilter`s. This was fixed in
                            //    https://github.com/MaterializeInc/materialize/pull/20920
                            //  - They were not being inserted for the first input of Differential
                            //    joins. This was fixed in
                            //    https://github.com/MaterializeInc/materialize/pull/16099)
                        }
                        let raw_plan = std::mem::replace(
                            input_plan,
                            Plan::Constant {
                                rows: Ok(Vec::new()),
                                lir_id: self.allocate_lir_id(),
                            },
                        );
                        *input_plan = self.arrange_by(raw_plan, missing, input_keys, arity);
                    }
                }
                // Return the plan, and no arrangements.
                (
                    Plan::Join {
                        inputs: plans,
                        plan,
                        lir_id: self.allocate_lir_id(),
                    },
                    AvailableCollections::new_raw(),
                )
            }
            MirRelationExpr::Reduce {
                input,
                group_key,
                aggregates,
                monotonic,
                expected_group_size,
            } => {
                let input_arity = input.arity();
                let (input, keys) = self.lower_mir_expr(input)?;
                let (input_key, permutation_and_new_arity) = if let Some((
                    input_key,
                    permutation,
                    thinning,
                )) = keys.arbitrary_arrangement()
                {
                    (
                        Some(input_key.clone()),
                        Some((permutation.clone(), thinning.len() + input_key.len())),
                    )
                } else {
                    (None, None)
                };
                let key_val_plan = KeyValPlan::new(
                    input_arity,
                    group_key,
                    aggregates,
                    permutation_and_new_arity,
                );
                let reduce_plan =
                    ReducePlan::create_from(aggregates.clone(), *monotonic, *expected_group_size);
                // Return the plan, and the keys it produces.
                let mfp_after;
                let output_arity;
                if self.enable_reduce_mfp_fusion {
                    (mfp_after, mfp, output_arity) =
                        reduce_plan.extract_mfp_after(mfp, group_key.len());
                } else {
                    (mfp_after, output_arity) = (
                        MapFilterProject::new(mfp.input_arity),
                        group_key.len() + aggregates.len(),
                    );
                    soft_assert_eq_or_log!(
                        mfp.input_arity,
                        output_arity,
                        "Output arity of reduce must match input arity for MFP on top of it"
                    );
                }
                let output_keys = reduce_plan.keys(group_key.len(), output_arity);
                (
                    Plan::Reduce {
                        input: Box::new(input),
                        key_val_plan,
                        plan: reduce_plan,
                        input_key,
                        mfp_after,
                        lir_id: self.allocate_lir_id(),
                    },
                    output_keys,
                )
            }
            MirRelationExpr::TopK {
                input,
                group_key,
                order_key,
                limit,
                offset,
                monotonic,
                expected_group_size,
            } => {
                let arity = input.arity();
                let (input, keys) = self.lower_mir_expr(input)?;

                let top_k_plan = TopKPlan::create_from(
                    group_key.clone(),
                    order_key.clone(),
                    *offset,
                    limit.clone(),
                    arity,
                    *monotonic,
                    *expected_group_size,
                );

                // We don't have an MFP here -- install an operator to permute the
                // input, if necessary.
                let input = if !keys.raw {
                    self.arrange_by(input, AvailableCollections::new_raw(), &keys, arity)
                } else {
                    input
                };
                // Return the plan, and no arrangements.
                (
                    Plan::TopK {
                        input: Box::new(input),
                        top_k_plan,
                        lir_id: self.allocate_lir_id(),
                    },
                    AvailableCollections::new_raw(),
                )
            }
            MirRelationExpr::Negate { input } => {
                let arity = input.arity();
                let (input, keys) = self.lower_mir_expr(input)?;

                // We don't have an MFP here -- install an operator to permute the
                // input, if necessary.
                let input = if !keys.raw {
                    self.arrange_by(input, AvailableCollections::new_raw(), &keys, arity)
                } else {
                    input
                };
                // Return the plan, and no arrangements.
                (
                    Plan::Negate {
                        input: Box::new(input),
                        lir_id: self.allocate_lir_id(),
                    },
                    AvailableCollections::new_raw(),
                )
            }
            MirRelationExpr::Threshold { input } => {
                let arity = input.arity();
                let (plan, keys) = self.lower_mir_expr(input)?;
                // We don't have an MFP here -- install an operator to permute the
                // input, if necessary.
                let plan = if !keys.raw {
                    self.arrange_by(plan, AvailableCollections::new_raw(), &keys, arity)
                } else {
                    plan
                };
                let (threshold_plan, required_arrangement) = ThresholdPlan::create_from(arity);
                let mut types = keys.types.clone();
                let plan = if !keys
                    .arranged
                    .iter()
                    .any(|(key, _, _)| key == &required_arrangement.0)
                {
                    types = Some(input.typ().column_types);
                    self.arrange_by(
                        plan,
                        AvailableCollections::new_arranged(
                            vec![required_arrangement],
                            types.clone(),
                        ),
                        &keys,
                        arity,
                    )
                } else {
                    plan
                };

                let output_keys = threshold_plan.keys(types);
                // Return the plan, and any produced keys.
                (
                    Plan::Threshold {
                        input: Box::new(plan),
                        threshold_plan,
                        lir_id: self.allocate_lir_id(),
                    },
                    output_keys,
                )
            }
            MirRelationExpr::Union { base, inputs } => {
                let arity = base.arity();
                let mut plans_keys = Vec::with_capacity(1 + inputs.len());
                let (plan, keys) = self.lower_mir_expr(base)?;
                plans_keys.push((plan, keys));
                for input in inputs.iter() {
                    let (plan, keys) = self.lower_mir_expr(input)?;
                    plans_keys.push((plan, keys));
                }
                let plans = plans_keys
                    .into_iter()
                    .map(|(plan, keys)| {
                        // We don't have an MFP here -- install an operator to permute the
                        // input, if necessary.
                        if !keys.raw {
                            self.arrange_by(plan, AvailableCollections::new_raw(), &keys, arity)
                        } else {
                            plan
                        }
                    })
                    .collect();
                // Return the plan and no arrangements.
                let plan = Plan::Union {
                    inputs: plans,
                    consolidate_output: false,
                    lir_id: self.allocate_lir_id(),
                };
                (plan, AvailableCollections::new_raw())
            }
            MirRelationExpr::ArrangeBy { input, keys } => {
                let arity = input.arity();
                let types = Some(input.typ().column_types);
                let (input, mut input_keys) = self.lower_mir_expr(input)?;
                input_keys.types = types;

                // Determine keys that are not present in `input_keys`.
                let new_keys = keys
                    .iter()
                    .filter(|k1| !input_keys.arranged.iter().any(|(k2, _, _)| k1 == &k2))
                    .cloned()
                    .collect::<Vec<_>>();
                if new_keys.is_empty() {
                    (input, input_keys)
                } else {
                    let new_keys = new_keys.iter().cloned().map(|k| {
                        let (permutation, thinning) = permutation_for_arrangement(&k, arity);
                        (k, permutation, thinning)
                    });
                    let (input_key, input_mfp) = if let Some((input_key, permutation, thinning)) =
                        input_keys.arbitrary_arrangement()
                    {
                        let mut mfp = MapFilterProject::new(arity);
                        mfp.permute(permutation.clone(), thinning.len() + input_key.len());
                        (Some(input_key.clone()), mfp)
                    } else {
                        (None, MapFilterProject::new(arity))
                    };
                    input_keys.arranged.extend(new_keys);
                    input_keys.arranged.sort_by(|k1, k2| k1.0.cmp(&k2.0));

                    // Return the plan and extended keys.
                    (
                        Plan::ArrangeBy {
                            input: Box::new(input),
                            forms: input_keys.clone(),
                            input_key,
                            input_mfp,
                            lir_id: self.allocate_lir_id(),
                        },
                        input_keys,
                    )
                }
            }
        };

        // If the plan stage did not absorb all linear operators, introduce a new stage to implement them.
        if !mfp.is_identity() {
            // Seek out an arrangement key that might be constrained to a literal.
            // TODO: Improve key selection heuristic.
            let key_val = keys
                .arranged
                .iter()
                .filter_map(|(key, permutation, thinning)| {
                    let mut mfp = mfp.clone();
                    mfp.permute(permutation.clone(), thinning.len() + key.len());
                    mfp.literal_constraints(key)
                        .map(|val| (key.clone(), permutation, thinning, val))
                })
                .max_by_key(|(key, _, _, _)| key.len());

            // Input key selection strategy:
            // (1) If we can read a key at a particular value, do so
            // (2) Otherwise, if there is a key that causes the MFP to be the identity, and
            // therefore allows us to avoid discarding the arrangement, use that.
            // (3) Otherwise, if there is _some_ key, use that,
            // (4) Otherwise just read the raw collection.
            let input_key_val = if let Some((key, permutation, thinning, val)) = key_val {
                mfp.permute(permutation.clone(), thinning.len() + key.len());

                Some((key, Some(val)))
            } else if let Some((key, permutation, thinning)) =
                keys.arranged.iter().find(|(key, permutation, thinning)| {
                    let mut mfp = mfp.clone();
                    mfp.permute(permutation.clone(), thinning.len() + key.len());
                    mfp.is_identity()
                })
            {
                mfp.permute(permutation.clone(), thinning.len() + key.len());
                Some((key.clone(), None))
            } else if let Some((key, permutation, thinning)) = keys.arbitrary_arrangement() {
                mfp.permute(permutation.clone(), thinning.len() + key.len());
                Some((key.clone(), None))
            } else {
                None
            };

            if mfp.is_identity() {
                // We have discovered a key
                // whose permutation causes the MFP to actually
                // be the identity! We can keep it around,
                // but without its permutation this time,
                // and with a trivial thinning of the right length.
                let (key, val) = input_key_val.unwrap();
                let (_old_key, old_permutation, old_thinning) = keys
                    .arranged
                    .iter_mut()
                    .find(|(key2, _, _)| key2 == &key)
                    .unwrap();
                *old_permutation = (0..mfp.input_arity).map(|i| (i, i)).collect();
                let old_thinned_arity = old_thinning.len();
                *old_thinning = (0..old_thinned_arity).collect();
                // Get rid of all other forms, as this is now the only one known to be valid.
                // TODO[btv] we can probably save the other arrangements too, if we adjust their permutations.
                // This is not hard to do, but leaving it for a quick follow-up to avoid making the present diff too unwieldy.
                keys.arranged.retain(|(key2, _, _)| key2 == &key);
                keys.raw = false;

                // Creating a Plan::Mfp node is now logically unnecessary, but we
                // should do so anyway when `val` is populated, so that
                // the `key_val` optimization gets applied.
                if val.is_some() {
                    plan = Plan::Mfp {
                        input: Box::new(plan),
                        mfp,
                        input_key_val: Some((key, val)),
                        lir_id: self.allocate_lir_id(),
                    }
                }
            } else {
                plan = Plan::Mfp {
                    input: Box::new(plan),
                    mfp,
                    input_key_val,
                    lir_id: self.allocate_lir_id(),
                };
                keys = AvailableCollections::new_raw();
            }
        }

        Ok((plan, keys))
    }

    /// Replace the plan with another one
    /// that has the collection in some additional forms.
    pub fn arrange_by<T>(
        &mut self,
        plan: Plan<T>,
        collections: AvailableCollections,
        old_collections: &AvailableCollections,
        arity: usize,
    ) -> Plan<T> {
        if let Plan::ArrangeBy {
            input,
            mut forms,
            input_key,
            input_mfp,
            lir_id,
        } = plan
        {
            forms.raw |= collections.raw;
            forms.arranged.extend(collections.arranged);
            forms.arranged.sort_by(|k1, k2| k1.0.cmp(&k2.0));
            forms.arranged.dedup_by(|k1, k2| k1.0 == k2.0);
            if forms.types.is_none() {
                forms.types = collections.types;
            } else {
                assert!(collections.types.is_none() || collections.types == forms.types);
            }
            Plan::ArrangeBy {
                input,
                forms,
                input_key,
                input_mfp,
                lir_id,
            }
        } else {
            let (input_key, input_mfp) = if let Some((input_key, permutation, thinning)) =
                old_collections.arbitrary_arrangement()
            {
                let mut mfp = MapFilterProject::new(arity);
                mfp.permute(permutation.clone(), thinning.len() + input_key.len());
                (Some(input_key.clone()), mfp)
            } else {
                (None, MapFilterProject::new(arity))
            };
            Plan::ArrangeBy {
                input: Box::new(plan),
                forms: collections,
                input_key,
                input_mfp,
                lir_id: self.allocate_lir_id(),
            }
        }
    }
}

/// Various bits of state to print along with error messages during LIR planning,
/// to aid debugging.
#[derive(Clone, Debug)]
pub struct LirDebugInfo {
    debug_name: String,
    id: GlobalId,
}

impl std::fmt::Display for LirDebugInfo {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Debug name: {}; id: {}", self.debug_name, self.id)
    }
}