From e13e36011887b29a7a206d8dedf31249b319c088 Mon Sep 17 00:00:00 2001
From: imaqtkatt <eduarda.so@proton.me>
Date: Fri, 3 May 2024 11:24:47 -0300
Subject: [PATCH] Add list comprehension

---
 src/term/flavour_py/mod.rs          |  2 ++
 src/term/flavour_py/order_kwargs.rs |  1 +
 src/term/flavour_py/parser.rs       | 30 +++++++++++++++++++++++++----
 src/term/flavour_py/to_lang.rs      |  1 +
 4 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/src/term/flavour_py/mod.rs b/src/term/flavour_py/mod.rs
index d64b722e..a69b7747 100644
--- a/src/term/flavour_py/mod.rs
+++ b/src/term/flavour_py/mod.rs
@@ -29,6 +29,8 @@ pub enum Term {
   Lst { els: Vec<Term> },
   // "(" ... ")"
   Tup { els: Vec<Term> },
+  // "[" {term} "for" {bind} "in" {iter} ("if" {cond})? "]"
+  Comprehension { term: Box<Term>, bind: Name, iter: Box<Term>, cond: Option<Box<Term>> },
 }
 
 #[derive(Clone, Debug)]
diff --git a/src/term/flavour_py/order_kwargs.rs b/src/term/flavour_py/order_kwargs.rs
index b28017ec..1121a268 100644
--- a/src/term/flavour_py/order_kwargs.rs
+++ b/src/term/flavour_py/order_kwargs.rs
@@ -88,6 +88,7 @@ impl Term {
         rhs.order_kwargs(ctx);
       }
       Term::Lst { els } | Term::Tup { els } => els.iter_mut().for_each(|e| e.order_kwargs(ctx)),
+      Term::Comprehension { .. } => {}
       Term::None | Term::Var { .. } | Term::Num { .. } | Term::Str { .. } => {}
     }
   }
diff --git a/src/term/flavour_py/parser.rs b/src/term/flavour_py/parser.rs
index bc42d7ee..90c357bf 100644
--- a/src/term/flavour_py/parser.rs
+++ b/src/term/flavour_py/parser.rs
@@ -65,10 +65,7 @@ impl<'a> PyParser<'a> {
           head
         }
       }
-      '[' => {
-        let els = self.list_like(|p| p.parse_term_py(), "[", "]", ",", true, 0)?;
-        Term::Lst { els }
-      }
+      '[' => self.list_or_comprehension()?,
       '\"' => {
         let str = self.parse_quoted_string()?;
         let val = STRINGS.get(str);
@@ -90,6 +87,28 @@ impl<'a> PyParser<'a> {
     Ok(res)
   }
 
+  fn list_or_comprehension(&mut self) -> Result<Term, String> {
+    self.consume("[")?;
+    let head = self.parse_term_py()?;
+    if self.try_consume_keyword("for") {
+      let bind = self.parse_hvml_name()?;
+      self.consume("in")?;
+      let iter = self.parse_term_py()?;
+      let mut cond = None;
+      if self.try_consume_keyword("if") {
+        cond = Some(Box::new(self.parse_term_py()?));
+      }
+      self.consume("]")?;
+      Ok(Term::Comprehension { term: Box::new(head), bind, iter: Box::new(iter), cond })
+    } else {
+      let mut head = vec![head];
+      self.try_consume(",");
+      let tail = self.list_like(|p| p.parse_term_py(), "", "]", ",", true, 0)?;
+      head.extend(tail);
+      Ok(Term::Lst { els: head })
+    }
+  }
+
   fn parse_term_py(&mut self) -> Result<Term, String> {
     self.skip_trivia();
     if self.try_consume_keyword("lambda") {
@@ -562,6 +581,9 @@ def inc(n):
   n += 1;
   return n;
 
+def inc_list(list):
+  return [x+1 for x in list];
+
 def lam():
   return lambda x, y: return x;;
 
diff --git a/src/term/flavour_py/to_lang.rs b/src/term/flavour_py/to_lang.rs
index 9005f5e9..065b32d1 100644
--- a/src/term/flavour_py/to_lang.rs
+++ b/src/term/flavour_py/to_lang.rs
@@ -109,6 +109,7 @@ impl Term {
         tag: lang::Tag::Static,
         els: els.into_iter().map(Self::to_lang).collect(),
       },
+      Term::Comprehension { .. } => todo!(),
     }
   }
 }