From 24b3219016f023b849bece74725a362c9d881bf6 Mon Sep 17 00:00:00 2001
From: imaqtkatt <eduarda.so@proton.me>
Date: Mon, 27 May 2024 08:15:36 -0300
Subject: [PATCH] Implement elif

---
 src/imp/parser.rs                             | 22 +++++++++++++++++++
 tests/golden_tests/compile_file/elif.bend     | 16 ++++++++++++++
 .../compile_file/elif_no_else.bend            |  5 +++++
 tests/snapshots/compile_file__elif.bend.snap  | 18 +++++++++++++++
 .../compile_file__elif_no_else.bend.snap      |  8 +++++++
 5 files changed, 69 insertions(+)
 create mode 100644 tests/golden_tests/compile_file/elif.bend
 create mode 100644 tests/golden_tests/compile_file/elif_no_else.bend
 create mode 100644 tests/snapshots/compile_file__elif.bend.snap
 create mode 100644 tests/snapshots/compile_file__elif_no_else.bend.snap

diff --git a/src/imp/parser.rs b/src/imp/parser.rs
index 565649f3..293e3428 100644
--- a/src/imp/parser.rs
+++ b/src/imp/parser.rs
@@ -532,6 +532,21 @@ impl<'a> PyParser<'a> {
     if nxt_indent != *indent {
       return self.expected_indent(*indent, nxt_indent);
     }
+    let mut elifs = Vec::new();
+    while self.try_parse_keyword("elif") {
+      let cond = self.parse_expr(true)?;
+      self.skip_trivia_inline();
+      self.consume_exactly(":")?;
+      indent.enter_level();
+      self.consume_indent_exactly(*indent)?;
+      let (then, nxt_indent) = self.parse_statement(indent)?;
+      indent.exit_level();
+
+      if nxt_indent != *indent {
+        return self.expected_indent(*indent, nxt_indent);
+      }
+      elifs.push((cond, then));
+    }
     self.parse_keyword("else")?;
     self.skip_trivia_inline();
     self.consume_exactly(":")?;
@@ -539,6 +554,13 @@ impl<'a> PyParser<'a> {
 
     self.consume_indent_exactly(*indent)?;
     let (otherwise, nxt_indent) = self.parse_statement(indent)?;
+    let otherwise = elifs.into_iter().fold(otherwise, |acc, (cond, then)| Stmt::If {
+      cond: Box::new(cond),
+      then: Box::new(then),
+      otherwise: Box::new(acc),
+      nxt: None,
+    });
+
     indent.exit_level();
     if nxt_indent == *indent {
       let (nxt, nxt_indent) = self.parse_statement(indent)?;
diff --git a/tests/golden_tests/compile_file/elif.bend b/tests/golden_tests/compile_file/elif.bend
new file mode 100644
index 00000000..222d4456
--- /dev/null
+++ b/tests/golden_tests/compile_file/elif.bend
@@ -0,0 +1,16 @@
+def main:
+  cond1 = 1 == 2
+  cond2 = 2 < 1
+  cond3 = 3 > 2
+  cond4 = 2 == 2
+  if cond1:
+    res = 1
+  elif cond2:
+    res = 2
+  elif cond3:
+    res = 3
+  elif cond4:
+    res = 4
+  else:
+    res = 0
+  return res
diff --git a/tests/golden_tests/compile_file/elif_no_else.bend b/tests/golden_tests/compile_file/elif_no_else.bend
new file mode 100644
index 00000000..0f25c55d
--- /dev/null
+++ b/tests/golden_tests/compile_file/elif_no_else.bend
@@ -0,0 +1,5 @@
+def main:
+  if 1 == 1:
+    return 0
+  elif 2 == 2:
+    return 1
diff --git a/tests/snapshots/compile_file__elif.bend.snap b/tests/snapshots/compile_file__elif.bend.snap
new file mode 100644
index 00000000..c2b50647
--- /dev/null
+++ b/tests/snapshots/compile_file__elif.bend.snap
@@ -0,0 +1,18 @@
+---
+source: tests/golden_tests.rs
+input_file: tests/golden_tests/compile_file/elif.bend
+---
+@main = d
+  & @main__C3 ~ (a (b (c d)))
+  & $(1 a) ~ [<2]
+  & $(2 b) ~ [>3]
+  & $(2 c) ~ [=2]
+
+@main__C0 = (?((0 (* 2)) a) a)
+
+@main__C1 = (a (?((@main__C0 (* (* 3))) (a b)) b))
+
+@main__C2 = (a (b (?((@main__C1 (* (* (* 4)))) (a (b c))) c)))
+
+@main__C3 = a
+  & $(2 ?((@main__C2 (* (* (* (* 1))))) a)) ~ [=1]
diff --git a/tests/snapshots/compile_file__elif_no_else.bend.snap b/tests/snapshots/compile_file__elif_no_else.bend.snap
new file mode 100644
index 00000000..4a120f89
--- /dev/null
+++ b/tests/snapshots/compile_file__elif_no_else.bend.snap
@@ -0,0 +1,8 @@
+---
+source: tests/golden_tests.rs
+input_file: tests/golden_tests/compile_file/elif_no_else.bend
+---
+Errors:
+In tests/golden_tests/compile_file/elif_no_else.bend :
+Indentation error. Expected 2 spaces, got end-of-input.
+  6 |