1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

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))
    }
}