// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Two-pass optimizer pipeline that pushes cheap expressions (like struct field
//! access `user['status']`) closer to data sources, enabling early data reduction
//! and source-level optimizations (e.g., Parquet column pruning). See
//! [`HamelinExtractLeafExpressions`] (pass 1) and [`HamelinPushDownLeafProjections`] (pass 2).
//!
//! # Why this is a Hamelin-local copy
//!
//! Forked from `datafusion-optimizer-53.0.0/src/extract_leaf_expressions.rs`
//! (Apache 2.0). The upstream rules emit invalid plans for several Hamelin
//! query shapes; these are the deltas vs. upstream. Each fix is annotated in
//! source with a `Hamelin fork:` comment for grep-ability.
//!
//! ## 1. `Unnest` barrier in `try_push_into_inputs` (line 1223)
//!
//! Upstream `try_push_into_inputs` will push an extraction projection past
//! `LogicalPlan::Unnest`. `Unnest` rewrites a column's type from `List<T>` to
//! `T`, so the rebuilt projection's `get_field(col, …)` fails type-checking
//! against the input schema where `col` is still `List<Struct>`. Fix: refuse
//! the pushdown when `node` is `LogicalPlan::Unnest(_)`.
//!
//! ## 2. Pass-through dedup by output field name (line 717)
//!
//! When merging into an existing Projection, upstream
//! `build_extraction_projection_impl` resolves a needed column through the
//! projection's rename map and appends it as an unaliased `Expr::Column`. If
//! the projection already provides that column under an alias (e.g.
//! `t.id AS id`), the merged schema ends up with both qualified `t.id` and
//! unqualified `id`, which DF 53's stricter `DFSchema` validation rejects as
//! ambiguous. Fix: track the existing projection's output field names and
//! skip pass-through additions whose unqualified name collides.
//!
//! ## 3. Refuse to push past a non-merge Projection (line 1111)
//!
//! Upstream `push_extraction_pairs` falls through to
//! `try_push_into_inputs(proj_input)` whenever the merge precondition
//! (`proj_exprs_captured == proj.expr.len()`) doesn't hold. When `proj_input`
//! is itself a `Projection` that redefines a column (e.g. an `EXPLODE` that
//! adds a new struct field), `try_push_into_inputs` routes the pairs into
//! the inner projection's input *without* translating them through the
//! intermediate projection's rename map — so an extraction like
//! `get_field(x, "c")` ends up referencing an `x` that lacks `c`, producing
//! `Field c not found in struct`. Fix: refuse the pushdown when
//! `proj_input` is `LogicalPlan::Projection(_)`; the extraction projection
//! stays one level higher (correct, just less optimal).
//!
//! ## Other deltas
//!
//! - `has_all_column_refs` (line 93) is inlined from
//!   `datafusion_optimizer::utils`, which marks it `pub(crate)`.
//! - The two rule structs are renamed `HamelinExtractLeafExpressions` /
//!   `HamelinPushDownLeafProjections` and report names
//!   `hamelin_extract_leaf_expressions` /
//!   `hamelin_push_down_leaf_projections` so they don't collide with the
//!   upstream defaults — the [`crate::optimizer::with_hamelin_optimizer_rules`]
//!   helper installs these in place of the upstream pair.
//! - Tests from the upstream file are dropped; coverage lives in
//!   `hamelin_it/` against real backends.

use indexmap::{IndexMap, IndexSet};
use std::collections::{HashMap, HashSet};
use std::sync::Arc;

use datafusion_common::alias::AliasGenerator;
use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
use datafusion_common::{qualified_name, Column, DFSchema, Result};
use datafusion_expr::logical_plan::LogicalPlan;
use datafusion_expr::{Expr, ExpressionPlacement, Projection};
use datafusion_optimizer::optimizer::ApplyOrder;
use datafusion_optimizer::push_down_filter::replace_cols_by_name;
use datafusion_optimizer::{OptimizerConfig, OptimizerRule};

/// Returns true if `expr` contains all columns in `schema_cols`.
///
/// Inlined from upstream `datafusion_optimizer::utils::has_all_column_refs`,
/// which is `pub(crate)`.
fn has_all_column_refs(expr: &Expr, schema_cols: &HashSet<Column>) -> bool {
    let column_refs = expr.column_refs();
    schema_cols
        .iter()
        .filter(|c| column_refs.contains(c))
        .count()
        == column_refs.len()
}

/// Prefix for aliases generated by the extraction optimizer passes.
///
/// This prefix is **reserved for internal optimizer use**. User-defined aliases
/// starting with this prefix may be misidentified as optimizer-generated
/// extraction aliases, leading to unexpected behavior. Do not use this prefix
/// in user queries.
const EXTRACTED_EXPR_PREFIX: &str = "__datafusion_extracted";

/// Returns `true` if any sub-expression in `exprs` has
/// [`ExpressionPlacement::MoveTowardsLeafNodes`] placement.
///
/// This is a lightweight pre-check that short-circuits as soon as one
/// extractable expression is found, avoiding the expensive allocations
/// (column HashSets, extractors, expression rewrites) that the full
/// extraction pipeline requires.
fn has_extractable_expr(exprs: &[Expr]) -> bool {
    exprs.iter().any(|expr| {
        expr.exists(|e| Ok(e.placement() == ExpressionPlacement::MoveTowardsLeafNodes))
            .unwrap_or(false)
    })
}

/// Extracts `MoveTowardsLeafNodes` sub-expressions from non-projection nodes
/// into **extraction projections** (pass 1 of 2).
///
/// This handles Filter, Sort, Limit, Aggregate, and Join nodes. For Projection
/// nodes, extraction and pushdown are handled by [`HamelinPushDownLeafProjections`].
///
/// # Key Concepts
///
/// **Extraction projection**: a projection inserted *below* a node that
/// pre-computes a cheap expression and exposes it under an alias
/// (`__datafusion_extracted_N`). The parent node then references the alias
/// instead of the original expression.
///
/// **Recovery projection**: a projection inserted *above* a node to restore
/// the original output schema when extraction changes it.
/// Schema-preserving nodes (Filter, Sort, Limit) gain extra columns from
/// the extraction projection that bubble up; the recovery projection selects
/// only the original columns to hide the extras.
///
/// # Example
///
/// Given a filter with a struct field access:
///
/// ```text
/// Filter: user['status'] = 'active'
///   TableScan: t [id, user]
/// ```
///
/// This rule:
/// 1. Inserts an **extraction projection** below the filter:
/// 2. Adds a **recovery projection** above to hide the extra column:
///
/// ```text
/// Projection: id, user                                                        <-- recovery projection
///   Filter: __datafusion_extracted_1 = 'active'
///     Projection: user['status'] AS __datafusion_extracted_1, id, user         <-- extraction projection
///       TableScan: t [id, user]
/// ```
///
/// **Important:** The `PushDownFilter` rule is aware of projections created by this rule
/// and will not push filters through them. It uses `ExpressionPlacement` to detect
/// `MoveTowardsLeafNodes` expressions and skip filter pushdown past them.
#[derive(Default, Debug)]
pub struct HamelinExtractLeafExpressions {}

impl HamelinExtractLeafExpressions {
    /// Create a new [`HamelinExtractLeafExpressions`]
    pub fn new() -> Self {
        Self {}
    }
}

impl OptimizerRule for HamelinExtractLeafExpressions {
    fn name(&self) -> &str {
        "hamelin_extract_leaf_expressions"
    }

    fn rewrite(
        &self,
        plan: LogicalPlan,
        config: &dyn OptimizerConfig,
    ) -> Result<Transformed<LogicalPlan>> {
        if !config.options().optimizer.enable_leaf_expression_pushdown {
            return Ok(Transformed::no(plan));
        }
        let alias_generator = config.alias_generator();

        // Advance the alias generator past any user-provided __datafusion_extracted_N
        // aliases to prevent collisions when generating new extraction aliases.
        advance_generator_past_existing(&plan, alias_generator)?;

        plan.transform_down_with_subqueries(|plan| extract_from_plan(plan, alias_generator))
    }
}

/// Scans the current plan node's expressions for pre-existing
/// `__datafusion_extracted_N` aliases and advances the generator
/// counter past them to avoid collisions with user-provided aliases.
fn advance_generator_past_existing(
    plan: &LogicalPlan,
    alias_generator: &AliasGenerator,
) -> Result<()> {
    plan.apply(|plan| {
        plan.expressions().iter().try_for_each(|expr| {
            expr.apply(|e| {
                if let Expr::Alias(alias) = e {
                    if let Some(id) = alias
                        .name
                        .strip_prefix(EXTRACTED_EXPR_PREFIX)
                        .and_then(|s| s.strip_prefix('_'))
                        .and_then(|s| s.parse().ok())
                    {
                        alias_generator.update_min_id(id);
                    }
                }
                Ok(TreeNodeRecursion::Continue)
            })?;
            Ok::<(), datafusion_common::error::DataFusionError>(())
        })?;
        Ok(TreeNodeRecursion::Continue)
    })
    .map(|_| ())
}

/// Extracts `MoveTowardsLeafNodes` sub-expressions from a plan node.
///
/// Works for any number of inputs (0, 1, 2, …N). For multi-input nodes
/// like Join, each extracted sub-expression is routed to the correct input
/// by checking which input's schema contains all of the expression's column
/// references.
fn extract_from_plan(
    plan: LogicalPlan,
    alias_generator: &Arc<AliasGenerator>,
) -> Result<Transformed<LogicalPlan>> {
    // Only extract from plan types whose output schema is predictable after
    // expression rewriting.  Nodes like Window derive column names from
    // their expressions, so rewriting `get_field` inside a window function
    // changes the output schema and breaks the recovery projection.
    if !matches!(
        &plan,
        LogicalPlan::Aggregate(_)
            | LogicalPlan::Filter(_)
            | LogicalPlan::Sort(_)
            | LogicalPlan::Limit(_)
            | LogicalPlan::Join(_)
    ) {
        return Ok(Transformed::no(plan));
    }

    let inputs = plan.inputs();
    if inputs.is_empty() {
        return Ok(Transformed::no(plan));
    }

    // Fast pre-check: skip all allocations if no extractable expressions exist
    if !has_extractable_expr(&plan.expressions()) {
        return Ok(Transformed::no(plan));
    }

    // Save original output schema before any transformation
    let original_schema = Arc::clone(plan.schema());

    // Build per-input schemas from borrowed inputs (before plan is consumed
    // by map_expressions). We only need schemas and column sets for routing;
    // the actual inputs are cloned later only if extraction succeeds.
    let input_schemas: Vec<Arc<DFSchema>> = inputs.iter().map(|i| Arc::clone(i.schema())).collect();

    // Build per-input extractors
    let mut extractors: Vec<LeafExpressionExtractor> = input_schemas
        .iter()
        .map(|schema| LeafExpressionExtractor::new(schema.as_ref(), alias_generator))
        .collect();

    // Build per-input column sets for routing expressions to the correct input
    let input_column_sets: Vec<std::collections::HashSet<Column>> = input_schemas
        .iter()
        .map(|schema| schema_columns(schema.as_ref()))
        .collect();

    // Transform expressions via map_expressions with routing
    let transformed =
        plan.map_expressions(|expr| routing_extract(expr, &mut extractors, &input_column_sets))?;

    // If no expressions were rewritten, nothing was extracted
    if !transformed.transformed {
        return Ok(transformed);
    }

    // Clone inputs now that we know extraction succeeded. Wrap in Arc
    // upfront since build_extraction_projection expects &Arc<LogicalPlan>.
    let owned_inputs: Vec<Arc<LogicalPlan>> = transformed
        .data
        .inputs()
        .into_iter()
        .map(|i| Arc::new(i.clone()))
        .collect();

    // Build per-input extraction projections (None means no extractions for that input)
    let new_inputs: Vec<LogicalPlan> = owned_inputs
        .into_iter()
        .zip(extractors.iter())
        .map(|(input_arc, extractor)| {
            match extractor.build_extraction_projection(&input_arc)? {
                Some(plan) => Ok(plan),
                // No extractions for this input — recover the LogicalPlan
                // without cloning (refcount is 1 since build returned None).
                None => Ok(Arc::try_unwrap(input_arc).unwrap_or_else(|arc| (*arc).clone())),
            }
        })
        .collect::<Result<Vec<_>>>()?;

    // Rebuild the plan keeping its rewritten expressions but replacing
    // inputs with the new extraction projections.
    let new_plan = transformed
        .data
        .with_new_exprs(transformed.data.expressions(), new_inputs)?;

    // Add recovery projection if the output schema changed
    let recovered = build_recovery_projection(original_schema.as_ref(), new_plan)?;

    Ok(Transformed::yes(recovered))
}

/// Given an expression, returns the index of the input whose columns fully
/// cover the expression's column references.
/// Returns `None` if the expression references columns from multiple inputs
/// or if multiple inputs match (ambiguous, e.g. unqualified columns present
/// in both sides of a join).
fn find_owning_input(
    expr: &Expr,
    input_column_sets: &[std::collections::HashSet<Column>],
) -> Option<usize> {
    let mut found = None;
    for (idx, cols) in input_column_sets.iter().enumerate() {
        if has_all_column_refs(expr, cols) {
            if found.is_some() {
                // Ambiguous — multiple inputs match
                return None;
            }
            found = Some(idx);
        }
    }
    found
}

/// Walks an expression tree top-down, extracting `MoveTowardsLeafNodes`
/// sub-expressions and routing each to the correct per-input extractor.
fn routing_extract(
    expr: Expr,
    extractors: &mut [LeafExpressionExtractor],
    input_column_sets: &[std::collections::HashSet<Column>],
) -> Result<Transformed<Expr>> {
    expr.transform_down(|e| {
        // Skip expressions already aliased with extracted expression pattern
        if let Expr::Alias(alias) = &e {
            if alias.name.starts_with(EXTRACTED_EXPR_PREFIX) {
                return Ok(Transformed {
                    data: e,
                    transformed: false,
                    tnr: TreeNodeRecursion::Jump,
                });
            }
        }

        // Don't extract Alias nodes directly — preserve the alias and let
        // transform_down recurse into the inner expression
        if matches!(&e, Expr::Alias(_)) {
            return Ok(Transformed::no(e));
        }

        match e.placement() {
            ExpressionPlacement::MoveTowardsLeafNodes => {
                if let Some(idx) = find_owning_input(&e, input_column_sets) {
                    let col_ref = extractors[idx].add_extracted(e)?;
                    Ok(Transformed::yes(col_ref))
                } else {
                    // References columns from multiple inputs — cannot extract
                    Ok(Transformed::no(e))
                }
            }
            ExpressionPlacement::Column => {
                // Track columns that the parent node references so the
                // extraction projection includes them as pass-through.
                // Without this, the extraction projection would only
                // contain __datafusion_extracted_N aliases, and the parent couldn't
                // resolve its other column references.
                if let Expr::Column(col) = &e {
                    if let Some(idx) = find_owning_input(&e, input_column_sets) {
                        extractors[idx].columns_needed.insert(col.clone());
                    }
                }
                Ok(Transformed::no(e))
            }
            _ => Ok(Transformed::no(e)),
        }
    })
}

/// Returns all columns in the schema (both qualified and unqualified forms)
fn schema_columns(schema: &DFSchema) -> std::collections::HashSet<Column> {
    schema
        .iter()
        .flat_map(|(qualifier, field)| {
            [
                Column::new(qualifier.cloned(), field.name()),
                Column::new_unqualified(field.name()),
            ]
        })
        .collect()
}

/// Rewrites extraction pairs and column references from one qualifier
/// space to another.
///
/// Builds a replacement map by zipping `from_schema` (whose qualifiers
/// currently appear in `pairs` / `columns`) with `to_schema` (the
/// qualifiers we want), then applies `replace_cols_by_name`.
///
/// Used for SubqueryAlias (alias-space -> input-space) and Union
/// (union output-space -> per-branch input-space).
fn remap_pairs_and_columns(
    pairs: &[(Expr, String)],
    columns: &IndexSet<Column>,
    from_schema: &DFSchema,
    to_schema: &DFSchema,
) -> Result<ExtractionTarget> {
    let mut replace_map = HashMap::new();
    for ((from_q, from_f), (to_q, to_f)) in from_schema.iter().zip(to_schema.iter()) {
        replace_map.insert(
            qualified_name(from_q, from_f.name()),
            Expr::Column(Column::new(to_q.cloned(), to_f.name())),
        );
    }
    let remapped_pairs: Vec<(Expr, String)> = pairs
        .iter()
        .map(|(expr, alias)| {
            Ok((
                replace_cols_by_name(expr.clone(), &replace_map)?,
                alias.clone(),
            ))
        })
        .collect::<Result<_>>()?;
    let remapped_columns: IndexSet<Column> = columns
        .iter()
        .filter_map(|col| {
            let rewritten = replace_cols_by_name(Expr::Column(col.clone()), &replace_map).ok()?;
            if let Expr::Column(c) = rewritten {
                Some(c)
            } else {
                Some(col.clone())
            }
        })
        .collect();
    Ok(ExtractionTarget {
        pairs: remapped_pairs,
        columns: remapped_columns,
    })
}

// =============================================================================
// Helper Types & Functions for Extraction Targeting
// =============================================================================

/// A bundle of extraction pairs (expression + alias) and standalone columns
/// that need to be pushed through a plan node.
struct ExtractionTarget {
    /// Extracted expressions paired with their generated aliases.
    pairs: Vec<(Expr, String)>,
    /// Standalone column references needed by the parent node.
    columns: IndexSet<Column>,
}

/// Build a replacement map from a projection: output_column_name -> underlying_expr.
///
/// This is used to resolve column references through a renaming projection.
/// For example, if a projection has `user AS x`, this maps `x` -> `col("user")`.
fn build_projection_replace_map(projection: &Projection) -> HashMap<String, Expr> {
    projection
        .schema
        .iter()
        .zip(projection.expr.iter())
        .map(|((qualifier, field), expr)| {
            let key = Column::from((qualifier, field)).flat_name();
            (key, expr.clone().unalias())
        })
        .collect()
}

/// Build a recovery projection to restore the original output schema.
///
/// After extraction, a node's output schema may differ from the original:
///
/// - **Schema-preserving nodes** (Filter/Sort/Limit): the extraction projection
///   below adds extra `__datafusion_extracted_N` columns that bubble up through
///   the node. Recovery selects only the original columns to hide the extras.
///   ```text
///   Original schema: [id, user]
///   After extraction: [__datafusion_extracted_1, id, user]   ← extra column leaked through
///   Recovery: SELECT id, user FROM ...                       ← hides __datafusion_extracted_1
///   ```
///
/// - **Schema-defining nodes** (Aggregate): same number of columns but names
///   may differ because extracted aliases replaced the original expressions.
///   Recovery maps positionally, aliasing where names changed.
///   ```text
///   Original: [SUM(user['balance'])]
///   After:    [SUM(__datafusion_extracted_1)]                ← name changed
///   Recovery: SUM(__datafusion_extracted_1) AS "SUM(user['balance'])"
///   ```
///
/// - **Schemas identical** → no recovery projection needed.
fn build_recovery_projection(
    original_schema: &DFSchema,
    input: LogicalPlan,
) -> Result<LogicalPlan> {
    let new_schema = input.schema();
    let orig_len = original_schema.fields().len();
    let new_len = new_schema.fields().len();

    if orig_len == new_len {
        // Same number of fields — check if schemas are identical
        let schemas_match = original_schema.iter().zip(new_schema.iter()).all(
            |((orig_q, orig_f), (new_q, new_f))| orig_f.name() == new_f.name() && orig_q == new_q,
        );
        if schemas_match {
            return Ok(input);
        }

        // Schema-defining nodes (Aggregate, Join): names may differ at some
        // positions because extracted aliases replaced the original expressions.
        // Map positionally, aliasing where the name changed.
        //
        // Invariant: `with_new_exprs` on all supported node types (Aggregate,
        // Filter, Sort, Limit, Join) preserves column order, so positional
        // mapping is safe here.
        debug_assert!(
            orig_len == new_len,
            "build_recovery_projection: positional mapping requires same field count, \
             got original={orig_len} vs new={new_len}"
        );
        let mut proj_exprs = Vec::with_capacity(orig_len);
        for (i, (orig_qualifier, orig_field)) in original_schema.iter().enumerate() {
            let (new_qualifier, new_field) = new_schema.qualified_field(i);
            if orig_field.name() == new_field.name() && orig_qualifier == new_qualifier {
                proj_exprs.push(Expr::from((orig_qualifier, orig_field)));
            } else {
                let new_col = Expr::Column(Column::from((new_qualifier, new_field)));
                proj_exprs
                    .push(new_col.alias_qualified(orig_qualifier.cloned(), orig_field.name()));
            }
        }
        let projection = Projection::try_new(proj_exprs, Arc::new(input))?;
        Ok(LogicalPlan::Projection(projection))
    } else {
        // Schema-preserving nodes: new schema has extra extraction columns.
        // Original columns still exist by name; select them to hide extras.
        let col_exprs: Vec<Expr> = original_schema.iter().map(Expr::from).collect();
        let projection = Projection::try_new(col_exprs, Arc::new(input))?;
        Ok(LogicalPlan::Projection(projection))
    }
}

/// Collects `MoveTowardsLeafNodes` sub-expressions found during expression
/// tree traversal and can build an extraction projection from them.
///
/// # Example
///
/// Given `Filter: user['status'] = 'active' AND user['name'] IS NOT NULL`:
/// - `add_extracted(user['status'])` → stores it, returns `col("__datafusion_extracted_1")`
/// - `add_extracted(user['name'])`   → stores it, returns `col("__datafusion_extracted_2")`
/// - `build_extraction_projection()` produces:
///   `Projection: user['status'] AS __datafusion_extracted_1, user['name'] AS __datafusion_extracted_2, <all input columns>`
struct LeafExpressionExtractor<'a> {
    /// Extracted expressions: maps expression -> alias
    extracted: IndexMap<Expr, String>,
    /// Columns referenced by extracted expressions or the parent node,
    /// included as pass-through in the extraction projection.
    columns_needed: IndexSet<Column>,
    /// Input schema
    input_schema: &'a DFSchema,
    /// Alias generator
    alias_generator: &'a Arc<AliasGenerator>,
}

impl<'a> LeafExpressionExtractor<'a> {
    fn new(input_schema: &'a DFSchema, alias_generator: &'a Arc<AliasGenerator>) -> Self {
        Self {
            extracted: IndexMap::new(),
            columns_needed: IndexSet::new(),
            input_schema,
            alias_generator,
        }
    }

    /// Adds an expression to extracted set, returns column reference.
    fn add_extracted(&mut self, expr: Expr) -> Result<Expr> {
        // Deduplication: reuse existing alias if same expression
        if let Some(alias) = self.extracted.get(&expr) {
            return Ok(Expr::Column(Column::new_unqualified(alias)));
        }

        // Track columns referenced by this expression
        for col in expr.column_refs() {
            self.columns_needed.insert(col.clone());
        }

        // Generate unique alias
        let alias = self.alias_generator.next(EXTRACTED_EXPR_PREFIX);
        self.extracted.insert(expr, alias.clone());

        Ok(Expr::Column(Column::new_unqualified(&alias)))
    }

    /// Builds an extraction projection above the given input, or merges into
    /// it if the input is already a projection. Delegates to
    /// [`build_extraction_projection_impl`].
    ///
    /// Returns `None` if there are no extractions.
    fn build_extraction_projection(&self, input: &Arc<LogicalPlan>) -> Result<Option<LogicalPlan>> {
        if self.extracted.is_empty() {
            return Ok(None);
        }
        let pairs: Vec<(Expr, String)> = self
            .extracted
            .iter()
            .map(|(e, a)| (e.clone(), a.clone()))
            .collect();
        let proj = build_extraction_projection_impl(
            &pairs,
            &self.columns_needed,
            input,
            self.input_schema,
        )?;
        Ok(Some(LogicalPlan::Projection(proj)))
    }
}

/// Build an extraction projection above the target node (shared by both passes).
///
/// If the target is an existing projection, merges into it. This requires
/// resolving column references through the projection's rename mapping:
/// if the projection has `user AS u`, and an extracted expression references
/// `u['name']`, we must rewrite it to `user['name']` since the merged
/// projection reads from the same input as the original.
///
/// Deduplicates by resolved expression equality and adds pass-through
/// columns as needed. Otherwise builds a fresh projection with extracted
/// expressions + ALL input schema columns.
fn build_extraction_projection_impl(
    extracted_exprs: &[(Expr, String)],
    columns_needed: &IndexSet<Column>,
    target: &Arc<LogicalPlan>,
    target_schema: &DFSchema,
) -> Result<Projection> {
    if let LogicalPlan::Projection(existing) = target.as_ref() {
        // Merge into existing projection
        let mut proj_exprs = existing.expr.clone();

        // Build a map of existing expressions (by Expr equality) to their aliases
        let existing_extractions: IndexMap<Expr, String> = existing
            .expr
            .iter()
            .filter_map(|e| {
                if let Expr::Alias(alias) = e {
                    if alias.name.starts_with(EXTRACTED_EXPR_PREFIX) {
                        return Some((*alias.expr.clone(), alias.name.clone()));
                    }
                }
                None
            })
            .collect();

        // Resolve column references through the projection's rename mapping
        let replace_map = build_projection_replace_map(existing);

        // Add new extracted expressions, resolving column refs through the projection
        for (expr, alias) in extracted_exprs {
            let resolved = replace_cols_by_name(expr.clone().alias(alias), &replace_map)?;
            let resolved_inner = if let Expr::Alias(a) = &resolved {
                a.expr.as_ref()
            } else {
                &resolved
            };
            if let Some(existing_alias) = existing_extractions.get(resolved_inner) {
                // Same expression already extracted under a different alias —
                // add the expression with the new alias so both names are
                // available in the output. We can't reference the existing alias
                // as a column within the same projection, so we duplicate the
                // computation.
                if existing_alias != alias {
                    proj_exprs.push(resolved);
                }
            } else {
                proj_exprs.push(resolved);
            }
        }

        // Add any new pass-through columns that aren't already in the projection.
        // We check against existing.input.schema() (the projection's source) rather
        // than target_schema (the projection's output) because columns produced
        // by alias expressions (e.g., CSE's __common_expr_N) exist in the output but
        // not the input, and cannot be added as pass-through Column references.
        let existing_cols: IndexSet<Column> = existing
            .expr
            .iter()
            .filter_map(|e| {
                if let Expr::Column(c) = e {
                    Some(c.clone())
                } else {
                    None
                }
            })
            .collect();

        // Hamelin fork: track output field names the existing projection already
        // provides. Skip pass-through columns whose unqualified name collides
        // with one — adding `t.id` next to `t.id AS id` would produce a schema
        // with both qualified `t.id` and unqualified `id`, which DF 53 rejects
        // as ambiguous.
        let existing_output_names: HashSet<String> = existing
            .schema
            .iter()
            .map(|(_, f)| f.name().to_string())
            .collect();

        let input_schema = existing.input.schema();
        for col in columns_needed {
            let col_expr = Expr::Column(col.clone());
            let resolved = replace_cols_by_name(col_expr, &replace_map)?;
            if let Expr::Column(resolved_col) = &resolved {
                if !existing_cols.contains(resolved_col)
                    && input_schema.has_column(resolved_col)
                    && !existing_output_names.contains(&resolved_col.name)
                {
                    proj_exprs.push(Expr::Column(resolved_col.clone()));
                }
            }
            // If resolved to non-column expr, it's already computed by existing projection
        }

        Projection::try_new(proj_exprs, Arc::clone(&existing.input))
    } else {
        // Build new projection with extracted expressions + all input columns
        let mut proj_exprs = Vec::new();
        for (expr, alias) in extracted_exprs {
            proj_exprs.push(expr.clone().alias(alias));
        }
        for (qualifier, field) in target_schema.iter() {
            proj_exprs.push(Expr::from((qualifier, field)));
        }
        Projection::try_new(proj_exprs, Arc::clone(target))
    }
}

// =============================================================================
// Pass 2: HamelinPushDownLeafProjections
// =============================================================================

/// Pushes extraction projections down through schema-preserving nodes towards
/// leaf nodes (pass 2 of 2, after [`HamelinExtractLeafExpressions`]).
///
/// Handles two types of projections:
/// - **Pure extraction projections** (all `__datafusion_extracted` aliases + columns):
///   pushes through Filter/Sort/Limit, merges into existing projections, or routes
///   into multi-input node inputs (Join, SubqueryAlias, etc.)
/// - **Mixed projections** (user projections containing `MoveTowardsLeafNodes`
///   sub-expressions): splits into a recovery projection + extraction projection,
///   then pushes the extraction projection down.
///
/// # Example: Pushing through a Filter
///
/// After pass 1, the extraction projection sits directly below the filter:
/// ```text
/// Projection: id, user                                                              <-- recovery
///   Filter: __datafusion_extracted_1 = 'active'
///     Projection: user['status'] AS __datafusion_extracted_1, id, user               <-- extraction
///       TableScan: t [id, user]
/// ```
///
/// Pass 2 pushes the extraction projection through the recovery and filter,
/// and a subsequent `OptimizeProjections` pass removes the (now-redundant)
/// recovery projection:
/// ```text
/// Filter: __datafusion_extracted_1 = 'active'
///   Projection: user['status'] AS __datafusion_extracted_1, id, user                 <-- extraction (pushed down)
///     TableScan: t [id, user]
/// ```
#[derive(Default, Debug)]
pub struct HamelinPushDownLeafProjections {}

impl HamelinPushDownLeafProjections {
    pub fn new() -> Self {
        Self {}
    }
}

impl OptimizerRule for HamelinPushDownLeafProjections {
    fn name(&self) -> &str {
        "hamelin_push_down_leaf_projections"
    }

    fn apply_order(&self) -> Option<ApplyOrder> {
        Some(ApplyOrder::TopDown)
    }

    fn rewrite(
        &self,
        plan: LogicalPlan,
        config: &dyn OptimizerConfig,
    ) -> Result<Transformed<LogicalPlan>> {
        if !config.options().optimizer.enable_leaf_expression_pushdown {
            return Ok(Transformed::no(plan));
        }
        let alias_generator = config.alias_generator();
        match try_push_input(&plan, alias_generator)? {
            Some(new_plan) => Ok(Transformed::yes(new_plan)),
            None => Ok(Transformed::no(plan)),
        }
    }
}

/// Attempts to push a projection's extractable expressions further down.
///
/// Returns `Some(new_subtree)` if the projection was pushed down or merged,
/// `None` if there is nothing to push or the projection sits above a barrier.
fn try_push_input(
    input: &LogicalPlan,
    alias_generator: &Arc<AliasGenerator>,
) -> Result<Option<LogicalPlan>> {
    let LogicalPlan::Projection(proj) = input else {
        return Ok(None);
    };
    split_and_push_projection(proj, alias_generator)
}

/// Splits a projection into extractable pieces, pushes them towards leaf
/// nodes, and adds a recovery projection if needed.
///
/// Handles both:
/// - **Pure extraction projections** (all `__datafusion_extracted` aliases + columns)
/// - **Mixed projections** (containing `MoveTowardsLeafNodes` sub-expressions)
///
/// Returns `Some(new_subtree)` if extractions were pushed down,
/// `None` if there is nothing to extract or push.
///
/// # Example: Mixed Projection
///
/// ```text
/// Input plan:
///   Projection: user['name'] IS NOT NULL AS has_name, id
///     Filter: ...
///       TableScan
///
/// Phase 1 (Split):
///   extraction_pairs: [(user['name'], "__datafusion_extracted_1")]
///   recovery_exprs:   [__datafusion_extracted_1 IS NOT NULL AS has_name, id]
///
/// Phase 2 (Push):
///   Push extraction projection through Filter toward TableScan
///
/// Phase 3 (Recovery):
///   Projection: __datafusion_extracted_1 IS NOT NULL AS has_name, id       <-- recovery
///     Filter: ...
///       Projection: user['name'] AS __datafusion_extracted_1, id           <-- extraction (pushed)
///         TableScan
/// ```
fn split_and_push_projection(
    proj: &Projection,
    alias_generator: &Arc<AliasGenerator>,
) -> Result<Option<LogicalPlan>> {
    // Fast pre-check: skip if there are no pre-existing extracted aliases
    // and no new extractable expressions.
    let has_existing_extracted = proj
        .expr
        .iter()
        .any(|e| matches!(e, Expr::Alias(alias) if alias.name.starts_with(EXTRACTED_EXPR_PREFIX)));
    if !has_existing_extracted && !has_extractable_expr(&proj.expr) {
        return Ok(None);
    }

    let input = &proj.input;
    let input_schema = input.schema();

    // ── Phase 1: Split ──────────────────────────────────────────────────
    // For each projection expression, collect extraction pairs and build
    // recovery expressions.
    //
    // Pre-existing `__datafusion_extracted` aliases are inserted into the
    // extractor's `IndexMap` with the **full** `Expr::Alias(…)` as the key,
    // so the alias name participates in equality. This prevents collisions
    // when CSE rewrites produce the same inner expression under different
    // alias names (e.g. `__common_expr_4 AS __datafusion_extracted_1` and
    // `__common_expr_4 AS __datafusion_extracted_3`). New extractions from
    // `routing_extract` use bare (non-Alias) keys and get normal dedup.
    //
    // When building the final `extraction_pairs`, the Alias wrapper is
    // stripped so consumers see the usual `(inner_expr, alias_name)` tuples.

    let mut extractors = vec![LeafExpressionExtractor::new(
        input_schema.as_ref(),
        alias_generator,
    )];
    let input_column_sets = vec![schema_columns(input_schema.as_ref())];

    let original_schema = proj.schema.as_ref();
    let mut recovery_exprs: Vec<Expr> = Vec::with_capacity(proj.expr.len());
    let mut needs_recovery = false;
    let mut has_new_extractions = false;
    let mut proj_exprs_captured: usize = 0;
    // Track standalone column expressions (Case B) to detect column refs
    // from extracted aliases (Case A) that aren't also standalone expressions.
    let mut standalone_columns: IndexSet<Column> = IndexSet::new();

    for (expr, (qualifier, field)) in proj.expr.iter().zip(original_schema.iter()) {
        let extracted_alias = if let Expr::Alias(alias) = expr {
            alias
                .name
                .starts_with(EXTRACTED_EXPR_PREFIX)
                .then_some(alias)
        } else {
            None
        };
        if let Some(alias) = extracted_alias {
            // Insert the full Alias expression as the key so that
            // distinct alias names don't collide in the IndexMap.
            let alias_name = alias.name.clone();

            for col_ref in alias.expr.column_refs() {
                extractors[0].columns_needed.insert(col_ref.clone());
            }

            extractors[0]
                .extracted
                .insert(expr.clone(), alias_name.clone());
            recovery_exprs.push(Expr::Column(Column::new_unqualified(&alias_name)));
            proj_exprs_captured += 1;
        } else if let Expr::Column(col) = expr {
            // Plain column pass-through — track it in the extractor
            extractors[0].columns_needed.insert(col.clone());
            standalone_columns.insert(col.clone());
            recovery_exprs.push(expr.clone());
            proj_exprs_captured += 1;
        } else {
            // Everything else: run through routing_extract
            let transformed = routing_extract(expr.clone(), &mut extractors, &input_column_sets)?;
            if transformed.transformed {
                has_new_extractions = true;
            }
            let transformed_expr = transformed.data;

            // Build recovery expression, aliasing back to original name if needed
            let original_name = field.name();
            let needs_alias = if let Expr::Column(col) = &transformed_expr {
                col.name.as_str() != original_name
            } else {
                let expr_name = transformed_expr.schema_name().to_string();
                original_name != &expr_name
            };
            let recovery_expr = if needs_alias {
                needs_recovery = true;
                transformed_expr
                    .clone()
                    .alias_qualified(qualifier.cloned(), original_name)
            } else {
                transformed_expr.clone()
            };

            // If the expression was transformed (i.e., has extracted sub-parts),
            // it differs from what the pushed projection outputs → needs recovery.
            // Also, any non-column, non-__datafusion_extracted expression needs recovery
            // because the pushed extraction projection won't output it directly.
            if transformed.transformed || !matches!(expr, Expr::Column(_)) {
                needs_recovery = true;
            }

            recovery_exprs.push(recovery_expr);
        }
    }

    // Build extraction_pairs, stripping the Alias wrapper from pre-existing
    // entries (they used the full Alias as the map key to avoid dedup).
    let extractor = &extractors[0];
    let extraction_pairs: Vec<(Expr, String)> = extractor
        .extracted
        .iter()
        .map(|(e, a)| match e {
            Expr::Alias(alias) => (*alias.expr.clone(), a.clone()),
            _ => (e.clone(), a.clone()),
        })
        .collect();
    let columns_needed = &extractor.columns_needed;

    // If no extractions found, nothing to do
    if extraction_pairs.is_empty() {
        return Ok(None);
    }

    // If columns_needed has entries that aren't standalone projection columns
    // (i.e., they came from column refs inside extracted aliases), a merge
    // into an inner projection will widen the schema with those extra columns,
    // requiring a recovery projection to restore the original schema.
    if columns_needed
        .iter()
        .any(|c| !standalone_columns.contains(c))
    {
        needs_recovery = true;
    }

    // ── Phase 2: Push down ──────────────────────────────────────────────
    let proj_input = Arc::clone(&proj.input);
    let pushed = push_extraction_pairs(
        &extraction_pairs,
        columns_needed,
        proj,
        &proj_input,
        alias_generator,
        proj_exprs_captured,
    )?;

    // ── Phase 3: Recovery ───────────────────────────────────────────────
    // Determine the base plan: either the pushed result or an in-place extraction.
    let base_plan = match pushed {
        Some(plan) => plan,
        None => {
            if !has_new_extractions {
                // Only pre-existing __datafusion_extracted aliases and columns, no new
                // extractions from routing_extract. The original projection is
                // already an extraction projection that couldn't be pushed
                // further. Return None.
                return Ok(None);
            }
            // Build extraction projection in-place (couldn't push down)
            let input_arc = Arc::clone(input);
            let extraction = build_extraction_projection_impl(
                &extraction_pairs,
                columns_needed,
                &input_arc,
                input_schema.as_ref(),
            )?;
            LogicalPlan::Projection(extraction)
        }
    };

    // Wrap with recovery projection if the output schema changed
    if needs_recovery {
        let recovery =
            LogicalPlan::Projection(Projection::try_new(recovery_exprs, Arc::new(base_plan))?);
        Ok(Some(recovery))
    } else {
        Ok(Some(base_plan))
    }
}

/// Returns true if the plan is a Projection where ALL expressions are either
/// `Alias(EXTRACTED_EXPR_PREFIX, ...)` or `Column`, with at least one extraction.
/// Such projections can safely be pushed further without re-extraction.
fn is_pure_extraction_projection(plan: &LogicalPlan) -> bool {
    let LogicalPlan::Projection(proj) = plan else {
        return false;
    };
    let mut has_extraction = false;
    for expr in &proj.expr {
        match expr {
            Expr::Alias(alias) if alias.name.starts_with(EXTRACTED_EXPR_PREFIX) => {
                has_extraction = true;
            }
            Expr::Column(_) => {}
            _ => return false,
        }
    }
    has_extraction
}

/// Pushes extraction pairs down through the projection's input node,
/// dispatching to the appropriate handler based on the input node type.
fn push_extraction_pairs(
    pairs: &[(Expr, String)],
    columns_needed: &IndexSet<Column>,
    proj: &Projection,
    proj_input: &Arc<LogicalPlan>,
    alias_generator: &Arc<AliasGenerator>,
    proj_exprs_captured: usize,
) -> Result<Option<LogicalPlan>> {
    match proj_input.as_ref() {
        // Merge into existing projection, then try to push the result further down.
        // Only merge when every expression in the outer projection is fully
        // captured as either an extraction pair (Case A: __datafusion_extracted
        // alias) or a plain column (Case B). Uncaptured expressions (e.g.
        // `col AS __common_expr_1` from CSE, or complex expressions with
        // extracted sub-parts) would be lost during the merge.
        LogicalPlan::Projection(_) if proj_exprs_captured == proj.expr.len() => {
            let target_schema = Arc::clone(proj_input.schema());
            let merged = build_extraction_projection_impl(
                pairs,
                columns_needed,
                proj_input,
                target_schema.as_ref(),
            )?;
            let merged_plan = LogicalPlan::Projection(merged);

            // After merging, try to push the result further down, but ONLY
            // if the merged result is still a pure extraction projection
            // (all __datafusion_extracted aliases + columns). If the merge inherited
            // bare MoveTowardsLeafNodes expressions from the inner projection,
            // pushing would re-extract them into new aliases and fail when
            // the (None, true) fallback can't find the original aliases.
            // This handles: Extraction → Recovery(cols) → Filter → ... → TableScan
            // by pushing through the recovery projection AND the filter in one pass.
            if is_pure_extraction_projection(&merged_plan) {
                if let Some(pushed) = try_push_input(&merged_plan, alias_generator)? {
                    return Ok(Some(pushed));
                }
            }
            Ok(Some(merged_plan))
        }
        // Hamelin fork: refuse to push past a Projection we couldn't merge
        // with. `try_push_into_inputs` doesn't translate the extraction pairs
        // through the inner projection's rename map, so a pair like
        // `get_field(x, "c")` would be routed into the projection's input
        // schema where `x` may not even have field `c` (e.g., the projection
        // redefined `x` to add new fields). The extraction projection stays
        // between this projection and its parent — correct, just less optimal.
        LogicalPlan::Projection(_) => Ok(None),
        // Generic: handles Filter/Sort/Limit (via recursion),
        // SubqueryAlias (with qualifier remap in try_push_into_inputs),
        // Join, and anything else.
        // Safely bails out for nodes that don't pass through extracted
        // columns (Aggregate, Window) via the output schema check.
        _ => try_push_into_inputs(pairs, columns_needed, proj_input.as_ref(), alias_generator),
    }
}

/// Routes extraction pairs and columns to the appropriate inputs.
///
/// - **Union**: broadcasts to every input via [`remap_pairs_and_columns`].
/// - **Other nodes**: routes each expression to the one input that owns
///   all of its column references (via [`find_owning_input`]).
///
/// Returns `None` if any expression can't be routed or no input has pairs.
fn route_to_inputs(
    pairs: &[(Expr, String)],
    columns: &IndexSet<Column>,
    node: &LogicalPlan,
    input_column_sets: &[std::collections::HashSet<Column>],
    input_schemas: &[Arc<DFSchema>],
) -> Result<Option<Vec<ExtractionTarget>>> {
    let num_inputs = input_schemas.len();
    let mut per_input: Vec<ExtractionTarget> = (0..num_inputs)
        .map(|_| ExtractionTarget {
            pairs: vec![],
            columns: IndexSet::new(),
        })
        .collect();

    if matches!(node, LogicalPlan::Union(_)) {
        // Union output schema and each input schema have the same fields by
        // index but may differ in qualifiers (e.g. output `s` vs input
        // `simple_struct.s`). Remap pairs/columns to each input's space.
        let union_schema = node.schema();
        for (idx, input_schema) in input_schemas.iter().enumerate() {
            per_input[idx] = remap_pairs_and_columns(pairs, columns, union_schema, input_schema)?;
        }
    } else {
        for (expr, alias) in pairs {
            match find_owning_input(expr, input_column_sets) {
                Some(idx) => per_input[idx].pairs.push((expr.clone(), alias.clone())),
                None => return Ok(None), // Cross-input expression — bail out
            }
        }
        for col in columns {
            let col_expr = Expr::Column(col.clone());
            match find_owning_input(&col_expr, input_column_sets) {
                Some(idx) => {
                    per_input[idx].columns.insert(col.clone());
                }
                None => return Ok(None), // Ambiguous column — bail out
            }
        }
    }

    // Check at least one input has extractions to push
    if per_input.iter().all(|t| t.pairs.is_empty()) {
        return Ok(None);
    }

    Ok(Some(per_input))
}

/// Pushes extraction expressions into a node's inputs by routing each
/// expression to the input that owns all of its column references.
///
/// Works for any number of inputs (1, 2, …N). For single-input nodes,
/// all expressions trivially route to that input. For multi-input nodes
/// (Join, etc.), each expression is routed to the side that owns its columns.
///
/// Returns `Some(new_node)` if all expressions could be routed AND the
/// rebuilt node's output schema contains all extracted aliases.
/// Returns `None` if any expression references columns from multiple inputs
/// or the node doesn't pass through the extracted columns.
///
/// # Example: Join with expressions from both sides
///
/// ```text
/// Extraction projection above a Join:
///   Projection: left.user['name'] AS __datafusion_extracted_1, right.order['total'] AS __datafusion_extracted_2, ...
///     Join: left.id = right.user_id
///       TableScan: left [id, user]
///       TableScan: right [user_id, order]
///
/// After routing each expression to its owning input:
///   Join: left.id = right.user_id
///     Projection: user['name'] AS __datafusion_extracted_1, id, user              <-- left-side extraction
///       TableScan: left [id, user]
///     Projection: order['total'] AS __datafusion_extracted_2, user_id, order      <-- right-side extraction
///       TableScan: right [user_id, order]
/// ```
fn try_push_into_inputs(
    pairs: &[(Expr, String)],
    columns_needed: &IndexSet<Column>,
    node: &LogicalPlan,
    alias_generator: &Arc<AliasGenerator>,
) -> Result<Option<LogicalPlan>> {
    let inputs = node.inputs();
    if inputs.is_empty() {
        return Ok(None);
    }

    // Hamelin fork: refuse to push extractions past Unnest. Unnest rewrites a
    // column's type from List<T> to T, so an extraction projection rebuilt
    // below the Unnest would type-check `get_field(col, …)` against an input
    // schema where `col` is still List<Struct>, which `get_field` rejects.
    if matches!(node, LogicalPlan::Unnest(_)) {
        return Ok(None);
    }

    // SubqueryAlias remaps qualifiers between input and output.
    // Rewrite pairs/columns from alias-space to input-space before routing.
    let remapped = if let LogicalPlan::SubqueryAlias(sa) = node {
        remap_pairs_and_columns(pairs, columns_needed, &sa.schema, sa.input.schema())?
    } else {
        ExtractionTarget {
            pairs: pairs.to_vec(),
            columns: columns_needed.clone(),
        }
    };
    let pairs = &remapped.pairs[..];
    let columns_needed = &remapped.columns;

    // Build per-input schemas and column sets for routing
    let input_schemas: Vec<Arc<DFSchema>> = inputs.iter().map(|i| Arc::clone(i.schema())).collect();
    let input_column_sets: Vec<std::collections::HashSet<Column>> =
        input_schemas.iter().map(|s| schema_columns(s)).collect();

    // Route pairs and columns to the appropriate inputs
    let per_input = match route_to_inputs(
        pairs,
        columns_needed,
        node,
        &input_column_sets,
        &input_schemas,
    )? {
        Some(routed) => routed,
        None => return Ok(None),
    };

    let num_inputs = inputs.len();

    // Build per-input extraction projections and push them as far as possible
    // immediately. This is critical because map_children preserves cached schemas,
    // so if the TopDown pass later pushes a child further (changing its output
    // schema), the parent node's schema becomes stale.
    let mut new_inputs: Vec<LogicalPlan> = Vec::with_capacity(num_inputs);
    for (idx, input) in inputs.into_iter().enumerate() {
        if per_input[idx].pairs.is_empty() {
            new_inputs.push(input.clone());
        } else {
            let input_arc = Arc::new(input.clone());
            let target_schema = Arc::clone(input.schema());
            let proj = build_extraction_projection_impl(
                &per_input[idx].pairs,
                &per_input[idx].columns,
                &input_arc,
                target_schema.as_ref(),
            )?;
            // Verify all requested aliases appear in the projection's output.
            // A merge may deduplicate if the same expression already exists
            // under a different alias, leaving the requested alias missing.
            let proj_schema = proj.schema.as_ref();
            for (_expr, alias) in &per_input[idx].pairs {
                if !proj_schema.fields().iter().any(|f| f.name() == alias) {
                    return Ok(None);
                }
            }
            let proj_plan = LogicalPlan::Projection(proj);
            // Try to push the extraction projection further down within
            // this input (e.g., through Filter → existing extraction projection).
            // This ensures the input's output schema is stable and won't change
            // when the TopDown pass later visits children.
            match try_push_input(&proj_plan, alias_generator)? {
                Some(pushed) => new_inputs.push(pushed),
                None => new_inputs.push(proj_plan),
            }
        }
    }

    // Rebuild the node with new inputs
    let new_node = node.with_new_exprs(node.expressions(), new_inputs)?;

    // Safety check: verify all extracted aliases appear in the rebuilt
    // node's output schema. Nodes like Aggregate define their own output
    // and won't pass through extracted columns — bail out for those.
    let output_schema = new_node.schema();
    for (_expr, alias) in pairs {
        if !output_schema.fields().iter().any(|f| f.name() == alias) {
            return Ok(None);
        }
    }

    Ok(Some(new_node))
}
