5 files changed, 268 insertions, 165 deletions

diff --git a/flake.nix b/flake.nix
index 50f9efc..b27c441 100644
--- a/flake.nix
+++ b/flake.nix
@@ -66,7 +66,7 @@
           pkgs.cargo
 
           pkgs.rust-bin.nightly.latest.default
-          pkgs.rust-bin.nightly.latest.rust-analyzer
+          pkgs.rust-analyzer
 
           pkgs.mermaid-cli
         ];
diff --git a/readme.txt b/readme.txt
index 106fd0f..129792a 100644
--- a/readme.txt
+++ b/readme.txt
@@ -188,11 +188,11 @@ roadmap:
   - [ ] bytecode VM?
   - [ ] look into embedding high perf VMs, lua etc.
 - pattern matching
-  - [ ] allow matching on tree-sitter queries
+  - [x] allow matching on tree-sitter queries
   - [ ] support captures
 - language features
   - [ ] arrays and loops
-  - [ ] access node children
+  - [x] access node children
   - [x] access node fields
   - [ ] repr for ranges
   - [ ] comments
diff --git a/src/ast.rs b/src/ast.rs
index 6d7d326..7e83b3d 100644
--- a/src/ast.rs
+++ b/src/ast.rs
@@ -1,4 +1,4 @@
-#[derive(Debug)]
+#[derive(Debug, Default)]
 pub struct Program {
     pub stanzas: Vec<Stanza>,
 }
@@ -17,6 +17,30 @@ impl Program {
         self.stanzas = stanzas;
         Ok(self)
     }
+
+    pub fn begin(&self) -> Option<&Block> {
+        self.stanzas
+            .iter()
+            .find(|stanza| stanza.pattern == Pattern::Begin)
+            .map(|s| &s.statements)
+    }
+
+    pub fn end(&self) -> Option<&Block> {
+        self.stanzas
+            .iter()
+            .find(|stanza| stanza.pattern == Pattern::End)
+            .map(|s| &s.statements)
+    }
+
+    pub fn stanza_by_node(&self, node: tree_sitter::Node, state: Modifier) -> Option<&Block> {
+        self.stanzas
+            .iter()
+            .find(|stanza| {
+                stanza.pattern.matches(node)
+                    && matches!(stanza.pattern, Pattern::Tree { modifier, .. } if modifier == state)
+            })
+            .map(|s| &s.statements)
+    }
 }
 
 #[derive(Debug, PartialEq, Eq)]
@@ -29,13 +53,19 @@ pub struct Stanza {
 pub enum Pattern {
     Begin,
     End,
-    Node(NodePattern),
+    Tree {
+        modifier: Modifier,
+        matcher: TreePattern,
+    },
 }
 
-#[derive(Debug, Eq, PartialEq, Clone)]
-pub struct NodePattern {
-    pub modifier: Modifier,
-    pub kind: String,
+impl Pattern {
+    pub fn matches(&self, node: tree_sitter::Node) -> bool {
+        match self {
+            Self::Begin | Self::End => false,
+            Self::Tree { matcher, .. } => matcher.matches(node),
+        }
+    }
 }
 
 #[derive(Default, Debug, Eq, PartialEq, Clone, Copy)]
@@ -45,6 +75,32 @@ pub enum Modifier {
     Leave,
 }
 
+#[derive(Debug, Eq, PartialEq, Clone)]
+pub enum TreePattern {
+    Atom(String),
+    List(Vec<TreePattern>),
+}
+
+impl TreePattern {
+    pub fn matches(&self, node: tree_sitter::Node) -> bool {
+        match self {
+            Self::Atom(kind) => node.kind() == kind,
+            Self::List(l) => match l.as_slice() {
+                &[] => panic!(),
+                [kind] => kind.matches(node),
+                [root, rest @ ..] => {
+                    let root_match = root.matches(node);
+                    let child_match = rest
+                        .iter()
+                        .zip(node.named_children(&mut node.walk()))
+                        .all(|(pat, child)| pat.matches(child));
+                    root_match && child_match
+                }
+            },
+        }
+    }
+}
+
 #[derive(Debug, Default, Eq, PartialEq, Clone)]
 pub struct Block {
     pub body: Vec<Statement>,
@@ -110,7 +166,7 @@ impl Expr {
     #[cfg(test)]
     pub fn list<const N: usize>(items: [Expr; N]) -> Expr {
         Self::List(List {
-            items: items.to_vec()
+            items: items.to_vec(),
         })
     }
 }
diff --git a/src/eval.rs b/src/eval.rs
index 7d6c64e..e9fbbf2 100644
--- a/src/eval.rs
+++ b/src/eval.rs
@@ -346,65 +346,6 @@ impl From<Vec<Value>> for Value {
     }
 }
 
-type NodeKind = u16;
-
-#[derive(Debug, Default)]
-struct Visitor {
-    enter: ast::Block,
-    leave: ast::Block,
-}
-
-#[derive(Debug)]
-struct Visitors {
-    visitors: HashMap<NodeKind, Visitor>,
-    begin: ast::Block,
-    end: ast::Block,
-}
-
-impl Default for Visitors {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
-impl Visitors {
-    pub fn new() -> Self {
-        Self {
-            visitors: HashMap::new(),
-            begin: ast::Block { body: vec![] },
-            end: ast::Block { body: vec![] },
-        }
-    }
-
-    pub fn insert(
-        &mut self,
-        stanza: ast::Stanza,
-        language: &tree_sitter::Language,
-    ) -> std::result::Result<(), Error> {
-        match &stanza.pattern {
-            ast::Pattern::Begin => self.begin = stanza.statements,
-            ast::Pattern::End => self.end = stanza.statements,
-            ast::Pattern::Node(ast::NodePattern { modifier, kind }) => {
-                let id = language.id_for_node_kind(&kind, true);
-                if id == 0 {
-                    return Err(Error::InvalidNodeKind(kind.to_owned()));
-                }
-                let v = self.visitors.entry(id).or_default();
-                match modifier {
-                    ast::Modifier::Enter => v.enter = stanza.statements.clone(),
-                    ast::Modifier::Leave => v.leave = stanza.statements.clone(),
-                };
-            }
-        }
-        Ok(())
-    }
-
-    pub fn get_by_node(&self, node: tree_sitter::Node) -> Option<&Visitor> {
-        let node_id = node.kind_id();
-        self.visitors.get(&node_id)
-    }
-}
-
 #[derive(Debug, PartialEq, Eq)]
 pub enum Error {
     FailedLookup(ast::Identifier),
@@ -440,7 +381,7 @@ pub type Result = std::result::Result<Value, Error>;
 pub struct Context {
     variables: HashMap<ast::Identifier, Variable>,
     language: tree_sitter::Language,
-    visitors: Visitors,
+    program: ast::Program,
     pub(crate) input_src: Option<String>,
     cursor: Option<tree_sitter::TreeCursor<'static>>,
     tree: Option<&'static tree_sitter::Tree>,
@@ -452,7 +393,6 @@ impl fmt::Debug for Context {
         f.debug_struct("Context")
             .field("variables", &self.variables)
             .field("language", &self.language)
-            .field("visitors", &self.visitors)
             .field("input_src", &self.input_src)
             .field(
                 "cursor",
@@ -469,7 +409,7 @@ impl fmt::Debug for Context {
 impl Context {
     pub fn new(language: tree_sitter::Language) -> Self {
         Self {
-            visitors: Default::default(),
+            program: Default::default(),
             variables: Default::default(),
             language,
             input_src: None,
@@ -512,11 +452,9 @@ impl Context {
         None
     }
 
-    pub fn with_program(mut self, program: ast::Program) -> std::result::Result<Self, Error> {
-        for stanza in program.stanzas.into_iter() {
-            self.visitors.insert(stanza, &self.language)?;
-        }
-        Ok(self)
+    pub fn with_program(mut self, program: ast::Program) -> Self {
+        self.program = program;
+        self
     }
 
     pub fn with_input(mut self, src: String) -> Self {
@@ -566,13 +504,19 @@ impl Context {
             .wrap_ok()
     }
 
-    pub(crate) fn lookup(&mut self, ident: &ast::Identifier) -> std::result::Result<&Variable, Error> {
+    pub(crate) fn lookup(
+        &mut self,
+        ident: &ast::Identifier,
+    ) -> std::result::Result<&Variable, Error> {
         self.variables
             .get(ident)
             .ok_or_else(|| Error::FailedLookup(ident.to_owned()))
     }
 
-    pub(crate) fn lookup_mut(&mut self, ident: &ast::Identifier) -> std::result::Result<&mut Variable, Error> {
+    pub(crate) fn lookup_mut(
+        &mut self,
+        ident: &ast::Identifier,
+    ) -> std::result::Result<&mut Variable, Error> {
         self.variables
             .get_mut(ident)
             .ok_or_else(|| Error::FailedLookup(ident.to_owned()))
@@ -701,9 +645,11 @@ impl Context {
 
     fn eval_call(&mut self, call: &ast::Call) -> Result {
         ((&*crate::builtins::BUILTINS)
-         .get(call.function.as_str())
-         .ok_or_else(|| Error::FailedLookup(call.function.to_owned()))?)
-            (self, call.parameters.as_slice())
+            .get(call.function.as_str())
+            .ok_or_else(|| Error::FailedLookup(call.function.to_owned()))?)(
+            self,
+            call.parameters.as_slice(),
+        )
     }
 
     fn eval_list(&mut self, list: &ast::List) -> Result {
@@ -774,42 +720,50 @@ impl Context {
     }
 
     pub fn eval(&mut self) -> Result {
-        let visitors = std::mem::take(&mut self.visitors);
+        let program = std::mem::take(&mut self.program);
         let mut has_next = true;
         let mut postorder = Vec::new();
 
         // BEGIN block
-        self.eval_block(&visitors.begin)?;
+        if let Some(block) = program.begin() {
+            self.eval_block(block)?;
+        }
 
         while has_next {
             let current_node = self.cursor.as_ref().unwrap().node();
             postorder.push(current_node);
 
-            let visitor = visitors.get_by_node(current_node);
-
-            if let Some(v) = visitor {
-                self.eval_block(&v.enter)?;
+            if let Some(block) = program.stanza_by_node(current_node, ast::Modifier::Enter) {
+                self.eval_block(block)?;
             }
 
             has_next = self.goto_first_child();
 
             if !has_next {
                 has_next = self.cursor.as_mut().unwrap().goto_next_sibling();
-                if let Some(v) = postorder.pop().and_then(|n| visitors.get_by_node(n)) {
-                    self.eval_block(&v.leave)?;
-                }
+                if let Some(block) = postorder
+                    .pop()
+                    .and_then(|n| program.stanza_by_node(n, ast::Modifier::Leave))
+                {
+                    self.eval_block(block)?;
+                };
             }
 
             while !has_next && self.cursor.as_mut().unwrap().goto_parent() {
                 has_next = self.cursor.as_mut().unwrap().goto_next_sibling();
-                if let Some(v) = postorder.pop().and_then(|n| visitors.get_by_node(n)) {
-                    self.eval_block(&v.leave)?;
-                }
+                if let Some(block) = postorder
+                    .pop()
+                    .and_then(|n| program.stanza_by_node(n, ast::Modifier::Leave))
+                {
+                    self.eval_block(block)?;
+                };
             }
         }
 
         // END block
-        self.eval_block(&visitors.end)?;
+        if let Some(block) = program.end() {
+            self.eval_block(block)?;
+        }
 
         Ok(Value::Unit)
     }
@@ -825,7 +779,7 @@ pub fn evaluate(file: &str, program: &str, language: tree_sitter::Language) -> R
     let mut ctx = Context::new(language)
         .with_input(file.to_owned())
         .with_tree(tree)
-        .with_program(program)?;
+        .with_program(program);
 
     ctx.eval()
 }
@@ -838,7 +792,7 @@ mod test {
     #[test]
     fn bin() {
         let language = tree_sitter_python::language();
-        let mut ctx = Context::new(language).with_program(Program::new()).unwrap();
+        let mut ctx = Context::new(language).with_program(Program::new());
         assert_eq!(
             ctx.eval_expr(&Expr::bin(Expr::int(5), "+", Expr::int(10),)),
             Ok(Value::Integer(15))
@@ -864,7 +818,7 @@ mod test {
     #[test]
     fn test_evaluate_blocks() {
         let language = tree_sitter_python::language();
-        let mut ctx = Context::new(language).with_program(Program::new()).unwrap();
+        let mut ctx = Context::new(language).with_program(Program::new());
         assert_eq!(
             ctx.eval_block(&Block {
                 body: vec![
@@ -891,7 +845,7 @@ mod test {
     #[test]
     fn test_evaluate_if() {
         let language = tree_sitter_python::language();
-        let mut ctx = Context::new(language).with_program(Program::new()).unwrap();
+        let mut ctx = Context::new(language).with_program(Program::new());
         assert_eq!(
             ctx.eval_block(&Block {
                 body: vec![
@@ -934,7 +888,7 @@ mod test {
     #[test]
     fn test_substring() {
         let language = tree_sitter_python::language();
-        let mut ctx = Context::new(language).with_program(Program::new()).unwrap();
+        let mut ctx = Context::new(language).with_program(Program::new());
         assert_eq!(
             ctx.eval_block(&Block {
                 body: vec![
@@ -971,7 +925,7 @@ mod test {
     #[test]
     fn test_list() {
         let language = tree_sitter_python::language();
-        let mut ctx = Context::new(language).with_program(Program::new()).unwrap();
+        let mut ctx = Context::new(language).with_program(Program::new());
         assert_eq!(
             ctx.eval_block(&Block {
                 body: vec![Statement::Declaration(Declaration {
@@ -1000,14 +954,10 @@ mod test {
     #[test]
     fn test_ts_builtins() {
         let language = tree_sitter_python::language();
-        let mut ctx = Context::new(language).with_program(Program::new()).unwrap();
+        let mut ctx = Context::new(language).with_program(Program::new());
         assert_eq!(
             ctx.eval_block(&Block {
-                body: vec![Statement::decl(
-                    Type::List,
-                    "a",
-                    Expr::list([Expr::int(5)]),
-                )]
+                body: vec![Statement::decl(Type::List, "a", Expr::list([Expr::int(5)]),)]
             }),
             Ok(Value::Unit)
         );
diff --git a/src/parser.rs b/src/parser.rs
index 15d03fe..8b04307 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -4,7 +4,7 @@ use nom::{
     character::complete::{alpha1, alphanumeric1, char, multispace0, multispace1, one_of},
     combinator::{map, opt, recognize, value},
     error::ParseError,
-    multi::{fold_many0, many0, many0_count, many1, separated_list0},
+    multi::{fold_many0, many0, many0_count, many1, separated_list0, separated_list1},
     sequence::{delimited, pair, preceded, terminated, tuple},
     IResult, Parser,
 };
@@ -373,11 +373,35 @@ fn parse_modifier<'a>(i: &str) -> IResult<&str, Modifier> {
 fn parse_pattern<'a>(i: &str) -> IResult<&str, Pattern> {
     let begin = value(Pattern::Begin, ws(tag("BEGIN")));
     let end = value(Pattern::End, ws(tag("END")));
-    let node = map(
-        tuple((parse_modifier, multispace0, parse_ident)),
-        |(modifier, _, kind)| Pattern::Node(NodePattern { modifier, kind }),
-    );
-    ws(alt((begin, end, node))).parse(i)
+    ws(alt((begin, end, parse_tree_pattern))).parse(i)
+}
+
+// fn parse_node_pattern<'a>(i: &str) -> IResult<&str, Pattern> {
+//     map(
+//         tuple((parse_modifier, multispace0, parse_ident)),
+//         |(modifier, _, kind)| Pattern::Node(NodePattern { modifier, kind }),
+//     )
+//     .parse(i)
+// }
+
+fn parse_tree_pattern<'a>(i: &str) -> IResult<&str, Pattern> {
+    let parse_matcher = alt((parse_tree_atom, parse_tree_list));
+    tuple((parse_modifier, multispace0, parse_matcher))
+        .map(|(modifier, _, matcher)| Pattern::Tree { modifier, matcher })
+        .parse(i)
+}
+
+fn parse_tree_atom<'a>(i: &str) -> IResult<&str, TreePattern> {
+    parse_ident.map(TreePattern::Atom).parse(i)
+}
+
+fn parse_tree_list<'a>(i: &str) -> IResult<&str, TreePattern> {
+    let open = terminated(char('('), multispace0);
+    let close = preceded(multispace0, char(')'));
+    let list = separated_list1(multispace1, alt((parse_tree_atom, parse_tree_list)));
+    tuple((open, list, close))
+        .map(|(_, list, _)| TreePattern::List(list))
+        .parse(i)
 }
 
 pub fn parse_stanza<'a>(i: &str) -> IResult<&str, Stanza> {
@@ -745,78 +769,78 @@ mod test {
             parse_pattern("enter function_definition"),
             Ok((
                 "",
-                Pattern::Node(NodePattern {
+                Pattern::Tree {
                     modifier: Modifier::Enter,
-                    kind: "function_definition".to_owned()
-                })
+                    matcher: TreePattern::Atom("function_definition".to_owned()),
+                }
             ))
         );
         assert_eq!(
             parse_pattern("function_definition"),
             Ok((
                 "",
-                Pattern::Node(NodePattern {
+                Pattern::Tree {
                     modifier: Modifier::Enter,
-                    kind: "function_definition".to_owned()
-                })
+                    matcher: TreePattern::Atom("function_definition".to_owned()),
+                }
             ))
         );
         assert_eq!(
             parse_pattern("leave function_definition"),
             Ok((
                 "",
-                Pattern::Node(NodePattern {
+                Pattern::Tree {
                     modifier: Modifier::Leave,
-                    kind: "function_definition".to_owned()
-                })
-            ))
-        );
-    }
-
-    #[test]
-    fn test_parse_stanza() {
-        assert_eq!(
-            parse_stanza("enter function_definition { true; }"),
-            Ok((
-                "",
-                Stanza {
-                    pattern: Pattern::Node(NodePattern {
-                        modifier: Modifier::Enter,
-                        kind: "function_definition".to_owned()
-                    }),
-                    statements: Block {
-                        body: vec![Statement::Bare(Expr::true_())]
-                    }
-                }
-            ))
-        );
-        assert_eq!(
-            parse_stanza("BEGIN { true; }"),
-            Ok((
-                "",
-                Stanza {
-                    pattern: Pattern::Begin,
-                    statements: Block {
-                        body: vec![Statement::Bare(Expr::true_())]
-                    }
-                }
-            ))
-        );
-        assert_eq!(
-            parse_block(
-                " { 
-                    true;
-                }"
-            ),
-            Ok((
-                "",
-                Block {
-                    body: vec![Statement::Bare(Expr::true_())]
+                    matcher: TreePattern::Atom("function_definition".to_owned()),
                 }
             ))
         );
     }
 
+    // #[test]
+    // fn test_parse_stanza() {
+    //     assert_eq!(
+    //         parse_stanza("enter function_definition { true; }"),
+    //         Ok((
+    //             "",
+    //             Stanza {
+    //                 pattern: Pattern::Node(NodePattern {
+    //                     modifier: Modifier::Enter,
+    //                     kind: "function_definition".to_owned()
+    //                 }),
+    //                 statements: Block {
+    //                     body: vec![Statement::Bare(Expr::true_())]
+    //                 }
+    //             }
+    //         ))
+    //     );
+    //     assert_eq!(
+    //         parse_stanza("BEGIN { true; }"),
+    //         Ok((
+    //             "",
+    //             Stanza {
+    //                 pattern: Pattern::Begin,
+    //                 statements: Block {
+    //                     body: vec![Statement::Bare(Expr::true_())]
+    //                 }
+    //             }
+    //         ))
+    //     );
+    //     assert_eq!(
+    //         parse_block(
+    //             " {
+    //                 true;
+    //             }"
+    //         ),
+    //         Ok((
+    //             "",
+    //             Block {
+    //                 body: vec![Statement::Bare(Expr::true_())]
+    //             }
+    //         ))
+    //     );
+    // }
+
     #[test]
     fn test_parse_if_statement_regression() {
         assert_eq!(
@@ -848,4 +872,77 @@ mod test {
             ))
         );
     }
+    #[test]
+    fn test_parse_tree_pattern() {
+        assert_eq!(
+            parse_tree_pattern("enter foo"),
+            Ok((
+                "",
+                Pattern::Tree {
+                    modifier: Modifier::Enter,
+                    matcher: TreePattern::Atom("foo".to_owned())
+                }
+            ))
+        );
+        assert_eq!(
+            parse_tree_pattern("enter (foo)"),
+            Ok((
+                "",
+                Pattern::Tree {
+                    modifier: Modifier::Enter,
+                    matcher: TreePattern::List(vec![TreePattern::Atom("foo".to_owned())])
+                }
+            ))
+        );
+        assert_eq!(
+            parse_tree_pattern("leave (foo bar baz)"),
+            Ok((
+                "",
+                Pattern::Tree {
+                    modifier: Modifier::Leave,
+                    matcher: TreePattern::List(vec![
+                        TreePattern::Atom("foo".to_owned()),
+                        TreePattern::Atom("bar".to_owned()),
+                        TreePattern::Atom("baz".to_owned()),
+                    ])
+                }
+            ))
+        );
+        assert_eq!(
+            parse_tree_pattern("leave (foo (bar quux) baz)"),
+            Ok((
+                "",
+                Pattern::Tree {
+                    modifier: Modifier::Leave,
+                    matcher: TreePattern::List(vec![
+                        TreePattern::Atom("foo".to_owned()),
+                        TreePattern::List(vec![
+                            TreePattern::Atom("bar".to_owned()),
+                            TreePattern::Atom("quux".to_owned())
+                        ]),
+                        TreePattern::Atom("baz".to_owned()),
+                    ])
+                }
+            ))
+        );
+        assert_eq!(
+            parse_tree_pattern("enter ( foo (bar quux ) baz)"),
+            Ok((
+                "",
+                Pattern::Tree {
+                    modifier: Modifier::Enter,
+                    matcher: TreePattern::List(vec![
+                        TreePattern::Atom("foo".to_owned()),
+                        TreePattern::List(vec![
+                            TreePattern::Atom("bar".to_owned()),
+                            TreePattern::Atom("quux".to_owned())
+                        ]),
+                        TreePattern::Atom("baz".to_owned()),
+                    ])
+                }
+            ))
+        );
+        assert!(parse_tree_pattern("(  )").is_err());
+        assert!(parse_tree_pattern("()").is_err());
+    }
 }