lutra-compiler 0.6.0

use std::{borrow::Cow, collections::HashMap};

use crate::intermediate;
use crate::sql::cr;
use crate::sql::utils::RelCols;
use crate::utils::IdGenerator;
use itertools::Itertools;
use lutra_bin::ir;

pub(super) struct Context<'t> {
    bindings: HashMap<u32, usize>,
    functions: HashMap<u32, FuncProvider>,

    defs: HashMap<&'t ir::Path, &'t ir::Ty>,

    scope_id_gen: IdGenerator,
    func_id_gen: IdGenerator,
}

enum FuncProvider {
    Expr(Vec<cr::ExprKind>),
    QueryParam,
}

pub fn compile(program: &ir::Program) -> (cr::Expr, HashMap<&ir::Path, &ir::Ty>) {
    // TODO: this clone is a bit unnecessary
    let (_, id_counts) = intermediate::IdCounter::run(program.clone());

    let mut ctx = Context {
        bindings: Default::default(),
        functions: Default::default(),
        defs: program
            .defs
            .iter()
            .map(|def| (&def.name, &def.ty))
            .collect(),
        scope_id_gen: Default::default(),
        func_id_gen: IdGenerator::new_at(id_counts.max_func_id as usize),
    };

    assert!(
        program.main.ty.kind.is_function(),
        "expected program.main to be a function, got: {:?}",
        program.main.ty
    );

    // find the top-level function
    let func = ctx.as_function_or_wrap(&program.main);
    ctx.functions.insert(func.id, FuncProvider::QueryParam);

    let body = ctx.compile_rel(&func.body);

    (body, ctx.defs)
}

impl<'a> Context<'a> {
    pub(super) fn get_ty_mat(&self, mut ty: &'a ir::Ty) -> &'a ir::Ty {
        while let ir::TyKind::Ident(path) = &ty.kind {
            if ir::TyStd::try_new(path).is_some() {
                return ty;
            }
            ty = self.defs.get(path).unwrap();
        }
        ty
    }

    fn new_binding(&mut self, rel: cr::Expr) -> Box<cr::BoundExpr> {
        Box::new(cr::BoundExpr {
            rel,
            id: self.scope_id_gen.next(),
        })
    }

    /// Constructs a reference to a column of a bound relational expr.
    /// !! Can only be used when reference really is exactly one column !!
    pub fn new_rel_col(
        &mut self,
        rel: &cr::BoundExpr,
        col_position: usize,
        col_ty: ir::Ty,
    ) -> cr::Expr {
        cr::Expr {
            kind: cr::ExprKind::Transform(
                self.new_binding(cr::Expr::new_rel_ref(rel)),
                cr::Transform::ProjectPick(vec![col_position]),
            ),
            ty: col_ty,
        }
    }

    /// Constructs a reference to an array item.
    /// If needed, the reference is also deserialized.
    pub fn new_array_item_ref(&mut self, rel: &cr::BoundExpr) -> cr::Expr {
        let array_ty = self.get_ty_mat(&rel.rel.ty);
        let item_ty = self.get_ty_mat(array_ty.kind.as_array().unwrap());
        let item_ty = item_ty.clone();

        let mut item_ref = cr::ExprKind::Transform(
            self.new_binding(cr::Expr::new_rel_ref(rel)),
            cr::Transform::ProjectDiscard(vec![0]),
        );

        // deserialize from JSON if needed
        if item_ty.kind.is_array() {
            item_ref = cr::ExprKind::From(cr::From::Deserialize(Box::new(cr::Expr {
                kind: item_ref,
                ty: ir::Ty::text(),
            })));
        }

        cr::Expr {
            kind: item_ref,
            ty: item_ty,
        }
    }

    fn compile_rel(&mut self, expr: &ir::Expr) -> cr::Expr {
        let kind = match &expr.kind {
            ir::ExprKind::Literal(lit) => cr::ExprKind::From(cr::From::Row(vec![cr::Expr {
                kind: cr::ExprKind::From(cr::From::Literal(lit.clone())),
                ty: expr.ty.clone(),
            }])),

            ir::ExprKind::Array(items)
                if items.is_empty()
                    && (self.get_ty_mat(&expr.ty).kind.as_ident())
                        .is_some_and(|x| x.is(&["std", "Text"])) =>
            {
                // special case: this is produced by `std::default(): Text`
                let lit = ir::Literal::Text(String::new());
                cr::ExprKind::From(cr::From::Row(vec![cr::Expr {
                    kind: cr::ExprKind::From(cr::From::Literal(lit)),
                    ty: expr.ty.clone(),
                }]))
            }

            ir::ExprKind::Tuple(fields) => {
                let ty_fields = self.get_ty_mat(&expr.ty).kind.as_tuple().unwrap();

                let mut res_rels = Vec::new();

                let mut res_cols = Vec::new();
                let mut res_ty_fields = Vec::new();
                for (field, ty_field) in std::iter::zip(fields, ty_fields) {
                    match self.compile_column_list(&field.expr) {
                        ColumnsOrUnpack::Columns(cols) => {
                            res_cols.extend(cols);
                            res_ty_fields.push(ty_field.clone());
                        }
                        ColumnsOrUnpack::Unpack(rel) => {
                            // finish prev rel
                            if !res_cols.is_empty() {
                                res_rels.push(cr::Expr {
                                    kind: cr::ExprKind::From(cr::From::Row(res_cols)),
                                    ty: ir::Ty::new(ir::TyKind::Tuple(res_ty_fields)),
                                });
                                res_cols = Vec::new();
                                res_ty_fields = Vec::new();
                            }

                            // push this rel
                            res_rels.push(rel);
                        }
                    }
                }

                if res_rels.is_empty() {
                    // simple case: only simple columns, just From::Row

                    cr::ExprKind::From(cr::From::Row(res_cols))
                } else {
                    // there are rels, Join needed

                    // finish last rel
                    if !res_cols.is_empty() {
                        res_rels.push(cr::Expr {
                            kind: cr::ExprKind::From(cr::From::Row(res_cols)),
                            ty: ir::Ty::new(ir::TyKind::Tuple(res_ty_fields)),
                        });
                    }

                    let mut res_rels = res_rels.into_iter().peekable();
                    let mut result = res_rels.next().unwrap();

                    if res_rels.peek().is_none() {
                        // a single rel, wrap it into a Transform::ProjectDiscard to update column names
                        cr::ExprKind::Transform(
                            self.new_binding(result),
                            cr::Transform::ProjectDiscard(vec![]),
                        )
                    } else {
                        // Join with all remaining rels
                        for r in res_rels {
                            result = cr::Expr {
                                ty: ty_concat_as_tuples(
                                    self.get_ty_mat(&result.ty).clone(),
                                    self.get_ty_mat(&r.ty).clone(),
                                ),
                                kind: cr::ExprKind::Join(
                                    self.new_binding(result),
                                    self.new_binding(r),
                                    None,
                                ),
                            };
                        }
                        result.kind
                    }
                }
            }

            ir::ExprKind::Array(items) => {
                let rows = items
                    .iter()
                    .enumerate()
                    .map(|(index, x)| {
                        let row = vec![new_index(index as i64)];

                        let x = self.compile_column_list(x);

                        cr::Expr {
                            kind: self.row_or_join(row, x),
                            ty: expr.ty.clone(),
                        }
                    })
                    .collect();

                cr::ExprKind::Union(rows)
            }

            ir::ExprKind::Call(call) => match &call.function.kind {
                ir::ExprKind::Pointer(ir::Pointer::External(ptr))
                    if ptr.id.starts_with("std::") =>
                {
                    return self.compile_rel_std(expr);
                }
                ir::ExprKind::Pointer(ir::Pointer::External(ptr)) => {
                    panic!("Function {} is not implemented", ptr.id)
                }
                ir::ExprKind::Pointer(ir::Pointer::Parameter(_)) => {
                    todo!("Call of a param")
                }
                ir::ExprKind::Pointer(ir::Pointer::Binding(id)) => {
                    todo!("Call of a binding: id={id}")
                }
                ir::ExprKind::Call(_) => todo!(),
                ir::ExprKind::Function(_) => todo!(),

                ir::ExprKind::TupleLookup(_) => todo!(),
                ir::ExprKind::Binding(_) => todo!(),

                ir::ExprKind::Literal(_)
                | ir::ExprKind::Tuple(_)
                | ir::ExprKind::Array(_)
                | ir::ExprKind::EnumVariant(_)
                | ir::ExprKind::EnumTag(_)
                | ir::ExprKind::EnumUnwrap(_)
                | ir::ExprKind::Switch(_) => {
                    unreachable!()
                }
            },

            ir::ExprKind::Pointer(ir::Pointer::Binding(ptr)) => {
                let name = self.bindings.get(ptr).unwrap();
                cr::ExprKind::From(cr::From::RelRef(*name))
            }
            ir::ExprKind::Pointer(ir::Pointer::Parameter(ptr)) => {
                let provider = self.functions.get(&ptr.function_id).unwrap();
                match provider {
                    FuncProvider::Expr(r_expr) => {
                        r_expr.get(ptr.param_position as usize).unwrap().clone()
                    }
                    FuncProvider::QueryParam => {
                        assert_eq!(ptr.param_position, 0);

                        match &self.get_ty_mat(&expr.ty).kind {
                            ir::TyKind::Primitive(_)
                            | ir::TyKind::Tuple(_)
                            | ir::TyKind::Enum(_)
                            | ir::TyKind::Ident(_) => {
                                let columns = self
                                    .rel_cols_ty_nested(&expr.ty)
                                    .enumerate()
                                    .map(|(p, field_ty)| cr::Expr {
                                        kind: cr::ExprKind::From(cr::From::Param(
                                            u16::try_from(p)
                                                .expect("too many flattened SQL parameters"),
                                        )),
                                        ty: field_ty.into_owned(),
                                    })
                                    .collect();
                                cr::ExprKind::From(cr::From::Row(columns))
                            }

                            ir::TyKind::Array(_) => {
                                // param will be encoded as JSON
                                cr::ExprKind::From(cr::From::Deserialize(Box::new(cr::Expr {
                                    kind: cr::ExprKind::From(cr::From::Param(0)),
                                    ty: expr.ty.clone(),
                                })))
                            }

                            ir::TyKind::Function(_) => unreachable!(),
                        }
                    }
                }
            }
            ir::ExprKind::Pointer(ir::Pointer::External(_)) => todo!(),

            ir::ExprKind::Function(_) => todo!(),

            ir::ExprKind::EnumVariant(variant) => {
                let ty_mat = self.get_ty_mat(&expr.ty);
                let ir::TyKind::Enum(ty_variants) = &ty_mat.kind else {
                    panic!("invalid program");
                };

                if self.is_option(ty_variants) {
                    // special case: nullable column

                    if variant.tag == 0 {
                        let null = cr::Expr {
                            kind: cr::ExprKind::From(cr::From::Null),
                            ty: ty_variants[1].ty.clone(),
                        };
                        cr::ExprKind::From(cr::From::Row(vec![null]))
                    } else {
                        self.compile_rel(&variant.inner).kind
                    }
                } else {
                    let cols = std::iter::zip(
                        self.rel_cols_nested(&expr.ty, "".into()),
                        self.rel_cols_ty_nested(&expr.ty),
                    );
                    let mut row = Vec::with_capacity(cols.size_hint().1.unwrap_or_default());
                    let mut cols = cols.peekable();

                    // tag
                    cols.next().unwrap();
                    row.push(new_tag(variant.tag as i16));

                    let inner_name = format!("_{}", variant.tag);

                    // spacing
                    row.extend(
                        // columns before inner
                        cols.peeking_take_while(|(n, _)| !n.starts_with(&inner_name))
                            .map(|(_, t)| t.into_owned())
                            .map(cr::Expr::null),
                    );
                    // consume inner columns
                    cols.peeking_take_while(|(n, _)| n.starts_with(&inner_name))
                        .count();
                    // column after inner
                    let spacing_after: Vec<_> = cols
                        .map(|(_, t)| t.into_owned())
                        .map(cr::Expr::null)
                        .collect();

                    // inner
                    let is_recursive = lutra_bin::layout::does_enum_variant_contain_recursive(
                        ty_mat,
                        variant.tag as u16,
                    );
                    if is_recursive {
                        let inner = self.compile_rel(&variant.inner);
                        row.push(cr::Expr::new_serialize(inner));
                    } else {
                        row.extend(self.compile_column_list(&variant.inner).unwrap_columns());
                    }

                    row.extend(spacing_after);

                    cr::ExprKind::From(cr::From::Row(row))
                }
            }
            ir::ExprKind::EnumTag(e) => {
                let base = self.compile_rel(&e.subject);

                let ir::TyKind::Enum(variants) = &self.get_ty_mat(&e.subject.ty).kind else {
                    panic!("invalid program");
                };
                if self.is_option(variants) {
                    // a nullable column: NULL is tag 0 (none), otherwise tag 1 (some)
                    let is_null = cr::Expr {
                        kind: cr::ExprKind::From(cr::From::FuncCall(
                            "is_not_null".into(),
                            vec![base],
                        )),
                        ty: ir::Ty::bool(),
                    };
                    // wrap into `_::int::int2`
                    let as_int = cr::Expr {
                        kind: cr::ExprKind::From(cr::From::Cast(Box::new(is_null))),
                        ty: ir::Ty::new_ident(&["std", "Int32"]),
                    };
                    cr::ExprKind::From(cr::From::Cast(Box::new(as_int)))
                } else {
                    // tag + one column for variant: pick the tag column
                    let base = self.new_binding(base);

                    cr::ExprKind::Transform(base, cr::Transform::ProjectPick(vec![0]))
                }
            }
            ir::ExprKind::EnumUnwrap(enum_unwrap) => {
                let base = self.compile_rel(&enum_unwrap.subject);

                let ir::TyKind::Enum(variants) = &self.get_ty_mat(&enum_unwrap.subject.ty).kind
                else {
                    panic!("invalid program");
                };

                if self.is_option(variants) {
                    // a nullable column
                    return base;
                } else {
                    // tag + one column for variant

                    let base = self.new_binding(base);

                    // count number of columns in front of selected variant and the first tag
                    let start = 1 + variants
                        .iter()
                        .take(enum_unwrap.tag as usize)
                        .map(|v| self.rel_cols_nested(&v.ty, String::new()).count())
                        .sum::<usize>();

                    // count number of columns of the selected variant
                    let end = start
                        + self
                            .rel_cols_nested(&variants[enum_unwrap.tag as usize].ty, String::new())
                            .count();

                    cr::ExprKind::Transform(
                        base,
                        cr::Transform::ProjectPick((start..end).collect()),
                    )
                }
            }

            ir::ExprKind::TupleLookup(lookup) => {
                let base = self.compile_rel(&lookup.base);
                let position = lookup.position as usize;

                let ir::TyKind::Tuple(fields) = &self.get_ty_mat(&lookup.base.ty).kind else {
                    panic!("invalid program");
                };

                // The target might be packed. Performing a tuple lookup of such target must
                // change the repr into normal SQL repr.
                // Currently this only happens for arrays, which change from JSON to SQL repr.
                let ty_mat = self.get_ty_mat(&expr.ty);
                let unpack = ty_mat.kind.is_array();

                // In simple case start == lookup.position. But this tuple might contains fields
                // before lookup.position which are represented with multiple columns.
                // So we must iterate over all preceding fields and count how many columns do they
                // produce.
                let start: usize = fields
                    .iter()
                    .take(position)
                    .map(|f| self.rel_cols_ty_nested(&f.ty).count())
                    .sum();

                // Also, target of lookup might not contain a single column
                // (e.g. it might be a nested tuple with two fields).
                let end = start
                    + if unpack {
                        1
                    } else {
                        self.rel_cols_ty_nested(&fields[position].ty).count()
                    };

                let mut rel = cr::ExprKind::Transform(
                    self.new_binding(base),
                    cr::Transform::ProjectPick((start..end).collect()),
                );

                if unpack {
                    rel = cr::ExprKind::From(cr::From::Deserialize(Box::new(cr::Expr {
                        kind: rel,
                        ty: ir::Ty::text(),
                    })));
                }

                rel
            }
            ir::ExprKind::Binding(binding) => {
                if let ir::TyKind::Function(_) = &binding.expr.ty.kind {
                    todo!()
                } else {
                    // compile expr
                    let expr = self.compile_rel(&binding.expr);
                    let expr = self.new_binding(expr);
                    self.bindings.insert(binding.id, expr.id);

                    // compile main with expr in ctx
                    let main = self.compile_rel(&binding.main);

                    self.bindings.remove(&binding.id).unwrap();

                    cr::ExprKind::Bind(expr, Box::new(main))
                }
            }

            ir::ExprKind::Switch(switch) => {
                let single_col = self.rel_cols(&expr.ty).nth(1).is_none();

                if single_col {
                    // CASE
                    // WHEN b0.condition THEN b0.value
                    // WHEN b1.condition THEN b1.value
                    //                   ELSE b2.value
                    let mut cases = Vec::with_capacity(switch.len());
                    for branch in switch {
                        let condition = self.compile_column(&branch.condition);
                        let value = self.compile_column(&branch.value);

                        cases.push((condition, value));
                    }
                    cr::ExprKind::From(cr::From::Case(cases))
                } else {
                    // WITH selector AS (CASE ...)
                    // b0.value WHERE selector = 0
                    // UNION ALL
                    // b1.value WHERE selector = 1
                    // UNION ALL
                    // b2.value WHERE selector = 2

                    fn new_branch_selector(index: usize) -> cr::Expr {
                        new_tag(index as i16)
                    }

                    let mut cases = Vec::with_capacity(switch.len());
                    for (index, branch) in switch.iter().enumerate() {
                        let condition = self.compile_column(&branch.condition);
                        let value = new_branch_selector(index);
                        cases.push((condition, value));
                    }
                    let selector = cr::Expr {
                        kind: cr::ExprKind::From(cr::From::Case(cases)),
                        ty: ty_tag(),
                    };
                    let selector = self.new_binding(selector);

                    let selector_ref = self.new_rel_col(&selector, 0, selector.rel.ty.clone());

                    let mut branches = Vec::with_capacity(switch.len());
                    for (index, branch) in switch.iter().enumerate() {
                        let condition = cr::Expr {
                            kind: cr::ExprKind::From(cr::From::FuncCall(
                                "std::ops::eq".into(),
                                vec![selector_ref.clone(), new_branch_selector(index)],
                            )),
                            ty: ir::Ty::bool(),
                        };

                        let value = self.compile_rel(&branch.value);
                        let value = self.new_binding(value);

                        branches.push(cr::Expr::new_iso_transform(
                            value,
                            cr::Transform::Where(Box::new(condition)),
                        ));
                    }
                    let union = cr::Expr {
                        kind: cr::ExprKind::Union(branches),
                        ty: expr.ty.clone(),
                    };
                    cr::ExprKind::Bind(selector, Box::new(union))
                }
            }
        };
        cr::Expr {
            kind,
            ty: expr.ty.clone(),
        }
    }

    fn row_or_join(&mut self, mut row: Vec<cr::Expr>, expr: ColumnsOrUnpack) -> cr::ExprKind {
        match expr {
            ColumnsOrUnpack::Columns(cols) => {
                row.extend(cols);
                cr::ExprKind::From(cr::From::Row(row))
            }
            ColumnsOrUnpack::Unpack(e) => cr::ExprKind::Join(
                self.new_binding(cr::Expr {
                    kind: cr::ExprKind::From(cr::From::Row(row)),
                    ty: ty_index(),
                }),
                self.new_binding(e),
                None,
            ),
        }
    }

    fn compile_rel_std(&mut self, expr: &ir::Expr) -> cr::Expr {
        let ir::ExprKind::Call(call) = &expr.kind else {
            unreachable!()
        };
        let ir::ExprKind::Pointer(ir::Pointer::External(ptr)) = &call.function.kind else {
            unreachable!()
        };

        let kind = match ptr.id.as_str() {
            "std::array::slice" => {
                let array = self.compile_rel(&call.args[0]);
                let start = self.compile_column(&call.args[1]);
                let end = self.compile_column(&call.args[2]);

                let array = self.new_binding(array);
                let index_col = self.new_rel_col(&array, 0, ty_index());

                let filtered = cr::Expr::new_iso_transform(
                    array,
                    cr::Transform::Where(Box::new(new_bin_op(
                        new_bin_op(start, "std::ops::lte", index_col.clone(), ir::Ty::bool()),
                        "std::ops::and",
                        new_bin_op(index_col, "std::ops::lt", end, ir::Ty::bool()),
                        ir::Ty::bool(),
                    ))),
                );

                // reindex
                let filtered = self.new_binding(filtered);
                let index_col = self.new_rel_col(&filtered, 0, ty_index());
                cr::ExprKind::Transform(filtered, cr::Transform::Reindex(vec![index_col]))
            }
            "std::array::index" => {
                let array = self.compile_rel(&call.args[0]);
                let index = self.compile_column(&call.args[1]);

                let item_ty = self.get_ty_mat(&array.ty).kind.as_array().unwrap();
                let item_ty = self.get_ty_mat(item_ty).clone();

                let array = self.new_binding(array);
                let index_col = self.new_rel_col(&array, 0, ty_index());

                let item_row = cr::Expr::new_iso_transform(
                    array,
                    cr::Transform::Where(Box::new(new_bin_op(
                        index_col,
                        "std::ops::eq",
                        index,
                        ir::Ty::bool(),
                    ))),
                );

                // drop index column
                let item = cr::ExprKind::Transform(
                    self.new_binding(item_row),
                    cr::Transform::ProjectDiscard(vec![0]),
                );

                let item = cr::Expr {
                    kind: item,
                    ty: if item_ty.kind.is_array() {
                        ir::Ty::text()
                    } else {
                        item_ty
                    },
                };

                let ty_variants = expr.ty.kind.as_enum().unwrap();

                if self.is_option(ty_variants) {
                    cr::ExprKind::From(cr::From::Row(vec![item]))
                } else {
                    // construct Option::Some
                    let some = cr::Expr {
                        kind: cr::ExprKind::Join(
                            self.new_binding(new_tag(1)),
                            self.new_binding(item),
                            None,
                        ),
                        ty: expr.ty.clone(),
                    };

                    // construct Option::None
                    let mut none_cols = vec![new_tag(0)];
                    none_cols.extend(self.rel_cols_ty_nested(&expr.ty).skip(1).map(|t| cr::Expr {
                        kind: cr::ExprKind::From(cr::From::Null),
                        ty: t.into_owned(),
                    }));
                    let none = cr::Expr {
                        kind: cr::ExprKind::From(cr::From::Row(none_cols)),
                        ty: expr.ty.clone(),
                    };

                    // union both variants
                    let union = self.new_binding(cr::Expr {
                        kind: cr::ExprKind::Union(vec![some, none]),
                        ty: expr.ty.clone(),
                    });

                    // take first
                    let tag = self.new_rel_col(&union, 0, ty_tag());
                    let neg_tag = new_un_op("std::ops::neg", tag, ty_tag());
                    cr::ExprKind::Transform(union, cr::Transform::Limit(1, Box::new(neg_tag)))
                }
            }
            "std::array::map" => {
                let array = self.compile_rel(&call.args[0]);
                let array = self.new_binding(array);
                let func = &call.args[1];

                // index
                let row = vec![self.new_rel_col(&array, 0, ty_index())];

                // compile func body
                let func = self.as_function_or_wrap(func);
                let item_ref = self.new_array_item_ref(&array).kind;

                self.functions
                    .insert(func.id, FuncProvider::Expr(vec![item_ref]));
                let mapped_item = self.compile_column_list(&func.body);
                self.functions.remove(&func.id);

                let item = self.row_or_join(row, mapped_item);

                cr::ExprKind::BindCorrelated(
                    array,
                    Box::new(cr::Expr {
                        kind: item,
                        ty: expr.ty.clone(),
                    }),
                )
            }
            "std::array::flat_map" => {
                let array = self.compile_rel(&call.args[0]);
                let array = self.new_binding(array);
                let func = &call.args[1];

                let output_item_ty = *expr.ty.kind.clone().into_array().unwrap();

                // index
                let output_row = vec![self.new_rel_col(&array, 0, ty_index())];

                // compile func body
                let func = func.kind.as_function().unwrap();
                let item_ref = self.new_array_item_ref(&array).kind;

                self.functions
                    .insert(func.id, FuncProvider::Expr(vec![item_ref]));
                let mapped_items = self.compile_rel(&func.body);
                self.functions.remove(&func.id);

                let mapped_items = self.new_binding(mapped_items);

                let mapped_items_ref = self.new_binding(cr::Expr::new_rel_ref(&mapped_items));
                // remove unneeded index of mapped_items
                let mapped_items_ref = self.to_column_list(cr::Expr {
                    kind: cr::ExprKind::Transform(
                        mapped_items_ref,
                        cr::Transform::ProjectDiscard(vec![0]),
                    ),
                    ty: output_item_ty,
                });

                cr::ExprKind::BindCorrelated(
                    array,
                    Box::new(cr::Expr {
                        kind: cr::ExprKind::BindCorrelated(
                            mapped_items,
                            Box::new(cr::Expr {
                                kind: self.row_or_join(output_row, mapped_items_ref),
                                ty: expr.ty.clone(),
                            }),
                        ),
                        ty: expr.ty.clone(),
                    }),
                )
            }
            "std::array::filter" => {
                let array = self.compile_rel(&call.args[0]);
                let array = self.new_binding(array);
                let func = &call.args[1];
                let func = func.kind.as_function().unwrap();

                let item = self.new_array_item_ref(&array);

                self.functions
                    .insert(func.id, FuncProvider::Expr(vec![item.kind]));
                let cond = self.compile_column(&func.body);
                self.functions.remove(&func.id);

                let filtered =
                    cr::Expr::new_iso_transform(array, cr::Transform::Where(Box::new(cond)));

                // reindex
                let filtered = self.new_binding(filtered);
                let index_col = self.new_rel_col(&filtered, 0, ty_index());
                cr::ExprKind::Transform(filtered, cr::Transform::Reindex(vec![index_col]))
            }
            "std::array::sort" => {
                let array = self.compile_rel(&call.args[0]);
                let array = self.new_binding(array);

                let func = &call.args[1];
                let func = func.kind.as_function().unwrap();

                let item = self.new_array_item_ref(&array);

                self.functions
                    .insert(func.id, FuncProvider::Expr(vec![item.kind]));
                let key = self.compile_column(&func.body);
                self.functions.remove(&func.id);

                let old_index = self.new_rel_col(&array, 0, ty_index());

                cr::ExprKind::Transform(array, cr::Transform::Reindex(vec![key, old_index]))
            }

            // aggregation functions
            "std::array::min" | "std::array::max" | "std::array::sum" | "std::array::mean"
            | "std::array::count" | "std::array::any" | "std::array::all" | "std::text::join" => {
                let array = self.compile_rel(&call.args[0]);
                let array = self.new_binding(array);

                let item = self.new_array_item_ref(&array);

                let mut args = vec![item];
                args.extend(call.args[1..].iter().map(|a| self.compile_column(a)));

                cr::ExprKind::Transform(
                    array,
                    cr::Transform::Aggregate(vec![cr::Expr {
                        kind: cr::ExprKind::From(cr::From::FuncCall(ptr.id.clone(), args)),
                        ty: expr.ty.clone(),
                    }]),
                )
            }

            // window functions
            "std::array::lead"
            | "std::array::lag"
            | "std::array::rolling_mean"
            | "std::array::rank"
            | "std::array::rank_dense"
            | "std::array::rank_percentile"
            | "std::array::cume_dist" => {
                let array = self.compile_rel(&call.args[0]);
                let array = self.new_binding(array);

                let item = self.new_array_item_ref(&array);

                let mut args = vec![item];
                args.extend(call.args[1..].iter().map(|a| self.compile_column(a)));

                let row = vec![
                    self.new_rel_col(&array, 0, ty_index()),
                    cr::Expr {
                        kind: cr::ExprKind::From(cr::From::FuncCall(ptr.id.clone(), args)),
                        ty: expr.ty.clone(),
                    },
                ];
                cr::ExprKind::Transform(array, cr::Transform::Aggregate(row))
            }

            "std::array::to_columnar" => {
                let array = self.compile_rel(&call.args[0]);
                let array = self.new_binding(array);

                let ty_out_fields = expr.ty.kind.as_tuple().unwrap();

                let mut aggregate_cols = Vec::with_capacity(ty_out_fields.len());
                for (index, ty_out_field) in ty_out_fields.iter().enumerate() {
                    aggregate_cols.push(cr::Expr::new_serialize(cr::Expr {
                        kind: cr::ExprKind::Transform(
                            self.new_binding(cr::Expr::new_rel_ref(&array)),
                            cr::Transform::ProjectPick(vec![
                                0,         // index
                                1 + index, // tuple field value
                            ]),
                        ),
                        ty: ty_out_field.ty.clone(),
                    }));
                }

                cr::ExprKind::Transform(array, cr::Transform::Aggregate(aggregate_cols))
            }

            "std::array::from_columnar" => {
                let tuple = self.compile_rel(&call.args[0]);
                let tuple = self.new_binding(tuple);

                let ty_in_fields = tuple.rel.ty.kind.as_tuple().unwrap();

                let ty_out_item = expr.ty.kind.as_array().unwrap();
                let ty_out_fields = ty_out_item.kind.as_tuple().unwrap();

                // construct correlated subqueries
                let mut fields = Vec::new();
                for (pos, ty_in_field) in ty_in_fields.iter().enumerate() {
                    let input_col = self.new_rel_col(&tuple, pos, ir::Ty::text());
                    fields.push(cr::Expr {
                        ty: ty_in_field.ty.clone(),
                        kind: cr::ExprKind::From(cr::From::Deserialize(Box::new(input_col))),
                    });
                }
                let fields_len = fields.len();
                let mut fields = fields.into_iter().enumerate();

                if let Some((_, mut joined)) = fields.next() {
                    for (f_index, curr) in fields {
                        let prev = self.new_binding(joined);
                        let curr = self.new_binding(curr);

                        let join_cond = cr::Expr {
                            ty: ir::Ty::bool(),
                            kind: cr::ExprKind::From(cr::From::FuncCall(
                                "std::ops::eq".into(),
                                vec![
                                    self.new_rel_col(&prev, 0, ty_index()),
                                    self.new_rel_col(&curr, 0, ty_index()),
                                ],
                            )),
                        };

                        let ty_index_field = ir::TyTupleField {
                            ty: ty_index(),
                            name: None,
                        };
                        let mut ty_join_fields = Vec::new();
                        for ty_out_field in &ty_out_fields[..f_index] {
                            ty_join_fields.push(ty_out_field.clone());
                        }
                        ty_join_fields.push(ty_index_field.clone());
                        ty_join_fields.push(ty_out_fields[f_index].clone());
                        let ty_join = ir::Ty::new(ir::TyKind::Array(Box::new(ir::Ty::new(
                            ir::TyKind::Tuple(ty_join_fields),
                        ))));

                        let ty_joined_fields = ty_out_fields[..f_index + 1].to_vec();
                        let ty_joined = ir::Ty::new(ir::TyKind::Array(Box::new(ir::Ty::new(
                            ir::TyKind::Tuple(ty_joined_fields),
                        ))));

                        joined = cr::Expr {
                            kind: cr::ExprKind::Transform(
                                self.new_binding(cr::Expr {
                                    kind: cr::ExprKind::Join(prev, curr, Some(Box::new(join_cond))),
                                    ty: ty_join,
                                }),
                                cr::Transform::ProjectDiscard(vec![f_index + 1]),
                            ),
                            ty: ty_joined,
                        };
                    }

                    if fields_len == 1 {
                        joined = cr::Expr {
                            kind: cr::ExprKind::Transform(
                                self.new_binding(joined),
                                cr::Transform::ProjectPick(vec![0, 1]),
                            ),
                            ty: expr.ty.clone(),
                        }
                    }

                    cr::ExprKind::BindCorrelated(tuple, Box::new(joined))
                } else {
                    cr::ExprKind::Union(vec![])
                }
            }

            "std::array::zip" => {
                let a = self.compile_rel(&call.args[0]);
                let a = self.new_binding(a);

                let num_cols_a = self.rel_cols(&a.rel.ty).count();

                let b = self.compile_rel(&call.args[1]);
                let b = self.new_binding(b);

                let a_index = self.new_rel_col(&a, 0, ty_index());
                let b_index = self.new_rel_col(&b, 0, ty_index());
                let condition = new_bin_op(a_index, "std::ops::eq", b_index, ir::Ty::bool());

                let join_ty = ty_zip(
                    self.get_ty_mat(&a.rel.ty).clone(),
                    self.get_ty_mat(&b.rel.ty).clone(),
                );
                let join = cr::Expr {
                    kind: cr::ExprKind::Join(a, b, Some(Box::new(condition))),
                    ty: join_ty,
                };

                cr::ExprKind::Transform(
                    self.new_binding(join),
                    cr::Transform::ProjectDiscard(vec![num_cols_a]), // discard b index
                )
            }

            "std::array::group" => {
                let array = self.compile_rel(&call.args[0]);
                let array = self.new_binding(array);

                let func = &call.args[1];
                let func = func.kind.as_function().unwrap();

                let item_ref = self.new_array_item_ref(&array).kind;

                self.functions
                    .insert(func.id, FuncProvider::Expr(vec![item_ref.clone()]));
                let key = self.compile_rel(&func.body);
                self.functions.remove(&func.id);

                let serialize = cr::Expr::new_serialize(cr::Expr::new_rel_ref(&array));

                let is_key_simple = is_simple(&key);
                let ty_key = self.get_ty_mat(&key.ty);
                let key = if ty_key.kind.is_array() {
                    cr::Expr::new_serialize(key.clone())
                } else {
                    key
                };
                if is_key_simple {
                    // simple key: cloning is fine
                    let values = vec![key.clone(), serialize];
                    let key = Box::new(key);
                    cr::ExprKind::Transform(array, cr::Transform::Group { key, values })
                } else {
                    // complex key: compute key for each row, join it with original rel,
                    // then group the joined relation outside the correlated scope.
                    //
                    // Structure:
                    //   Transform(
                    //     BindCorrelated(array, Join(RelRef(array), key)),
                    //     Group {
                    //       key: key_ref,
                    //       values: [
                    //         key_ref,
                    //         Serialize(array_ref),
                    //       ]
                    //     }
                    //   )

                    let (joined, array_ref, key_ref) = self.correlated_join(array, key);

                    let serialized = cr::Expr::new_serialize(array_ref);

                    cr::ExprKind::Transform(
                        joined,
                        cr::Transform::Group {
                            key: Box::new(key_ref.clone()),
                            values: vec![key_ref, serialized],
                        },
                    )
                }
            }

            "std::array::append" => {
                // compute each rel and apply order to it
                let first = self.compile_rel(&call.args[0]);
                let first =
                    cr::Expr::new_iso_transform(self.new_binding(first), cr::Transform::Order);

                let second = self.compile_rel(&call.args[1]);
                let second =
                    cr::Expr::new_iso_transform(self.new_binding(second), cr::Transform::Order);

                // union
                let union = self.new_binding(cr::Expr {
                    kind: cr::ExprKind::Union(vec![first, second]),
                    ty: expr.ty.clone(),
                });

                // reindex
                cr::ExprKind::Transform(union, cr::Transform::Reindex(vec![]))
            }

            "std::array::fold" => {
                let (inputs, iteration) = self.compile_std_scan(call);

                // order by -index, limit 1

                let iteration = self.new_binding(iteration);
                let index = self.new_rel_col(&iteration, 0, ty_index());
                let neg_index = new_un_op("std::ops::neg", index, ty_index());
                let iteration = cr::Expr::new_iso_transform(
                    iteration,
                    cr::Transform::Limit(1, Box::new(neg_index)),
                );

                // drop index
                let iteration = cr::Expr {
                    kind: cr::ExprKind::Transform(
                        self.new_binding(iteration),
                        cr::Transform::ProjectDiscard(vec![0]),
                    ),
                    ty: expr.ty.clone(),
                };

                cr::ExprKind::Bind(inputs, Box::new(iteration))
            }

            "std::array::scan" => {
                let (inputs, iteration) = self.compile_std_scan(call);

                // filter out initial accumulator
                let iteration = self.new_binding(iteration);
                let filter_initial_acc = new_bin_op(
                    new_index(0),
                    "std::ops::lt",
                    self.new_rel_col(&iteration, 0, ty_index()),
                    ir::Ty::bool(),
                );
                let iteration = cr::Expr::new_iso_transform(
                    iteration,
                    cr::Transform::Where(Box::new(filter_initial_acc)),
                );

                cr::ExprKind::Bind(inputs, Box::new(iteration))
            }

            "std::array::loop_until_empty" => {
                // compute each rel and apply order to it
                let initial = Box::new(self.compile_rel(&call.args[0]));
                let operation = &call.args[1];

                // prepare iterator
                let step_id = self.scope_id_gen.next();
                let iterator_ref = cr::ExprKind::From(cr::From::RelRef(step_id));

                // compile step
                let operation = self.as_function_or_wrap(operation);
                self.functions
                    .insert(operation.id, FuncProvider::Expr(vec![iterator_ref]));
                let step = self.compile_rel(&operation.body);
                self.functions.remove(&operation.id);

                let step = Box::new(cr::BoundExpr {
                    id: step_id,
                    rel: step,
                });

                cr::ExprKind::Iteration(initial, step)
            }

            "std::sql::from" => {
                let table_ident = &call.args[0];
                let ir::ExprKind::Literal(ir::Literal::Text(name)) = &table_ident.kind else {
                    panic!("table identifier must be const")
                };

                cr::ExprKind::From(cr::From::Table {
                    name: name.clone(),
                    use_row_id: false,
                })
            }
            "std::sql::insert" => {
                let rows = self.compile_rel(&call.args[0]);
                let rows = self.new_binding(rows);

                let table_ident = &call.args[1];
                let ir::ExprKind::Literal(ir::Literal::Text(table_ident)) = &table_ident.kind
                else {
                    panic!("table identifier must be const")
                };

                cr::ExprKind::Transform(rows, cr::Transform::Insert(table_ident.clone()))
            }
            "std::sql::update" => {
                // compiles to this:
                //   BindCorrelated(
                //     From(subject_table),
                //     Join(
                //       subject.rowid,
                //       ProjectDiscard(
                //         Filter(
                //           updater,
                //           _t == 1
                //         ),
                //         [0] // option tag
                //       )
                //     )
                //   )

                // subject (the table being updated)
                let table = &call.args[0];
                let ir::ExprKind::Literal(ir::Literal::Text(table)) = &table.kind else {
                    panic!("table identifier must be const")
                };
                let updater_ty = call.args[1].ty.kind.as_function().unwrap();
                let row_ty = updater_ty.params[0].clone();

                let subject = cr::Expr {
                    kind: cr::ExprKind::From(cr::From::Table {
                        name: table.clone(),
                        use_row_id: true,
                    }),
                    ty: ir::Ty::new(ir::TyKind::Array(Box::new(row_ty.clone()))),
                };
                let subject = self.new_binding(subject);

                // updater function (with subject in scope)
                let updater_fn = &call.args[1];
                let updater_fn = updater_fn.kind.as_function().unwrap();

                let row_ref = self.new_array_item_ref(&subject);
                self.functions
                    .insert(updater_fn.id, FuncProvider::Expr(vec![row_ref.kind]));
                let update = self.compile_rel(&updater_fn.body);
                self.functions.remove(&updater_fn.id);

                // filter updates by .some
                let update = self.new_binding(update);
                let tag_is_one = new_bin_op(
                    self.new_rel_col(&update, 0, ty_tag()),
                    "std::ops::eq",
                    new_tag(1),
                    ir::Ty::bool(),
                );
                let update =
                    cr::Expr::new_iso_transform(update, cr::Transform::Where(Box::new(tag_is_one)));

                // discard tag column
                let update = cr::Expr {
                    kind: cr::ExprKind::Transform(
                        self.new_binding(update),
                        cr::Transform::ProjectDiscard(vec![0]),
                    ),
                    ty: row_ty.clone(),
                };

                // join with row id
                let row_id = self.new_rel_col(&subject, 0, ty_index());
                let update = cr::Expr {
                    kind: cr::ExprKind::Join(
                        self.new_binding(row_id),
                        self.new_binding(update),
                        None,
                    ),
                    ty: subject.rel.ty.clone(),
                };

                let updates = Box::new(cr::Expr {
                    ty: subject.rel.ty.clone(),
                    kind: cr::ExprKind::BindCorrelated(subject, Box::new(update)),
                });

                cr::ExprKind::Update {
                    table: table.clone(),
                    updates,
                }
            }
            "std::sql::raw" => {
                let source = &call.args[0];
                let ir::ExprKind::Literal(ir::Literal::Text(source)) = &source.kind else {
                    panic!("sql_source must be const")
                };
                cr::ExprKind::From(cr::From::SQLSource(source.clone()))
            }

            _ => return self.compile_expr_std(expr),
        };
        cr::Expr {
            kind,
            ty: expr.ty.clone(),
        }
    }

    /// Compiles call to `std::scan` to [cr::ExprKind::Iteration].
    /// Result also includes an item for the initial value.
    fn compile_std_scan(&mut self, call: &ir::Call) -> (Box<cr::BoundExpr>, cr::Expr) {
        let inputs = self.compile_rel(&call.args[0]);
        let inputs = self.new_binding(inputs);

        let iteration_id = self.scope_id_gen.next();

        let initial = &call.args[1];
        let operation = &call.args[2];

        let ty_acc = ir::Ty::new(ir::TyKind::Array(Box::new(initial.ty.clone())));

        // -- initial state
        let initial = self.compile_column_list(initial);
        let initial = cr::Expr {
            kind: self.row_or_join(
                vec![new_index(0)], // index
                initial,
            ),
            ty: ty_acc.clone(),
        };

        // -- iteration step

        // a relational reference to the recursive CTE
        let prev_acc_rel = self.new_binding(cr::Expr {
            kind: cr::ExprKind::From(cr::From::RelRef(iteration_id)),
            ty: ty_acc.clone(),
        });

        // reference input rel & pick the row by index
        let input_rel = self.new_binding(cr::Expr::new_rel_ref(&inputs));
        let find_by_index = cr::Transform::Where(Box::new(new_bin_op(
            self.new_rel_col(&input_rel, 0, ty_index()),
            "std::ops::eq",
            self.new_rel_col(&prev_acc_rel, 0, ty_index()),
            ir::Ty::bool(),
        )));
        let input_rel = self.new_binding(cr::Expr::new_iso_transform(input_rel, find_by_index));

        //  prepare operation args
        let acc_ref = cr::ExprKind::Transform(
            self.new_binding(cr::Expr::new_rel_ref(&prev_acc_rel)),
            cr::Transform::ProjectDiscard(vec![0]), // discard index
        );
        let input_ref = cr::ExprKind::Transform(
            self.new_binding(cr::Expr::new_rel_ref(&input_rel)),
            cr::Transform::ProjectDiscard(vec![0]), // discard index
        );

        // compile reduced state
        let operation = self.as_function_or_wrap(operation);
        self.functions
            .insert(operation.id, FuncProvider::Expr(vec![acc_ref, input_ref]));
        let op_result = self.compile_column_list(&operation.body);
        self.functions.remove(&operation.id);

        // compose the new state
        let new_index = new_bin_op(
            self.new_rel_col(&prev_acc_rel, 0, ty_index()),
            "std::ops::add",
            new_index(1),
            ty_index(),
        );
        let new_acc = cr::Expr {
            kind: self.row_or_join(vec![new_index], op_result),
            ty: ty_acc.clone(),
        };

        let step = cr::Expr::new_bind_correlated(
            prev_acc_rel,
            cr::Expr::new_bind_correlated(input_rel, new_acc),
        );

        let iteration = cr::Expr {
            kind: cr::ExprKind::Iteration(
                Box::new(initial),
                Box::new(cr::BoundExpr {
                    id: iteration_id,
                    rel: step,
                }),
            ),
            ty: ty_acc,
        };
        (inputs, iteration)
    }

    /// Compiles an expression into a column
    #[track_caller]
    fn compile_column(&mut self, expr: &ir::Expr) -> cr::Expr {
        let cols = self.compile_column_list(expr);
        let ColumnsOrUnpack::Columns(cols) = cols else {
            panic!("expected columns, found: {cols:?}");
        };
        assert!(cols.len() == 1, "expected a single column, found: {cols:?}");
        cols.into_iter().next().unwrap()
    }

    /// Compiles an expression into a list of columns
    fn compile_column_list(&mut self, expr: &ir::Expr) -> ColumnsOrUnpack {
        // compile as if this was a rel, which can sink complex operations
        // into the relational expression
        let rel = self.compile_rel(expr);

        self.to_column_list(rel)
    }

    fn to_column_list(&self, rel: cr::Expr) -> ColumnsOrUnpack {
        let ty_mat = self.get_ty_mat(&rel.ty);

        match rel.kind {
            // it is just one constructed row: unwrap
            cr::ExprKind::From(cr::From::Row(row)) if !ty_mat.kind.is_array() => {
                ColumnsOrUnpack::Columns(row)
            }

            // it is actually complex: subquery
            _ => {
                match &ty_mat.kind {
                    ir::TyKind::Primitive(_) | ir::TyKind::Ident(_) => {
                        // return as a subquery
                        ColumnsOrUnpack::Columns(vec![rel])
                    }
                    ir::TyKind::Array(_) => {
                        // arrays need to packed to JSON
                        ColumnsOrUnpack::Columns(vec![cr::Expr {
                            ty: ir::Ty::text(),
                            kind: cr::ExprKind::From(cr::From::Serialize(Box::new(rel))),
                        }])
                    }
                    ir::TyKind::Tuple(_) => {
                        // non-constructed tuples need to be declared as a rel var and unpacked

                        ColumnsOrUnpack::Unpack(rel)
                    }
                    ir::TyKind::Enum(_) => {
                        let is_single_col = self.rel_cols_ty_nested(ty_mat).nth(1).is_none();
                        if is_single_col {
                            ColumnsOrUnpack::Columns(vec![rel])
                        } else {
                            ColumnsOrUnpack::Unpack(rel)
                        }
                    }
                    ir::TyKind::Function(_) => unreachable!(),
                }
            }
        }
    }

    /// Compiles `ir::Call { function: Pointer(External), args }` to cr::From(FuncCall { .. })
    fn compile_expr_std(&mut self, expr: &ir::Expr) -> cr::Expr {
        let ir::ExprKind::Call(call) = &expr.kind else {
            unreachable!()
        };

        let ir::ExprKind::Pointer(ir::Pointer::External(ptr)) = &call.function.kind else {
            unreachable!()
        };

        let args = call.args.iter().map(|x| self.compile_column(x)).collect();
        cr::Expr {
            kind: cr::ExprKind::From(cr::From::FuncCall(ptr.id.clone(), args)),
            ty: expr.ty.clone(),
        }
    }

    /// Ensures that an expression is a function - not just in type, but also in form.
    /// When it is not, it wraps the expressions into a new function that just calls it.
    ///
    /// Example: `ident` of type `func(a, b) -> c` is wrapped into
    ///          `func (x: a, y: b) -> ident(a, b)`
    fn as_function_or_wrap<'e>(&mut self, expr: &'e ir::Expr) -> Cow<'e, ir::Function> {
        if let ir::ExprKind::Function(func) = &expr.kind {
            return Cow::Borrowed(func.as_ref());
        }
        let expr = expr.clone();
        let ty_func = expr.ty.kind.as_function().unwrap().clone();

        let function_id = self.func_id_gen.next() as u32;

        Cow::Owned({
            ir::Function {
                id: function_id,
                body: ir::Expr {
                    ty: ty_func.body,
                    kind: ir::ExprKind::Call(Box::new(ir::Call {
                        function: expr.clone(),
                        args: ty_func
                            .params
                            .into_iter()
                            .enumerate()
                            .map(|(position, ty)| ir::Expr {
                                kind: ir::ExprKind::Pointer(ir::Pointer::Parameter(
                                    ir::ParameterPtr {
                                        function_id,
                                        param_position: position as u8,
                                    },
                                )),
                                ty,
                            })
                            .collect(),
                    })),
                },
            }
        })
    }

    /// Adds a `derived` expression for each row of `base`
    /// by joining them inside a `BindCorrelated`.
    fn correlated_join(
        &mut self,
        base: Box<cr::BoundExpr>,
        derived: cr::Expr,
    ) -> (Box<cr::BoundExpr>, cr::Expr, cr::Expr) {
        let base_ty = base.rel.ty.clone();
        let base_item_ty = {
            let mat = self.get_ty_mat(&base.rel.ty);
            mat.kind.as_array().unwrap().clone()
        };
        let derived_ty = derived.ty.clone();

        // Count columns for each part
        let base_nested = self.rel_cols_ty_nested(&base_item_ty).count();
        let base_col_count = 1 + base_nested; // index + item cols
        let derived_col_count = self.rel_cols_ty_nested(&derived_ty).count();

        // Construct enriched type: [{base_item, derived}]
        let enriched_item_ty = ty_concat_as_tuples((*base_item_ty).clone(), derived_ty.clone());
        let enriched_ty = ir::Ty::new(ir::TyKind::Array(Box::new(enriched_item_ty)));

        // Inner: for each row, join original data with derived value
        let enriched_inner = cr::Expr {
            kind: cr::ExprKind::Join(
                self.new_binding(cr::Expr::new_rel_ref(&base)),
                self.new_binding(derived),
                None,
            ),
            ty: enriched_ty.clone(),
        };

        // BindCorrelated produces all enriched rows
        let enriched = cr::Expr {
            kind: cr::ExprKind::BindCorrelated(base, Box::new(enriched_inner)),
            ty: enriched_ty,
        };
        let enriched_bound = self.new_binding(enriched);

        // Pick base columns from enriched
        let base_positions: Vec<usize> = (0..base_col_count).collect();
        let base_ref = cr::Expr {
            kind: cr::ExprKind::Transform(
                self.new_binding(cr::Expr::new_rel_ref(&enriched_bound)),
                cr::Transform::ProjectPick(base_positions),
            ),
            ty: base_ty,
        };

        // Pick derived columns from enriched
        let derived_positions: Vec<usize> =
            (base_col_count..base_col_count + derived_col_count).collect();
        let derived_ref = cr::Expr {
            kind: cr::ExprKind::Transform(
                self.new_binding(cr::Expr::new_rel_ref(&enriched_bound)),
                cr::Transform::ProjectPick(derived_positions),
            ),
            ty: derived_ty,
        };

        (enriched_bound, base_ref, derived_ref)
    }
}

#[derive(Debug)]
enum ColumnsOrUnpack {
    Columns(Vec<cr::Expr>),

    // A relation whose columns should be unpacked
    Unpack(cr::Expr),
}

impl ColumnsOrUnpack {
    #[track_caller]
    fn unwrap_columns(self) -> Vec<cr::Expr> {
        match self {
            ColumnsOrUnpack::Columns(cols) => cols,
            ColumnsOrUnpack::Unpack(rel) => panic!("expected columns, found: {rel:?}"),
        }
    }
}

/// Returns true if the key expression is simple enough to clone without
/// causing duplicate SQL compilation issues. Simple keys are chains of
/// ProjectPick/ProjectDiscard transforms on RelRef/Literal/Null nodes.
fn is_simple(expr: &cr::Expr) -> bool {
    match &expr.kind {
        cr::ExprKind::From(from) => matches!(
            from,
            cr::From::RelRef(_) | cr::From::Literal(_) | cr::From::Null
        ),
        cr::ExprKind::Transform(bound, transform) => {
            matches!(
                transform,
                cr::Transform::ProjectPick(_) | cr::Transform::ProjectDiscard(_)
            ) && is_simple(&bound.rel)
        }
        _ => false,
    }
}

fn new_bin_op(left: cr::Expr, op: &str, right: cr::Expr, ty: ir::Ty) -> cr::Expr {
    let kind = cr::ExprKind::From(cr::From::FuncCall(op.to_string(), vec![left, right]));
    cr::Expr { kind, ty }
}

fn new_un_op(op: &str, operand: cr::Expr, ty: ir::Ty) -> cr::Expr {
    let kind = cr::ExprKind::From(cr::From::FuncCall(op.to_string(), vec![operand]));
    cr::Expr { kind, ty }
}

fn new_index(index: i64) -> cr::Expr {
    let kind = cr::ExprKind::From(cr::From::Literal(ir::Literal::new_int64(index)));
    cr::Expr {
        kind,
        ty: ty_index(),
    }
}

fn ty_index() -> ir::Ty {
    ir::Ty::new(ir::TyPrimitive::Prim64)
}

/// Creates an int16 literal expression, used for enum tags.
/// Note: we use int16 instead of uint16 for PostgreSQL compatibility,
/// since PostgreSQL doesn't have unsigned types and uint16 would map to int4.
fn new_tag(tag: i16) -> cr::Expr {
    let kind = cr::ExprKind::From(cr::From::Literal(ir::Literal::new_int16(tag)));
    cr::Expr { kind, ty: ty_tag() }
}

fn ty_tag() -> ir::Ty {
    ir::Ty::new(ir::TyPrimitive::Prim16)
}

fn ty_concat_as_tuples(a: ir::Ty, b: ir::Ty) -> ir::Ty {
    let mut fields = Vec::new();
    if let ir::TyKind::Tuple(a) = a.kind {
        fields.extend(a);
    } else {
        fields.push(ir::TyTupleField { name: None, ty: a });
    }
    if let ir::TyKind::Tuple(b) = b.kind {
        fields.extend(b);
    } else {
        fields.push(ir::TyTupleField { name: None, ty: b });
    }
    ir::Ty::new(ir::TyKind::Tuple(fields))
}

fn ty_zip(a: ir::Ty, b: ir::Ty) -> ir::Ty {
    let mut fields = Vec::new();

    fn ty_tuple_field(ty: ir::Ty) -> ir::TyTupleField {
        ir::TyTupleField { name: None, ty }
    }

    let ir::TyKind::Array(a_item) = a.kind else {
        unreachable!()
    };
    fields.push(ty_tuple_field(*a_item));

    fields.push(ty_tuple_field(ty_index()));
    let ir::TyKind::Array(b_item) = b.kind else {
        unreachable!()
    };
    fields.push(ty_tuple_field(*b_item));

    let tuple = ir::Ty::new(ir::TyKind::Tuple(fields));

    ir::Ty::new(ir::TyKind::Array(Box::new(tuple)))
}