7 files changed, 211 insertions, 58 deletions

diff --git a/.gitignore b/.gitignore
index 4075fb2..3906c5c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,2 +1,3 @@
 target
 result
+.direnv
diff --git a/Cargo.lock b/Cargo.lock
index afa07fe..164709d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -49,6 +49,22 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "066fce287b1d4eafef758e89e09d724a24808a9196fe9756b8ca90e86d0719a2"
 
 [[package]]
+name = "dissimilar"
+version = "1.0.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "59f8e79d1fbf76bdfbde321e902714bf6c49df88a7dda6fc682fc2979226962d"
+
+[[package]]
+name = "expect-test"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9e0be0a561335815e06dab7c62e50353134c796e7a6155402a64bcff66b6a5e0"
+dependencies = [
+ "dissimilar",
+ "once_cell",
+]
+
+[[package]]
 name = "la-arena"
 version = "0.3.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -77,6 +93,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "once_cell"
+version = "1.20.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1261fe7e33c73b354eab43b1273a57c8f967d0391e80353e51f764ac02cf6775"
+
+[[package]]
 name = "proc-macro2"
 version = "1.0.86"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -159,6 +181,7 @@ name = "tbsp"
 version = "0.1.0"
 dependencies = [
  "argh",
+ "expect-test",
  "la-arena",
  "nom",
  "regex",
diff --git a/Cargo.toml b/Cargo.toml
index 5d43ae0..894ea00 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -29,3 +29,4 @@ python = ["tree-sitter-python"]
 rust = ["tree-sitter-rust"]
 
 [dev-dependencies]
+expect-test = "1.5.0"
diff --git a/flake.nix b/flake.nix
index b27c441..531d73f 100644
--- a/flake.nix
+++ b/flake.nix
@@ -68,7 +68,6 @@
           pkgs.rust-bin.nightly.latest.default
           pkgs.rust-analyzer
 
-          pkgs.mermaid-cli
         ];
         RUST_LOG = "info";
         RUST_BACKTRACE = 1;
diff --git a/src/builtins.rs b/src/eval/builtins.rs
index 0c9dbf3..7820a8e 100644
--- a/src/builtins.rs
+++ b/src/eval/builtins.rs
@@ -24,6 +24,7 @@ macro_rules! builtins {
 
 builtins! {
     print,
+    println,
 
     // string
     isupper,
@@ -41,16 +42,27 @@ builtins! {
     member,
     push,
     pop,
+    isempty,
 }
 
 fn print(ctx: &mut Context, args: &[ast::Expr]) -> Result {
     for arg in args {
         let val = ctx.eval_expr(arg)?;
-        print!("{val}");
+        let mut default_stream =Box::new(std::io::stdout()) as Box<dyn std::io::Write> ;
+        let stream = ctx
+            .output_stream
+            .as_mut()
+            .unwrap_or(&mut default_stream);
+        write!(stream, "{val}").unwrap();
     }
     Ok(Value::Unit)
 }
 
+fn println(ctx: &mut Context, args: &[ast::Expr]) -> Result {
+    print(ctx, args)?;
+    print(ctx, &[ast::Expr::Lit(ast::Literal::Str("\n".to_owned()))])
+}
+
 fn isupper(ctx: &mut Context, args: &[ast::Expr]) -> Result {
     Ok(ctx
         .eval_expr(&get_args::<1>(args)?[0])?
@@ -169,10 +181,7 @@ fn push(ctx: &mut Context, args: &[ast::Expr]) -> Result {
             l.push(element);
             Ok(Value::Unit)
         }
-        _ => Err(Error::TypeMismatch {
-            expected: ast::Type::List,
-            got: v.ty().clone(),
-        }),
+        _ => Err(v.ty().expected([ast::Type::List])),
     })
 }
 
@@ -189,16 +198,59 @@ fn pop(ctx: &mut Context, args: &[ast::Expr]) -> Result {
         Value::List(l) => l
             .pop()
             .ok_or_else(|| Error::ArrayOutOfBounds { idx: 0, len: 0 }),
-        _ => Err(Error::TypeMismatch {
-            expected: ast::Type::List,
-            got: v.ty().clone(),
-        }),
+        _ => Err(v.ty().expected([ast::Type::List])),
     })
 }
 
+fn isempty(ctx: &mut Context, args: &[ast::Expr]) -> Result {
+    let v = ctx.eval_expr(&get_args::<1>(args)?[0])?;
+    match v.ty() {
+        ast::Type::List => v
+            .as_list()
+            .unwrap()
+            .is_empty()
+            .wrap(Value::Boolean)
+            .wrap(Ok),
+        ast::Type::String => v
+            .as_str()
+            .unwrap()
+            .is_empty()
+            .wrap(Value::Boolean)
+            .wrap(Ok),
+        _ => Err(v.ty().expected([ast::Type::List])),
+    }
+}
+
 fn get_args<const N: usize>(args: &[ast::Expr]) -> std::result::Result<&[ast::Expr; N], Error> {
     args.try_into().map_err(|_| Error::IncorrectArgFormat {
         wanted: N,
         got: args.len(),
     })
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::{ast::*, eval::*};
+
+    #[test]
+    fn test_ts_builtins() {
+        let language = tree_sitter_python::language();
+        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)]),)]
+            }),
+            Ok(Value::Unit)
+        );
+        assert_eq!(
+            ctx.lookup(&String::from("a")).unwrap().clone(),
+            Variable {
+                ty: Type::List,
+                name: "a".to_owned(),
+                value: vec![5usize.into()].into(),
+            }
+        );
+    }
+}
diff --git a/src/eval.rs b/src/eval/mod.rs
index e9fbbf2..c4460c0 100644
--- a/src/eval.rs
+++ b/src/eval/mod.rs
@@ -1,7 +1,9 @@
 //! tree walking interpreter for tbsp
 
 use crate::{ast, Wrap};
-use std::{collections::HashMap, fmt};
+use std::{collections::HashMap, fmt, io};
+
+mod builtins;
 
 #[derive(Debug, PartialEq, Eq, Clone)]
 pub struct Variable {
@@ -15,19 +17,16 @@ impl Variable {
         &self.value
     }
 
-    pub(crate) fn ty(&self) -> &ast::Type {
-        &self.ty
+    pub(crate) fn ty(&self) -> ast::Type {
+        self.ty
     }
 
     fn assign(&mut self, value: Value) -> Result {
-        if self.ty() == &value.ty() {
+        if self.ty() == value.ty() {
             self.value = value;
             Ok(self.value.clone())
         } else {
-            Err(Error::TypeMismatch {
-                expected: self.ty().clone(),
-                got: value.ty(),
-            })
+            Err(value.ty().expected([self.ty()]))
         }
     }
 
@@ -90,50 +89,35 @@ impl Value {
     fn as_boolean(&self) -> std::result::Result<bool, Error> {
         match self {
             Self::Boolean(b) => Ok(*b),
-            v => Err(Error::TypeMismatch {
-                expected: ast::Type::Boolean,
-                got: v.ty(),
-            }),
+            v => Err(v.ty().expected([ast::Type::Boolean])),
         }
     }
 
     pub(crate) fn as_str(&self) -> std::result::Result<&str, Error> {
         match self {
             Self::String(s) => Ok(s.as_str()),
-            v => Err(Error::TypeMismatch {
-                expected: ast::Type::String,
-                got: v.ty(),
-            }),
+            v => Err(v.ty().expected([ast::Type::String])),
         }
     }
 
     pub(crate) fn as_int(&self) -> std::result::Result<i128, Error> {
         match self {
             Self::Integer(i) => Ok(*i),
-            v => Err(Error::TypeMismatch {
-                expected: ast::Type::Integer,
-                got: v.ty(),
-            }),
+            v => Err(v.ty().expected([ast::Type::Integer])),
         }
     }
 
     pub(crate) fn as_node(&self) -> std::result::Result<NodeId, Error> {
         match self {
             Self::Node(id) => Ok(*id),
-            v => Err(Error::TypeMismatch {
-                expected: ast::Type::Node,
-                got: v.ty(),
-            }),
+            v => Err(v.ty().expected([ast::Type::Node])),
         }
     }
 
     pub(crate) fn as_list(&self) -> std::result::Result<Vec<Value>, Error> {
         match self {
             Self::List(values) => Ok(values.clone()),
-            v => Err(Error::TypeMismatch {
-                expected: ast::Type::List,
-                got: v.ty(),
-            }),
+            v => Err(v.ty().expected([ast::Type::List])),
         }
     }
 
@@ -350,7 +334,7 @@ impl From<Vec<Value>> for Value {
 pub enum Error {
     FailedLookup(ast::Identifier),
     TypeMismatch {
-        expected: ast::Type,
+        expected: Vec<ast::Type>,
         got: ast::Type,
     },
     UndefinedBinOp(ast::BinOp, ast::Type, ast::Type),
@@ -376,9 +360,18 @@ pub enum Error {
     CurrentNodeNotPresent,
 }
 
+impl ast::Type {
+    pub fn expected<const N: usize>(self, expected: [Self; N]) -> Error {
+        Error::TypeMismatch {
+            expected: expected.to_vec(),
+            got: self,
+        }
+    }
+}
+
 pub type Result = std::result::Result<Value, Error>;
 
-pub struct Context {
+pub struct Context<'s> {
     variables: HashMap<ast::Identifier, Variable>,
     language: tree_sitter::Language,
     program: ast::Program,
@@ -386,9 +379,10 @@ pub struct Context {
     cursor: Option<tree_sitter::TreeCursor<'static>>,
     tree: Option<&'static tree_sitter::Tree>,
     cache: HashMap<NodeId, tree_sitter::Node<'static>>,
+    output_stream: Option<Box<dyn io::Write + 's>>,
 }
 
-impl fmt::Debug for Context {
+impl<'s> fmt::Debug for Context<'s> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("Context")
             .field("variables", &self.variables)
@@ -406,7 +400,7 @@ impl fmt::Debug for Context {
     }
 }
 
-impl Context {
+impl<'s> Context<'s> {
     pub fn new(language: tree_sitter::Language) -> Self {
         Self {
             program: Default::default(),
@@ -416,6 +410,7 @@ impl Context {
             cursor: None,
             tree: None,
             cache: HashMap::default(),
+            output_stream: Some(Box::new(io::stdout()) as Box<dyn io::Write + 's>),
         }
     }
 
@@ -469,6 +464,11 @@ impl Context {
         self
     }
 
+    pub fn with_output_stream(mut self, stream: Box<dyn io::Write + 's>) -> Self {
+        self.output_stream = Some(stream);
+        self
+    }
+
     pub(crate) fn eval_expr(&mut self, expr: &ast::Expr) -> Result {
         match expr {
             ast::Expr::Unit => Ok(Value::Unit),
@@ -644,7 +644,7 @@ impl Context {
     }
 
     fn eval_call(&mut self, call: &ast::Call) -> Result {
-        ((&*crate::builtins::BUILTINS)
+        ((&*builtins::BUILTINS)
             .get(call.function.as_str())
             .ok_or_else(|| Error::FailedLookup(call.function.to_owned()))?)(
             self,
@@ -779,6 +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_output_stream(Box::new(io::stdout()))
         .with_program(program);
 
     ctx.eval()
@@ -788,6 +789,28 @@ pub fn evaluate(file: &str, program: &str, language: tree_sitter::Language) -> R
 mod test {
     use super::*;
     use crate::ast::*;
+    use expect_test::{expect, Expect};
+
+    fn gen_test(file: &str, program: &str, expected: Expect) {
+        let language = tree_sitter_python::language();
+        let mut parser = tree_sitter::Parser::new();
+        let _ = parser.set_language(&language);
+        let tree = parser.parse(file, None).unwrap();
+        let program = ast::Program::new().from_str(program).unwrap();
+
+        let mut output = Vec::new();
+        let mut ctx = Context::new(language)
+            .with_input(file.to_owned())
+            .with_tree(tree)
+            .with_program(program)
+            .with_output_stream(Box::new(&mut output) as Box<dyn io::Write>);
+
+        ctx.eval().unwrap();
+
+        drop(ctx);
+
+        expected.assert_eq(&String::from_utf8(output).unwrap())
+    }
 
     #[test]
     fn bin() {
@@ -952,22 +975,77 @@ mod test {
     }
 
     #[test]
-    fn test_ts_builtins() {
-        let language = tree_sitter_python::language();
-        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)]),)]
-            }),
-            Ok(Value::Unit)
+    fn list_builtins() {
+        gen_test(
+            "",
+            "BEGIN {
+                list a = [5];
+                print(a);
+            }
+            ",
+            expect!["[5]"]
         );
-        assert_eq!(
-            ctx.lookup(&String::from("a")).unwrap().clone(),
-            Variable {
-                ty: Type::List,
-                name: "a".to_owned(),
-                value: vec![5usize.into()].into(),
+
+        gen_test(
+            "",
+            "BEGIN {
+                list a = [5, 4, 3];
+                print(length(a));
+            }
+            ",
+            expect!["3"]
+        );
+
+        gen_test(
+            "",
+            r#"BEGIN {
+                list a = [5, 4, 3];
+                print(member(a, 3));
+                print(", ");
+                print(member(a, 6));
             }
+            "#,
+            expect!["true, false"]
         );
+
+        gen_test(
+            "",
+            r#"BEGIN {
+                list a = [5];
+                println(a);
+                push(a, 4);
+                println(a);
+                push(a, 3);
+                println(a);
+                pop(a);
+                println(a);
+                pop(a);
+                println(a);
+            }
+            "#,
+            expect![[r#"
+                [5]
+                [5, 4]
+                [5, 4, 3]
+                [5, 4]
+                [5]
+            "#]]
+        );
+
+        gen_test(
+            "",
+            r#"BEGIN {
+                list a = [5];
+                println(isempty(a));
+                pop(a);
+                println(isempty(a));
+            }
+            "#,
+            expect![[r#"
+                false
+                true
+            "#]]
+        );
+
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
index afce26c..ec4e2d9 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,5 +1,4 @@
 mod ast;
-mod builtins;
 mod eval;
 mod parser;
 mod string;