dust/src/abstract_tree/select.rs

use serde::{Deserialize, Serialize};

use crate::{AbstractTree, Expression, Identifier, Item, Result, Table, Value, VariableMap};

#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, PartialOrd, Ord)]
pub struct Select {
    identifiers: Vec<Identifier>,
    expression: Expression,
    item: Option<Item>,
}

impl AbstractTree for Select {
    fn from_syntax_node(source: &str, node: tree_sitter::Node) -> crate::Result<Self> {
        let child_count = node.child_count();
        let mut identifiers = Vec::new();

        for index in 2..child_count - 4 {
            let node = node.child(index).unwrap();

            if node.kind() == "identifier" {
                let identifier = Identifier::from_syntax_node(source, node)?;
                identifiers.push(identifier);
            }

            if node.kind() == ">" {
                break;
            }
        }

        let final_node = node.child(child_count - 1).unwrap();

        let item = if final_node.kind() == "}" {
            let item_node = node.child(child_count - 2).unwrap();

            Some(Item::from_syntax_node(source, item_node)?)
        } else {
            None
        };

        let expression_node = if item.is_some() {
            node.child(child_count - 4).unwrap()
        } else {
            node.child(child_count - 1).unwrap()
        };

        let expression = Expression::from_syntax_node(source, expression_node)?;

        Ok(Select {
            identifiers,
            expression,
            item,
        })
    }

    fn run(&self, source: &str, context: &mut VariableMap) -> Result<Value> {
        let value = self.expression.run(source, context)?;
        let old_table = value.as_table()?;
        let column_names = if self.identifiers.len() > 0 {
            self.identifiers
                .iter()
                .cloned()
                .map(|identifierier| identifierier.take_inner())
                .collect()
        } else {
            old_table.headers().clone()
        };

        let mut new_table = Table::new(column_names.to_vec());

        for row in old_table.rows() {
            let mut new_row = Vec::new();
            let mut row_context = VariableMap::new();

            for (i, value) in row.iter().enumerate() {
                let column_name = old_table.headers().get(i).unwrap();

                row_context.set_value(column_name.clone(), value.clone())?;

                let new_table_column_index =
                    new_table
                        .headers()
                        .iter()
                        .enumerate()
                        .find_map(|(index, new_column_name)| {
                            if new_column_name == column_name {
                                Some(index)
                            } else {
                                None
                            }
                        });

                if let Some(index) = new_table_column_index {
                    while new_row.len() < index + 1 {
                        new_row.push(Value::Empty);
                    }
                    new_row[index] = value.clone();
                }
            }

            if let Some(where_clause) = &self.item {
                let should_include = where_clause.run(source, &mut row_context)?.as_boolean()?;

                if should_include {
                    new_table.insert(new_row)?;
                }
            } else {
                new_table.insert(new_row)?;
            }
        }

        Ok(Value::Table(new_table))
    }
}