3 files changed, 346 insertions, 3 deletions

diff --git a/compiler/src/ast/models.rs b/compiler/src/ast/models.rs
index 64fec19..1cb705d 100644
--- a/compiler/src/ast/models.rs
+++ b/compiler/src/ast/models.rs
@@ -31,6 +31,11 @@ pub enum Expr {
         vars: Vec<Spanned<LetVar>>,
         body: Vec<Spanned<Expr>>,
     },
+    If {
+        condition: Spanned<Box<Expr>>,
+        then_expr: Spanned<Box<Expr>>,
+        else_expr: Option<Spanned<Box<Expr>>>,
+    },
     For {
         loop_var: Spanned<Rc<str>>,
         from: Spanned<Box<Expr>>,
diff --git a/compiler/src/ast/parser.rs b/compiler/src/ast/parser.rs
index 12a6f0d..88ba8c7 100644
--- a/compiler/src/ast/parser.rs
+++ b/compiler/src/ast/parser.rs
@@ -244,14 +244,40 @@ where
         Ok(Spanned::new(let_vars, span))
     }
 
+    fn parse_if(&mut self, open_span: Span) -> Result<Spanned<Expr>, Spanned<Error>> {
+        self.consume()?;
+
+        let condition = self.parse_expr()?.map(Into::into);
+        let then_expr = self.parse_expr()?.map(Into::into);
+        let else_expr = match self.peek_token()?.inner {
+            Token::RightPar => {
+                self.consume()?;
+                None
+            }
+            _ => {
+                let expr = self.parse_expr()?.map(Into::into);
+                self.require_right_par()?;
+                Some(expr)
+            }
+        };
+
+        let if_stmt = Expr::If {
+            condition,
+            then_expr,
+            else_expr,
+        };
+        let span = Span::new(open_span.start, self.last_token_span.end);
+        Ok(Spanned::new(if_stmt, span))
+    }
+
     fn parse_for(&mut self, open_span: Span) -> Result<Spanned<Expr>, Spanned<Error>> {
         self.consume()?;
 
         let loop_var = self.parse_sym()?.map(Into::into);
         self.require_sym("from")?;
-        let from = self.parse_expr()?.map(Box::new);
+        let from = self.parse_expr()?.map(Into::into);
         self.require_sym("to")?;
-        let to = self.parse_expr()?.map(Box::new);
+        let to = self.parse_expr()?.map(Into::into);
         let body = self.parse_body(open_span, true)?;
 
         let for_loop = Expr::For {
@@ -268,7 +294,7 @@ where
         self.consume()?;
 
         let target_var = self.parse_sym()?.map(Into::into);
-        let expr = self.parse_expr()?.map(Box::new);
+        let expr = self.parse_expr()?.map(Into::into);
         self.require_right_par()?;
 
         let set = Expr::Set { target_var, expr };
@@ -312,6 +338,7 @@ where
             "fn" => self.parse_fn(open_span),
             "const" => self.parse_const(open_span),
             "let" => self.parse_let(open_span),
+            "if" => self.parse_if(open_span),
             "for" => self.parse_for(open_span),
             "set" => self.parse_set(open_span),
             "do" => self.parse_do(open_span),
diff --git a/compiler/src/ast/tests.rs b/compiler/src/ast/tests.rs
index 5b2399f..d2a27a1 100644
--- a/compiler/src/ast/tests.rs
+++ b/compiler/src/ast/tests.rs
@@ -16,6 +16,7 @@ enum E {
     Fn(&'static str, Vec<&'static str>, Vec<E>),
     Const(Vec<(&'static str, E)>),
     Let(Vec<(&'static str, E)>, Vec<E>),
+    If(Box<E>, Box<E>, Option<Box<E>>),
     For(&'static str, Box<E>, Box<E>, Vec<E>),
     Set(&'static str, Box<E>),
     Do(Vec<E>),
@@ -50,6 +51,15 @@ impl From<Expr> for E {
                 vars.into_iter().map(|v| convert_var(v.inner)).collect(),
                 body.into_iter().map(|e| e.inner.into()).collect(),
             ),
+            Expr::If {
+                condition,
+                then_expr,
+                else_expr,
+            } => If(
+                Box::new(E::from(*condition.inner)),
+                Box::new(E::from(*then_expr.inner)),
+                else_expr.map(|e| Box::new(E::from(*e.inner))),
+            ),
             Expr::For {
                 loop_var,
                 from,
@@ -388,6 +398,141 @@ fn test_let() {
 }
 
 #[test]
+fn test_if() {
+    let cases = vec![
+        // (if true 1 2)
+        (
+            vec![
+                LeftPar,
+                Symbol("if"),
+                Symbol("true"),
+                Number("1"),
+                Number("2"),
+                RightPar,
+            ],
+            vec![If(
+                Box::new(Bool(true)),
+                Box::new(Int(1)),
+                Some(Box::new(Int(2))),
+            )],
+        ),
+        // (if true 1)
+        (
+            vec![LeftPar, Symbol("if"), Symbol("true"), Number("1"), RightPar],
+            vec![If(Box::new(Bool(true)), Box::new(Int(1)), None)],
+        ),
+        // (if (= x 0) a b)
+        (
+            vec![
+                LeftPar,
+                Symbol("if"),
+                LeftPar,
+                Symbol("="),
+                Symbol("x"),
+                Number("0"),
+                RightPar,
+                Symbol("a"),
+                Symbol("b"),
+                RightPar,
+            ],
+            vec![If(
+                Box::new(Call("=", vec![Sym("x"), Int(0)])),
+                Box::new(Sym("a")),
+                Some(Box::new(Sym("b"))),
+            )],
+        ),
+        // (if true (do 1 2) (do 3 4))
+        (
+            vec![
+                LeftPar,
+                Symbol("if"),
+                Symbol("true"),
+                LeftPar,
+                Symbol("do"),
+                Number("1"),
+                Number("2"),
+                RightPar,
+                LeftPar,
+                Symbol("do"),
+                Number("3"),
+                Number("4"),
+                RightPar,
+                RightPar,
+            ],
+            vec![If(
+                Box::new(Bool(true)),
+                Box::new(Do(vec![Int(1), Int(2)])),
+                Some(Box::new(Do(vec![Int(3), Int(4)]))),
+            )],
+        ),
+        // (if nil nil)
+        (
+            vec![
+                LeftPar,
+                Symbol("if"),
+                Symbol("nil"),
+                Symbol("nil"),
+                RightPar,
+            ],
+            vec![If(Box::new(Nil), Box::new(Nil), None)],
+        ),
+        // (if true (if false 1 2) 3)
+        (
+            vec![
+                LeftPar,
+                Symbol("if"),
+                Symbol("true"),
+                LeftPar,
+                Symbol("if"),
+                Symbol("false"),
+                Number("1"),
+                Number("2"),
+                RightPar,
+                Number("3"),
+                RightPar,
+            ],
+            vec![If(
+                Box::new(Bool(true)),
+                Box::new(If(
+                    Box::new(Bool(false)),
+                    Box::new(Int(1)),
+                    Some(Box::new(Int(2))),
+                )),
+                Some(Box::new(Int(3))),
+            )],
+        ),
+        // (if true 1 (if false 2 3))
+        (
+            vec![
+                LeftPar,
+                Symbol("if"),
+                Symbol("true"),
+                Number("1"),
+                LeftPar,
+                Symbol("if"),
+                Symbol("false"),
+                Number("2"),
+                Number("3"),
+                RightPar,
+                RightPar,
+            ],
+            vec![If(
+                Box::new(Bool(true)),
+                Box::new(Int(1)),
+                Some(Box::new(If(
+                    Box::new(Bool(false)),
+                    Box::new(Int(2)),
+                    Some(Box::new(Int(3))),
+                ))),
+            )],
+        ),
+    ];
+    for (tokens, ast) in cases {
+        assert_eq!(parse(tokens.clone()), ast, "input: {tokens:?}");
+    }
+}
+
+#[test]
 fn test_for() {
     let cases = vec![
         // (for i from 0 to 10 i)
@@ -947,6 +1092,9 @@ fn test_keyword_prefixed_symbols() {
         (vec![Symbol("set?")], vec![Sym("set?")]),
         (vec![Symbol("for-each")], vec![Sym("for-each")]),
         (vec![Symbol("const-x")], vec![Sym("const-x")]),
+        (vec![Symbol("if-then")], vec![Sym("if-then")]),
+        (vec![Symbol("iffy")], vec![Sym("iffy")]),
+        (vec![Symbol("if?")], vec![Sym("if?")]),
         // (fn-foo 1)
         (
             vec![LeftPar, Symbol("fn-foo"), Number("1"), RightPar],
@@ -1150,6 +1298,16 @@ fn test_special_form_errors() {
             vec![LeftPar, Symbol("const"), LeftPar, RightPar, RightPar],
             Error::UnexpectedToken,
         ),
+        // (if)
+        (
+            vec![LeftPar, Symbol("if"), RightPar],
+            Error::UnexpectedClosePar,
+        ),
+        // (if true)
+        (
+            vec![LeftPar, Symbol("if"), Symbol("true"), RightPar],
+            Error::UnexpectedClosePar,
+        ),
     ];
     for (tokens, expected) in cases {
         assert_eq!(parse_err(tokens.clone()), expected, "input: {tokens:?}");
@@ -1301,6 +1459,79 @@ fn test_nested_special_forms() {
                 vec![Sym("x"), Call("+", vec![Sym("x"), Int(1)])],
             )],
         ),
+        // (if (if a b c) (if d e) (if f g h))
+        (
+            vec![
+                LeftPar,
+                Symbol("if"),
+                LeftPar,
+                Symbol("if"),
+                Symbol("a"),
+                Symbol("b"),
+                Symbol("c"),
+                RightPar,
+                LeftPar,
+                Symbol("if"),
+                Symbol("d"),
+                Symbol("e"),
+                RightPar,
+                LeftPar,
+                Symbol("if"),
+                Symbol("f"),
+                Symbol("g"),
+                Symbol("h"),
+                RightPar,
+                RightPar,
+            ],
+            vec![If(
+                Box::new(If(
+                    Box::new(Sym("a")),
+                    Box::new(Sym("b")),
+                    Some(Box::new(Sym("c"))),
+                )),
+                Box::new(If(Box::new(Sym("d")), Box::new(Sym("e")), None)),
+                Some(Box::new(If(
+                    Box::new(Sym("f")),
+                    Box::new(Sym("g")),
+                    Some(Box::new(Sym("h"))),
+                ))),
+            )],
+        ),
+        // (fn abs (x) (if (< x 0) (- 0 x) x))
+        (
+            vec![
+                LeftPar,
+                Symbol("fn"),
+                Symbol("abs"),
+                LeftPar,
+                Symbol("x"),
+                RightPar,
+                LeftPar,
+                Symbol("if"),
+                LeftPar,
+                Symbol("<"),
+                Symbol("x"),
+                Number("0"),
+                RightPar,
+                LeftPar,
+                Symbol("-"),
+                Number("0"),
+                Symbol("x"),
+                RightPar,
+                Symbol("x"),
+                RightPar,
+                RightPar,
+            ],
+            vec![Fn(
+                "abs",
+                vec!["x"],
+                vec![If(
+                    Box::new(Call("<", vec![Sym("x"), Int(0)])),
+                    Box::new(Call("-", vec![Int(0), Sym("x")])),
+                    Some(Box::new(Sym("x"))),
+                )],
+            )],
+        ),
     ];
     for (tokens, ast) in cases {
         assert_eq!(parse(tokens.clone()), ast, "input: {tokens:?}");
@@ -1471,6 +1702,86 @@ fn test_span_let() {
 }
 
 #[test]
+fn test_span_if() {
+    // (if true 1 2)
+    let lp = sp((1, 0, 0), (1, 1, 1));
+    let kw = sp((1, 1, 1), (1, 3, 3));
+    let cond = sp((1, 4, 4), (1, 8, 8));
+    let then_t = sp((1, 9, 9), (1, 10, 10));
+    let else_t = sp((1, 11, 11), (1, 12, 12));
+    let rp = sp((1, 12, 12), (1, 13, 13));
+
+    let tokens = vec![
+        tsp(LeftPar, lp),
+        tsp(Symbol("if"), kw),
+        tsp(Symbol("true"), cond),
+        tsp(Number("1"), then_t),
+        tsp(Number("2"), else_t),
+        tsp(RightPar, rp),
+    ];
+    let prog = parse_sp(tokens.clone());
+
+    assert_eq!(
+        prog[0].span,
+        sp((1, 0, 0), (1, 13, 13)),
+        "input: {tokens:?}"
+    );
+    if let Expr::If {
+        condition,
+        then_expr,
+        else_expr,
+    } = &prog[0].inner
+    {
+        assert_eq!(condition.span, cond, "input: {tokens:?}");
+        assert_eq!(then_expr.span, then_t, "input: {tokens:?}");
+        assert_eq!(
+            else_expr.as_ref().unwrap().span,
+            else_t,
+            "input: {tokens:?}"
+        );
+    } else {
+        panic!("expected If, input: {tokens:?}");
+    }
+}
+
+#[test]
+fn test_span_if_no_else() {
+    // (if true 1)
+    let lp = sp((1, 0, 0), (1, 1, 1));
+    let kw = sp((1, 1, 1), (1, 3, 3));
+    let cond = sp((1, 4, 4), (1, 8, 8));
+    let then_t = sp((1, 9, 9), (1, 10, 10));
+    let rp = sp((1, 10, 10), (1, 11, 11));
+
+    let tokens = vec![
+        tsp(LeftPar, lp),
+        tsp(Symbol("if"), kw),
+        tsp(Symbol("true"), cond),
+        tsp(Number("1"), then_t),
+        tsp(RightPar, rp),
+    ];
+    let prog = parse_sp(tokens.clone());
+
+    assert_eq!(
+        prog[0].span,
+        sp((1, 0, 0), (1, 11, 11)),
+        "input: {tokens:?}"
+    );
+    if let Expr::If {
+        condition,
+        then_expr,
+        else_expr,
+    } = &prog[0].inner
+    {
+        assert_eq!(condition.span, cond, "input: {tokens:?}");
+        assert_eq!(then_expr.span, then_t, "input: {tokens:?}");
+        assert!(else_expr.is_none(), "input: {tokens:?}");
+    } else {
+        panic!("expected If, input: {tokens:?}");
+    }
+}
+
+#[test]
 fn test_span_for() {
     // (for i from 0 to 10 i)
     let lp = sp((1, 0, 0), (1, 1, 1));