From 7a85cf6f860e30a769fcf26c7671a4875f8c43e8 Mon Sep 17 00:00:00 2001
From: Jeff <dev@jeffa.io>
Date: Tue, 17 Oct 2023 17:52:41 -0400
Subject: [PATCH] Implement filter loop

---
 examples/filter_loop.ds        |  8 ++++++
 src/abstract_tree/filter.rs    | 49 ++++++++++++++++++++++++++++++++++
 src/abstract_tree/mod.rs       |  5 ++--
 src/abstract_tree/statement.rs | 11 +++++---
 tree-sitter-dust               |  2 +-
 5 files changed, 69 insertions(+), 6 deletions(-)
 create mode 100644 examples/filter_loop.ds
 create mode 100644 src/abstract_tree/filter.rs

diff --git a/examples/filter_loop.ds b/examples/filter_loop.ds
new file mode 100644
index 0000000..c75d051
--- /dev/null
+++ b/examples/filter_loop.ds
@@ -0,0 +1,8 @@
+list = [1 2 1 2]
+
+new_list = filter i in list {
+	i == 2
+}
+
+(assert_equal [2 2] new_list)
+
diff --git a/src/abstract_tree/filter.rs b/src/abstract_tree/filter.rs
new file mode 100644
index 0000000..b0bbb41
--- /dev/null
+++ b/src/abstract_tree/filter.rs
@@ -0,0 +1,49 @@
+use serde::{Deserialize, Serialize};
+
+use crate::{AbstractTree, Expression, Identifier, Item, Value};
+
+#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, PartialOrd, Ord)]
+pub struct Filter {
+    identifier: Identifier,
+    expression: Expression,
+    item: Item,
+}
+
+impl AbstractTree for Filter {
+    fn from_syntax_node(source: &str, node: tree_sitter::Node) -> crate::Result<Self> {
+        let identifier_node = node.child(1).unwrap();
+        let identifier = Identifier::from_syntax_node(source, identifier_node)?;
+
+        let expression_node = node.child(3).unwrap();
+        let expression = Expression::from_syntax_node(source, expression_node)?;
+
+        let item_node = node.child(5).unwrap();
+        let item = Item::from_syntax_node(source, item_node)?;
+
+        Ok(Filter {
+            identifier,
+            expression,
+            item,
+        })
+    }
+
+    fn run(&self, source: &str, context: &mut crate::VariableMap) -> crate::Result<crate::Value> {
+        let value = self.expression.run(source, context)?;
+        let list = value.as_list()?;
+        let key = self.identifier.inner();
+        let mut context = context.clone();
+        let mut new_list = Vec::with_capacity(list.len());
+
+        for value in list {
+            context.set_value(key.clone(), value.clone())?;
+
+            let should_include = self.item.run(source, &mut context)?.as_boolean()?;
+
+            if should_include {
+                new_list.push(value.clone());
+            }
+        }
+
+        Ok(Value::List(new_list))
+    }
+}
diff --git a/src/abstract_tree/mod.rs b/src/abstract_tree/mod.rs
index 1b4e5f6..fab6300 100644
--- a/src/abstract_tree/mod.rs
+++ b/src/abstract_tree/mod.rs
@@ -9,6 +9,7 @@
 pub mod assignment;
 pub mod r#async;
 pub mod expression;
+pub mod filter;
 pub mod r#for;
 pub mod function_call;
 pub mod identifier;
@@ -24,8 +25,8 @@ pub mod value_node;
 pub mod r#while;
 
 pub use {
-    assignment::*, expression::*, function_call::*, identifier::*, if_else::*, item::*, logic::*,
-    math::*, r#async::*, r#for::*, r#match::*, r#while::*, statement::*, transform::*,
+    assignment::*, expression::*, filter::*, function_call::*, identifier::*, if_else::*, item::*,
+    logic::*, math::*, r#async::*, r#for::*, r#match::*, r#while::*, statement::*, transform::*,
 };
 
 use tree_sitter::Node;
diff --git a/src/abstract_tree/statement.rs b/src/abstract_tree/statement.rs
index 8546ab0..b5065cb 100644
--- a/src/abstract_tree/statement.rs
+++ b/src/abstract_tree/statement.rs
@@ -2,8 +2,8 @@ use serde::{Deserialize, Serialize};
 use tree_sitter::Node;
 
 use crate::{
-    AbstractTree, Assignment, Async, Error, Expression, For, IfElse, Match, Result, Transform,
-    Value, VariableMap, While,
+    AbstractTree, Assignment, Async, Error, Expression, Filter, For, IfElse, Match, Result,
+    Transform, Value, VariableMap, While,
 };
 
 /// Abstract representation of a statement.
@@ -20,6 +20,7 @@ pub enum Statement {
     Async(Box<Async>),
     For(Box<For>),
     Transform(Box<Transform>),
+    Filter(Box<Filter>),
 }
 
 impl AbstractTree for Statement {
@@ -53,8 +54,11 @@ impl AbstractTree for Statement {
             "transform" => Ok(Statement::Transform(Box::new(Transform::from_syntax_node(
                 source, child,
             )?))),
+            "filter" => Ok(Statement::Filter(Box::new(Filter::from_syntax_node(
+                source, child,
+            )?))),
             _ => Err(Error::UnexpectedSyntaxNode {
-                expected: "assignment, expression, if...else, while, for, transform, tool or async",
+                expected: "assignment, expression, if...else, while, for, transform, filter, tool or async",
                 actual: child.kind(),
                 location: child.start_position(),
                 relevant_source: source[child.byte_range()].to_string(),
@@ -72,6 +76,7 @@ impl AbstractTree for Statement {
             Statement::Async(run) => run.run(source, context),
             Statement::For(r#for) => r#for.run(source, context),
             Statement::Transform(transform) => transform.run(source, context),
+            Statement::Filter(filter) => filter.run(source, context),
         }
     }
 }
diff --git a/tree-sitter-dust b/tree-sitter-dust
index 6aab8e8..2f9590e 160000
--- a/tree-sitter-dust
+++ b/tree-sitter-dust
@@ -1 +1 @@
-Subproject commit 6aab8e8f65a29ada418b09c5cdae0583c37b7b23
+Subproject commit 2f9590e006243d45ba10b5a002ac139f7f0a221d