turso/core/translate/optimizer/access_method.rs

use std::sync::Arc;

use limbo_sqlite3_parser::ast::SortOrder;

use crate::{
    schema::Index,
    translate::plan::{IterationDirection, JoinOrderMember, TableReference},
    Result,
};

use super::{
    constraints::{usable_constraints_for_join_order, ConstraintRef, TableConstraints},
    cost::{estimate_cost_for_scan_or_seek, Cost, IndexInfo},
    order::OrderTarget,
};

#[derive(Debug, Clone)]
/// Represents a way to access a table.
pub struct AccessMethod<'a> {
    /// The estimated number of page fetches.
    /// We are ignoring CPU cost for now.
    pub cost: Cost,
    pub kind: AccessMethodKind<'a>,
}

impl<'a> AccessMethod<'a> {
    pub fn index(&self) -> Option<&Index> {
        match &self.kind {
            AccessMethodKind::Scan { index, .. } => index.as_ref().map(|i| i.as_ref()),
            AccessMethodKind::Search { index, .. } => index.as_ref().map(|i| i.as_ref()),
        }
    }
    pub fn iter_dir(&self) -> IterationDirection {
        match &self.kind {
            AccessMethodKind::Scan { iter_dir, .. } => *iter_dir,
            AccessMethodKind::Search { iter_dir, .. } => *iter_dir,
        }
    }
}

#[derive(Debug, Clone)]
/// Represents the kind of access method.
pub enum AccessMethodKind<'a> {
    /// A full scan, which can be an index scan or a table scan.
    Scan {
        index: Option<Arc<Index>>,
        iter_dir: IterationDirection,
    },
    /// A search, which can be an index seek or a rowid-based search.
    Search {
        index: Option<Arc<Index>>,
        iter_dir: IterationDirection,
        constraint_refs: &'a [ConstraintRef],
    },
}

/// Return the best [AccessMethod] for a given join order.
pub fn find_best_access_method_for_join_order<'a>(
    rhs_table: &TableReference,
    rhs_constraints: &'a TableConstraints,
    join_order: &[JoinOrderMember],
    maybe_order_target: Option<&OrderTarget>,
    input_cardinality: f64,
) -> Result<AccessMethod<'a>> {
    let table_no = join_order.last().unwrap().table_no;
    let cost_of_full_table_scan = estimate_cost_for_scan_or_seek(None, &[], &[], input_cardinality);
    let mut best_access_method = AccessMethod {
        cost: cost_of_full_table_scan,
        kind: AccessMethodKind::Scan {
            index: None,
            iter_dir: IterationDirection::Forwards,
        },
    };
    let rowid_column_idx = rhs_table.columns().iter().position(|c| c.is_rowid_alias);

    // Estimate cost for each candidate index (including the rowid index) and replace best_access_method if the cost is lower.
    for candidate in rhs_constraints.candidates.iter() {
        let index_info = match candidate.index.as_ref() {
            Some(index) => IndexInfo {
                unique: index.unique,
                covering: rhs_table.index_is_covering(index),
                column_count: index.columns.len(),
            },
            None => IndexInfo {
                unique: true, // rowids are always unique
                covering: false,
                column_count: 1,
            },
        };
        let usable_constraint_refs = usable_constraints_for_join_order(
            &rhs_constraints.constraints,
            &candidate.refs,
            join_order,
        );
        let cost = estimate_cost_for_scan_or_seek(
            Some(index_info),
            &rhs_constraints.constraints,
            &usable_constraint_refs,
            input_cardinality,
        );

        // All other things being equal, prefer an access method that satisfies the order target.
        let (iter_dir, order_satisfiability_bonus) = if let Some(order_target) = maybe_order_target
        {
            // If the index delivers rows in the same direction (or the exact reverse direction) as the order target, then it
            // satisfies the order target.
            let mut all_same_direction = true;
            let mut all_opposite_direction = true;
            for i in 0..order_target.0.len().min(index_info.column_count) {
                let correct_table = order_target.0[i].table_no == table_no;
                let correct_column = {
                    match &candidate.index {
                        Some(index) => index.columns[i].pos_in_table == order_target.0[i].column_no,
                        None => {
                            rowid_column_idx.map_or(false, |idx| idx == order_target.0[i].column_no)
                        }
                    }
                };
                if !correct_table || !correct_column {
                    all_same_direction = false;
                    all_opposite_direction = false;
                    break;
                }
                let correct_order = {
                    match &candidate.index {
                        Some(index) => order_target.0[i].order == index.columns[i].order,
                        None => order_target.0[i].order == SortOrder::Asc,
                    }
                };
                if correct_order {
                    all_opposite_direction = false;
                } else {
                    all_same_direction = false;
                }
            }
            if all_same_direction || all_opposite_direction {
                (
                    if all_same_direction {
                        IterationDirection::Forwards
                    } else {
                        IterationDirection::Backwards
                    },
                    Cost(1.0),
                )
            } else {
                (IterationDirection::Forwards, Cost(0.0))
            }
        } else {
            (IterationDirection::Forwards, Cost(0.0))
        };
        if cost < best_access_method.cost + order_satisfiability_bonus {
            best_access_method = AccessMethod {
                cost,
                kind: if usable_constraint_refs.is_empty() {
                    AccessMethodKind::Scan {
                        index: candidate.index.clone(),
                        iter_dir,
                    }
                } else {
                    AccessMethodKind::Search {
                        index: candidate.index.clone(),
                        iter_dir,
                        constraint_refs: &usable_constraint_refs,
                    }
                },
            };
        }
    }

    Ok(best_access_method)
}