feat: (partial) comprehension support

crates/erg_compiler/context/initialize/funcs.rs

@@ -37,6 +37,13 @@ impl Context {
             Bool,
         );
         let t_ascii = nd_func(vec![kw(KW_OBJECT, Obj)], None, Str);
+        let t_array = func(
+            vec![],
+            None,
+            vec![kw(KW_ITERABLE, poly(ITERABLE, vec![ty_tp(T.clone())]))],
+            array_t(T.clone(), TyParam::erased(Nat)),
+        )
+        .quantify();
         let t_assert = func(
             vec![kw(KW_TEST, Bool)],
             None,
@@ -67,6 +74,16 @@ impl Context {
             T.clone(),
         )
         .quantify();
+        let t_dict = func(
+            vec![],
+            None,
+            vec![kw(
+                KW_ITERABLE,
+                poly(ITERABLE, vec![ty_tp(tuple_t(vec![T.clone(), U.clone()]))]),
+            )],
+            dict! { T.clone() => U.clone() }.into(),
+        )
+        .quantify();
         let t_discard = nd_func(vec![kw(KW_OBJ, Obj)], None, NoneType);
         let t_enumerate = func(
             vec![kw(KW_ITERABLE, poly(ITERABLE, vec![ty_tp(T.clone())]))],
@@ -221,6 +238,13 @@ impl Context {
         )
         .quantify();
         let t_round = nd_func(vec![kw(KW_NUMBER, Float)], None, Int);
+        let t_set = func(
+            vec![],
+            None,
+            vec![kw(KW_ITERABLE, poly(ITERABLE, vec![ty_tp(T.clone())]))],
+            set_t(T.clone(), TyParam::erased(Nat)),
+        )
+        .quantify();
         let t_slice = func(
             vec![kw(KW_START, Int)],
             None,
@@ -256,6 +280,7 @@ impl Context {
         self.register_py_builtin(FUNC_ABS, t_abs, Some(FUNC_ABS), 11);
         self.register_py_builtin(FUNC_ALL, t_all, Some(FUNC_ALL), 22);
         self.register_py_builtin(FUNC_ANY, t_any, Some(FUNC_ANY), 33);
+        self.register_py_builtin(FUNC_ARRAY, t_array, Some(FUNC_LIST), 215);
         self.register_py_builtin(FUNC_ASCII, t_ascii, Some(FUNC_ASCII), 53);
         // Leave as `Const`, as it may negatively affect assert casting.
         self.register_builtin_erg_impl(FUNC_ASSERT, t_assert, Const, vis.clone());
@@ -283,6 +308,7 @@ impl Context {
             Some(FUNC_COMPILE),
         );
         self.register_builtin_erg_impl(KW_COND, t_cond, Immutable, vis.clone());
+        self.register_py_builtin(FUNC_DICT, t_dict, Some(FUNC_DICT), 224);
         self.register_builtin_py_impl(
             FUNC_ENUMERATE,
             t_enumerate,
@@ -401,6 +427,7 @@ impl Context {
             vis.clone(),
             Some(FUNC_ROUND),
         );
+        self.register_py_builtin(FUNC_SET, t_set, Some(FUNC_SET), 233);
         self.register_builtin_py_impl(
             FUNC_SLICE,
             t_slice,

crates/erg_compiler/context/initialize/mod.rs

@@ -341,6 +341,7 @@ const GENERATOR: &str = "Generator";
 const FUNC_RANGE: &str = "range";
 const FUNC_ALL: &str = "all";
 const FUNC_ANY: &str = "any";
+const FUNC_ARRAY: &str = "array";
 const FUNC_ASCII: &str = "ascii";
 const FUNC_ASSERT: &str = "assert";
 const FUNC_BIN: &str = "bin";
@@ -360,6 +361,7 @@ const FUNC_POW: &str = "pow";
 const FUNC_QUIT: &str = "quit";
 const FUNC_REPR: &str = "repr";
 const FUNC_ROUND: &str = "round";
+const FUNC_SET: &str = "set";
 const FUNC_SLICE: &str = "slice";
 const FUNC_SORTED: &str = "sorted";
 const FUNC_SUM: &str = "sum";

crates/erg_compiler/context/instantiate_spec.rs

@@ -980,16 +980,18 @@ impl Context {
                 Ok(TyParam::Set(tp_set))
             }
             ast::ConstExpr::Set(ConstSet::Comprehension(set)) => {
-                if set.op.is(TokenKind::Colon) {
+                if set.layout.is_none() && set.generators.len() == 1 && set.guard.is_some() {
+                    let (ident, expr) = set.generators.get(0).unwrap();
                     let iter = self.instantiate_const_expr(
-                        &set.iter,
+                        expr,
                         erased_idx,
                         tmp_tv_cache,
                         not_found_is_qvar,
                     )?;
-                    let pred = self.instantiate_pred_from_expr(&set.pred, tmp_tv_cache)?;
+                    let pred =
+                        self.instantiate_pred_from_expr(set.guard.as_ref().unwrap(), tmp_tv_cache)?;
                     if let Ok(t) = self.instantiate_tp_as_type(iter, set) {
-                        return Ok(TyParam::t(refinement(set.var.inspect().clone(), t, pred)));
+                        return Ok(TyParam::t(refinement(ident.inspect().clone(), t, pred)));
                     }
                 }
                 type_feature_error!(

crates/erg_compiler/lib/pystd/builtins.d.er

@@ -204,3 +204,30 @@ opened in a binary mode.
   newline := Str or NoneType,
   closefd := Bool,
 ) => File!
+
+'''
+Convert `iterable` into an array.
+'''
+'''erg
+assert array() == []
+assert array((1, 2)) == [1, 2]
+'''
+.array: |T| (iterable := Iterable(T)) -> [T; _]
+
+'''
+Convert `iterable` into a dict.
+'''
+'''erg
+assert dict() == {:}
+assert dict([("a", 1), ("b": 2)]) == {"a": 1, "b": 2}
+'''
+.dict: |K, V| (iterable := Iterable((K, V))) -> {K: V}
+
+'''
+Convert `iterable` into a set.
+'''
+'''erg
+assert set() == {}
+assert set([1, 2]) == {1, 2}
+'''
+.set: |T| (iterable := Iterable(T)) -> {T; _}

crates/erg_parser/ast.rs

@@ -629,44 +629,44 @@ impl ArrayWithLength {
 pub struct ArrayComprehension {
     pub l_sqbr: Token,
     pub r_sqbr: Token,
-    pub elem: Box<Expr>,
+    pub layout: Option<Box<Expr>>,
     pub generators: Vec<(Identifier, Expr)>,
-    pub guards: Vec<Expr>,
+    pub guard: Option<Box<Expr>>,
 }
 
 impl NestedDisplay for ArrayComprehension {
     fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, _level: usize) -> fmt::Result {
         let mut generators = String::new();
         for (name, gen) in self.generators.iter() {
-            write!(generators, "{name} <- {gen}, ")?;
+            write!(generators, "{name} <- {gen}; ")?;
         }
         write!(
             f,
-            "[{}| {}{}]",
+            "[{}{}{}]",
-            self.elem,
+            fmt_option!(self.layout, post " | "),
             generators,
-            fmt_vec(&self.guards)
+            fmt_option!(pre " | ", &self.guard)
         )
     }
 }
 
 impl_display_from_nested!(ArrayComprehension);
-impl_locational!(ArrayComprehension, l_sqbr, elem, r_sqbr);
+impl_locational!(ArrayComprehension, l_sqbr, r_sqbr);
 
 impl ArrayComprehension {
     pub fn new(
         l_sqbr: Token,
         r_sqbr: Token,
-        elem: Expr,
+        layout: Option<Expr>,
         generators: Vec<(Identifier, Expr)>,
-        guards: Vec<Expr>,
+        guard: Option<Expr>,
     ) -> Self {
         Self {
             l_sqbr,
             r_sqbr,
-            elem: Box::new(elem),
+            layout: layout.map(Box::new),
             generators,
-            guards,
+            guard: guard.map(Box::new),
         }
     }
 }
@@ -779,27 +779,43 @@ impl NormalDict {
 pub struct DictComprehension {
     l_brace: Token,
     r_brace: Token,
-    pub attrs: Args,
+    pub kv: Box<KeyValue>,
-    guards: Vec<Expr>,
+    pub generators: Vec<(Identifier, Expr)>,
+    pub guard: Option<Box<Expr>>,
 }
 
-// TODO:
 impl NestedDisplay for DictComprehension {
     fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, _level: usize) -> fmt::Result {
-        write!(f, "{{{} | {}}}", self.attrs, fmt_vec(&self.guards))
+        let mut generators = String::new();
+        for (name, gen) in self.generators.iter() {
+            write!(generators, "{name} <- {gen}; ")?;
+        }
+        write!(
+            f,
+            "{{{} | {generators}{}}}",
+            self.kv,
+            fmt_option!(pre " | ", &self.guard)
+        )
     }
 }
 
 impl_display_from_nested!(DictComprehension);
-impl_locational!(DictComprehension, l_brace, attrs, r_brace);
+impl_locational!(DictComprehension, l_brace, kv, r_brace);
 
 impl DictComprehension {
-    pub const fn new(l_brace: Token, r_brace: Token, attrs: Args, guards: Vec<Expr>) -> Self {
+    pub fn new(
+        l_brace: Token,
+        r_brace: Token,
+        kv: KeyValue,
+        generators: Vec<(Identifier, Expr)>,
+        guard: Option<Expr>,
+    ) -> Self {
         Self {
             l_brace,
             r_brace,
-            attrs,
+            kv: Box::new(kv),
-            guards,
+            generators,
+            guard: guard.map(Box::new),
         }
     }
 }
@@ -1075,18 +1091,22 @@ impl SetWithLength {
 pub struct SetComprehension {
     pub l_brace: Token,
     pub r_brace: Token,
-    pub var: Token,
+    pub layout: Option<Box<Expr>>,
-    pub op: Token, // <- or :
+    pub generators: Vec<(Identifier, Expr)>,
-    pub iter: Box<Expr>,
+    pub guard: Option<Box<Expr>>,
-    pub pred: Box<Expr>,
 }
 
 impl NestedDisplay for SetComprehension {
     fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, _level: usize) -> fmt::Result {
+        let mut generators = String::new();
+        for (name, gen) in self.generators.iter() {
+            write!(generators, "{name} <- {gen}; ")?;
+        }
         write!(
             f,
-            "{{{} {} {} | {}}}",
+            "{{{}{generators}{}}}",
-            self.var, self.op, self.iter, self.pred
+            fmt_option!(self.layout, post " | "),
+            fmt_option!(pre " | ", self.guard)
         )
     }
 }
@@ -1098,18 +1118,16 @@ impl SetComprehension {
     pub fn new(
         l_brace: Token,
         r_brace: Token,
-        var: Token,
+        layout: Option<Expr>,
-        op: Token,
+        generators: Vec<(Identifier, Expr)>,
-        iter: Expr,
+        guard: Option<Expr>,
-        pred: Expr,
     ) -> Self {
         Self {
             l_brace,
             r_brace,
-            var,
+            layout: layout.map(Box::new),
-            op,
+            generators,
-            iter: Box::new(iter),
+            guard: guard.map(Box::new),
-            pred: Box::new(pred),
         }
     }
 }
@@ -1641,41 +1659,43 @@ impl ConstNormalSet {
 pub struct ConstSetComprehension {
     pub l_brace: Token,
     pub r_brace: Token,
-    pub var: Token,
+    pub layout: Option<Box<ConstExpr>>,
-    pub op: Token,
+    pub generators: Vec<(ConstIdentifier, ConstExpr)>,
-    pub iter: Box<ConstExpr>,
+    pub guard: Option<Box<ConstExpr>>,
-    pub pred: Box<ConstExpr>,
 }
 
 impl NestedDisplay for ConstSetComprehension {
     fn fmt_nest(&self, f: &mut fmt::Formatter<'_>, _level: usize) -> fmt::Result {
+        let mut generators = String::new();
+        for (name, gen) in self.generators.iter() {
+            write!(generators, "{name} <- {gen}, ")?;
+        }
         write!(
             f,
-            "{{{} {} {} | {}}}",
+            "{{{}{generators}{}}}",
-            self.var, self.op, self.iter, self.pred
+            fmt_option!(self.layout, post " | "),
+            fmt_option!(pre " | ", self.guard)
         )
     }
 }
 
 impl_display_from_nested!(ConstSetComprehension);
-impl_locational!(ConstSetComprehension, l_brace, var, r_brace);
+impl_locational!(ConstSetComprehension, l_brace, r_brace);
 
 impl ConstSetComprehension {
     pub fn new(
         l_brace: Token,
         r_brace: Token,
-        var: Token,
+        elem: Option<ConstExpr>,
-        op: Token,
+        generators: Vec<(ConstIdentifier, ConstExpr)>,
-        iter: ConstExpr,
+        guard: Option<ConstExpr>,
-        pred: ConstExpr,
     ) -> Self {
         Self {
             l_brace,
             r_brace,
-            var,
+            layout: elem.map(Box::new),
-            op,
+            generators,
-            iter: Box::new(iter),
+            guard: guard.map(Box::new),
-            pred: Box::new(pred),
         }
     }
 
@@ -1683,10 +1703,12 @@ impl ConstSetComprehension {
         SetComprehension::new(
             self.l_brace,
             self.r_brace,
-            self.var,
+            self.layout.map(|ex| ex.downgrade()),
-            self.op,
+            self.generators
-            self.iter.downgrade(),
+                .into_iter()
-            self.pred.downgrade(),
+                .map(|(name, gen)| (name, gen.downgrade()))
+                .collect(),
+            self.guard.map(|ex| ex.downgrade()),
         )
     }
 }
@@ -3353,6 +3375,26 @@ impl Identifier {
     pub fn trim_end_proc_mark(&mut self) {
         self.name.trim_end_proc_mark();
     }
+
+    pub fn call1(self, arg: Expr) -> Call {
+        Call::new(
+            self.into(),
+            None,
+            Args::pos_only(vec![PosArg::new(arg)], None),
+        )
+    }
+
+    pub fn call2(self, arg1: Expr, arg2: Expr) -> Call {
+        Call::new(
+            self.into(),
+            None,
+            Args::pos_only(vec![PosArg::new(arg1), PosArg::new(arg2)], None),
+        )
+    }
+
+    pub fn call(self, args: Args) -> Call {
+        Call::new(self.into(), None, args)
+    }
 }
 
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]

crates/erg_parser/desugar.rs

@@ -7,7 +7,7 @@
 use erg_common::error::Location;
 use erg_common::fresh::FreshNameGenerator;
 use erg_common::traits::{Locational, Stream};
-use erg_common::Str;
+use erg_common::{debug_power_assert, Str};
 use erg_common::{enum_unwrap, get_hash, log, set};
 
 use crate::ast::{
@@ -81,6 +81,7 @@ impl Desugarer {
         let module = Self::desugar_shortened_record(module);
         let module = Self::desugar_acc(module);
         let module = Self::desugar_operator(module);
+        let module = Self::desugar_comprehension(module);
         log!(info "AST (desugared):\n{module}");
         log!(info "the desugaring process has completed.");
         module
@@ -203,15 +204,15 @@ impl Desugarer {
                     Expr::Array(Array::WithLength(arr))
                 }
                 Array::Comprehension(arr) => {
-                    let elem = desugar(*arr.elem);
+                    let layout = arr.layout.map(|ex| desugar(*ex));
                     let generators = arr
                         .generators
                         .into_iter()
                         .map(|(ident, gen)| (ident, desugar(gen)))
                         .collect();
-                    let guards = arr.guards.into_iter().map(desugar).collect();
+                    let guard = arr.guard.map(|ex| desugar(*ex));
                     let arr =
-                        ArrayComprehension::new(arr.l_sqbr, arr.r_sqbr, elem, generators, guards);
+                        ArrayComprehension::new(arr.l_sqbr, arr.r_sqbr, layout, generators, guard);
                     Expr::Array(Array::Comprehension(arr))
                 }
             },
@@ -245,15 +246,18 @@ impl Desugarer {
                     Expr::Set(astSet::WithLength(set))
                 }
                 astSet::Comprehension(set) => {
-                    let iter = desugar(*set.iter);
+                    let elem = set.layout.map(|ex| desugar(*ex));
-                    let pred = desugar(*set.pred);
+                    let mut new_generators = vec![];
+                    for (ident, gen) in set.generators.into_iter() {
+                        new_generators.push((ident, desugar(gen)));
+                    }
+                    let new_guard = set.guard.map(|ex| desugar(*ex));
                     let set = SetComprehension::new(
                         set.l_brace,
                         set.r_brace,
-                        set.var,
+                        elem,
-                        set.op,
+                        new_generators,
-                        iter,
+                        new_guard,
-                        pred,
                     );
                     Expr::Set(astSet::Comprehension(set))
                 }
@@ -1586,6 +1590,100 @@ impl Desugarer {
             expr => Self::perform_desugar(Self::rec_desugar_operator, expr),
         }
     }
+
+    fn desugar_comprehension(module: Module) -> Module {
+        Self::desugar_all_chunks(module, Self::rec_desugar_comprehension)
+    }
+
+    /// ```erg
+    /// [y | x <- xs] ==> array(map(x -> y, xs))
+    /// [(a, b) | x <- xs; y <- ys] ==> array(map(((x, y),) -> (a, b), itertools.product(xs, ys)))
+    /// {k: v | x <- xs} ==> dict(map(x -> (k, v), xs))
+    /// {y | x <- xs} ==> set(map(x -> y, xs))
+    /// {x <- xs | x <= 10} ==> set(filter(x -> x <= 10, xs))
+    /// {x + 1 | x <- xs | x <= 10} ==> set(map(x -> x + 1, filter(x -> x <= 10, xs)))
+    /// ```
+    fn rec_desugar_comprehension(expr: Expr) -> Expr {
+        match expr {
+            Expr::Array(Array::Comprehension(mut comp)) => {
+                debug_power_assert!(comp.generators.len(), >, 0);
+                if comp.generators.len() != 1 {
+                    return Expr::Array(Array::Comprehension(comp));
+                }
+                let (ident, iter) = comp.generators.remove(0);
+                let iterator = Self::desugar_layout_and_guard(ident, iter, comp.layout, comp.guard);
+                Identifier::private("array".into())
+                    .call1(iterator.into())
+                    .into()
+            }
+            Expr::Dict(Dict::Comprehension(mut comp)) => {
+                debug_power_assert!(comp.generators.len(), >, 0);
+                if comp.generators.len() != 1 {
+                    return Expr::Dict(Dict::Comprehension(comp));
+                }
+                let (ident, iter) = comp.generators.remove(0);
+                let params = Params::single(NonDefaultParamSignature::new(
+                    ParamPattern::VarName(ident.name),
+                    None,
+                ));
+                let sig = LambdaSignature::new(params, None, TypeBoundSpecs::empty());
+                let tuple = Tuple::Normal(NormalTuple::new(Args::pos_only(
+                    vec![PosArg::new(comp.kv.key), PosArg::new(comp.kv.value)],
+                    None,
+                )));
+                let body = Block::new(vec![tuple.into()]);
+                let lambda = Lambda::new(sig, Token::DUMMY, body, DefId(0));
+                let map = Identifier::private("map".into()).call2(lambda.into(), iter);
+                Identifier::private("dict".into()).call1(map.into()).into()
+            }
+            Expr::Set(astSet::Comprehension(mut comp)) => {
+                debug_power_assert!(comp.generators.len(), >, 0);
+                if comp.generators.len() != 1 {
+                    return Expr::Set(astSet::Comprehension(comp));
+                }
+                let (ident, iter) = comp.generators.remove(0);
+                let iterator = Self::desugar_layout_and_guard(ident, iter, comp.layout, comp.guard);
+                Identifier::private("set".into())
+                    .call1(iterator.into())
+                    .into()
+            }
+            expr => Self::perform_desugar(Self::rec_desugar_comprehension, expr),
+        }
+    }
+
+    fn desugar_layout_and_guard(
+        ident: Identifier,
+        iter: Expr,
+        layout: Option<Box<Expr>>,
+        guard: Option<Box<Expr>>,
+    ) -> Call {
+        let params = Params::single(NonDefaultParamSignature::new(
+            ParamPattern::VarName(ident.name),
+            None,
+        ));
+        let sig = LambdaSignature::new(params, None, TypeBoundSpecs::empty());
+        match (layout, guard) {
+            (Some(elem), Some(guard)) => {
+                let f_body = Block::new(vec![*guard]);
+                let f_lambda = Lambda::new(sig.clone(), Token::DUMMY, f_body, DefId(0));
+                let filter = Identifier::private("filter".into()).call2(f_lambda.into(), iter);
+                let m_body = Block::new(vec![*elem]);
+                let m_lambda = Lambda::new(sig, Token::DUMMY, m_body, DefId(0));
+                Identifier::private("map".into()).call2(m_lambda.into(), filter.into())
+            }
+            (Some(elem), None) => {
+                let body = Block::new(vec![*elem]);
+                let lambda = Lambda::new(sig, Token::DUMMY, body, DefId(0));
+                Identifier::private("map".into()).call2(lambda.into(), iter)
+            }
+            (None, Some(guard)) => {
+                let body = Block::new(vec![*guard]);
+                let lambda = Lambda::new(sig, Token::DUMMY, body, DefId(0));
+                Identifier::private("filter".into()).call2(lambda.into(), iter)
+            }
+            (None, None) => todo!(),
+        }
+    }
 }
 
 impl Default for Desugarer {

crates/erg_parser/parse.rs

@@ -131,12 +131,26 @@ pub enum ArrayInner {
     Normal(Args),
     WithLength(PosArg, Expr),
     Comprehension {
-        elem: PosArg,
+        layout: Option<Expr>,
         generators: Vec<(Identifier, Expr)>,
-        guards: Vec<Expr>,
+        guard: Option<Expr>,
     },
 }
 
+impl ArrayInner {
+    pub const fn comp(
+        layout: Option<Expr>,
+        generators: Vec<(Identifier, Expr)>,
+        guard: Option<Expr>,
+    ) -> Self {
+        Self::Comprehension {
+            layout,
+            generators,
+            guard,
+        }
+    }
+}
+
 pub enum BraceContainer {
     Set(Set),
     Dict(Dict),
@@ -757,6 +771,29 @@ impl Parser {
         Ok(rest)
     }
 
+    fn try_reduce_ident(&mut self) -> ParseResult<Identifier> {
+        debug_call_info!(self);
+        let ident = match self.peek_kind() {
+            Some(Symbol) => {
+                let symbol = self.lpop();
+                Identifier::private_from_token(symbol)
+            }
+            Some(Dot) => {
+                let dot = self.lpop();
+                let symbol = expect_pop!(self, Symbol);
+                Identifier::public_from_token(dot, symbol)
+            }
+            _ => {
+                let err = self.skip_and_throw_syntax_err(line!(), caused_by!());
+                self.errs.push(err);
+                debug_exit_info!(self);
+                return Err(());
+            }
+        };
+        debug_exit_info!(self);
+        Ok(ident)
+    }
+
     fn try_reduce_acc_lhs(&mut self) -> ParseResult<Accessor> {
         debug_call_info!(self);
         let acc = match self.peek_kind() {
@@ -827,15 +864,61 @@ impl Parser {
                 debug_exit_info!(self);
                 return Ok(ArrayInner::WithLength(elems.remove_pos(0), len));
             }
-            Some(VBar) => {
+            Some(Inclusion) => {
-                let caused_by = caused_by!();
-                log!(err "error caused by: {caused_by}");
-                let err =
-                    ParseError::feature_error(line!() as usize, self.lpop().loc(), "comprehension");
                 self.lpop();
+                let Expr::Accessor(Accessor::Ident(sym)) = elems.remove_pos(0).expr else {
+                    let err = self.skip_and_throw_invalid_seq_err(
+                        caused_by!(),
+                        line!() as usize,
+                        &["identifier"],
+                        Inclusion,
+                    );
                     self.errs.push(err);
                     debug_exit_info!(self);
                     return Err(());
+                };
+                let mut generators = vec![];
+                let expr = self
+                    .try_reduce_expr(false, false, false, false)
+                    .map_err(|_| self.stack_dec(fn_name!()))?;
+                generators.push((sym, expr));
+                let _ = expect_pop!(self, VBar);
+                let guard = self
+                    .try_reduce_expr(false, false, false, false)
+                    .map_err(|_| self.stack_dec(fn_name!()))?;
+                debug_exit_info!(self);
+                return Ok(ArrayInner::comp(None, generators, Some(guard)));
+            }
+            Some(VBar) => {
+                self.lpop();
+                let elem = elems.remove_pos(0).expr;
+                let mut generators = vec![];
+                loop {
+                    let sym = self
+                        .try_reduce_ident()
+                        .map_err(|_| self.stack_dec(fn_name!()))?;
+                    let _ = expect_pop!(self, Inclusion);
+                    let expr = self
+                        .try_reduce_expr(false, false, false, false)
+                        .map_err(|_| self.stack_dec(fn_name!()))?;
+                    generators.push((sym, expr));
+                    if !self.cur_is(Semi) {
+                        break;
+                    } else {
+                        self.lpop();
+                    }
+                }
+                let guard = if self.cur_is(VBar) {
+                    self.lpop();
+                    let expr = self
+                        .try_reduce_expr(false, false, false, false)
+                        .map_err(|_| self.stack_dec(fn_name!()))?;
+                    Some(expr)
+                } else {
+                    None
+                };
+                debug_exit_info!(self);
+                return Ok(ArrayInner::comp(Some(elem), generators, guard));
             }
             Some(RParen | RSqBr | RBrace | Dedent | Comma) => {}
             Some(_) => {
@@ -2663,17 +2746,13 @@ impl Parser {
             ArrayInner::WithLength(elem, len) => {
                 Array::WithLength(ArrayWithLength::new(l_sqbr, r_sqbr, elem, len))
             }
-            ArrayInner::Comprehension { .. } => {
+            ArrayInner::Comprehension {
-                let caused_by = caused_by!();
+                layout,
-                log!(err "error caused by: {caused_by}");
+                generators,
-                self.errs.push(ParseError::feature_error(
+                guard,
-                    line!() as usize,
+            } => Array::Comprehension(ArrayComprehension::new(
-                    Location::concat(&l_sqbr, &r_sqbr),
+                l_sqbr, r_sqbr, layout, generators, guard,
-                    "array comprehension",
+            )),
-                ));
-                debug_exit_info!(self);
-                return Err(());
-            }
         };
         debug_exit_info!(self);
         Ok(arr)
@@ -2821,18 +2900,79 @@ impl Parser {
                 debug_exit_info!(self);
                 Ok(BraceContainer::Record(record))
             }
-            other if self.cur_is(Colon) => {
+            other => {
+                match self.peek_kind() {
+                    Some(Colon) => {
                         let res = self
-                    .try_reduce_normal_dict_or_set_comp(l_brace, other)
+                            .try_reduce_normal_dict_or_refine_type(l_brace, other)
                             .map_err(|_| self.stack_dec(fn_name!()))?;
                         debug_exit_info!(self);
-                Ok(res)
+                        return Ok(res);
                     }
-            other => {
+                    Some(Inclusion) => {
-                match self.peek() {
+                        self.skip();
-                    Some(r_brace) if r_brace.is(RBrace) => {
+                        let mut generators = vec![];
+                        let Expr::Accessor(Accessor::Ident(ident)) = other else {
+                            let caused_by = caused_by!();
+                            log!(err "error caused by: {caused_by}");
+                            let err = ParseError::invalid_record_element_err(
+                                line!() as usize,
+                                other.loc(),
+                            );
+                            self.errs.push(err);
+                            self.next_expr();
+                            debug_exit_info!(self);
+                            return Err(());
+                        };
+                        let expr = self
+                            .try_reduce_expr(false, false, false, false)
+                            .map_err(|_| self.stack_dec(fn_name!()))?;
+                        generators.push((ident, expr));
+                        let _ = expect_pop!(self, VBar);
+                        let guard = self
+                            .try_reduce_expr(false, false, false, false)
+                            .map_err(|_| self.stack_dec(fn_name!()))?;
+                        let r_brace = expect_pop!(self, fail_next RBrace);
+                        debug_exit_info!(self);
+                        let comp =
+                            SetComprehension::new(l_brace, r_brace, None, generators, Some(guard));
+                        return Ok(BraceContainer::Set(Set::Comprehension(comp)));
+                    }
+                    Some(VBar) => {
+                        self.skip();
+                        let mut generators = vec![];
+                        loop {
+                            let ident = self.try_reduce_ident()?;
+                            let _ = expect_pop!(self, fail_next Inclusion);
+                            let expr = self
+                                .try_reduce_expr(false, false, false, false)
+                                .map_err(|_| self.stack_dec(fn_name!()))?;
+                            generators.push((ident, expr));
+                            if self.cur_is(Semi) {
+                                self.skip();
+                            } else {
+                                break;
+                            }
+                        }
+                        let guard = if self.cur_is(VBar) {
+                            self.skip();
+                            let expr = self
+                                .try_reduce_expr(false, false, false, false)
+                                .map_err(|_| self.stack_dec(fn_name!()))?;
+                            Some(expr)
+                        } else {
+                            None
+                        };
+                        let r_brace = expect_pop!(self, fail_next RBrace);
+                        debug_exit_info!(self);
+                        let comp =
+                            SetComprehension::new(l_brace, r_brace, Some(other), generators, guard);
+                        return Ok(BraceContainer::Set(Set::Comprehension(comp)));
+                    }
+                    Some(RBrace) => {
                         let arg = Args::new(vec![PosArg::new(other)], None, vec![], None);
                         let r_brace = self.lpop();
+                        debug_exit_info!(self);
                         return Ok(BraceContainer::Set(Set::Normal(NormalSet::new(
                             l_brace, r_brace, arg,
                         ))));
@@ -2959,13 +3099,13 @@ impl Parser {
         }
     }
 
-    fn try_reduce_normal_dict_or_set_comp(
+    fn try_reduce_normal_dict_or_refine_type(
         &mut self,
         l_brace: Token,
         lhs: Expr,
     ) -> ParseResult<BraceContainer> {
         debug_call_info!(self);
-        let colon = expect_pop!(self, fail_next Colon);
+        let _colon = expect_pop!(self, fail_next Colon);
         let rhs = self
             .try_reduce_expr(false, true, false, false)
             .map_err(|_| self.stack_dec(fn_name!()))?;
@@ -2977,12 +3117,12 @@ impl Parser {
                 debug_exit_info!(self);
                 return Err(());
             };
-            let pred = self
+            let generators = vec![(var, rhs)];
+            let guard = self
                 .try_reduce_chunk(false, false)
                 .map_err(|_| self.stack_dec(fn_name!()))?;
             let r_brace = expect_pop!(self, fail_next RBrace);
-            let set_comp =
+            let set_comp = SetComprehension::new(l_brace, r_brace, None, generators, Some(guard));
-                SetComprehension::new(l_brace, r_brace, var.name.into_token(), colon, rhs, pred);
             debug_exit_info!(self);
             Ok(BraceContainer::Set(Set::Comprehension(set_comp)))
         } else {

crates/erg_parser/typespec.rs

@@ -63,15 +63,25 @@ impl Parser {
                     Ok(ConstExpr::Set(ConstSet::Normal(const_set)))
                 }
                 Set::Comprehension(set) => {
-                    let iter = Self::validate_const_expr(*set.iter)?;
+                    let elem = set
-                    let pred = Self::validate_const_expr(*set.pred)?;
+                        .layout
+                        .map(|ex| Self::validate_const_expr(*ex))
+                        .transpose()?;
+                    let mut generators = vec![];
+                    for (name, gen) in set.generators.into_iter() {
+                        let pred = Self::validate_const_expr(gen)?;
+                        generators.push((name, pred));
+                    }
+                    let guard = set
+                        .guard
+                        .map(|ex| Self::validate_const_expr(*ex))
+                        .transpose()?;
                     let const_set_comp = ConstSetComprehension::new(
                         set.l_brace,
                         set.r_brace,
-                        set.var,
+                        elem,
-                        set.op,
+                        generators,
-                        iter,
+                        guard,
-                        pred,
                     );
                     Ok(ConstExpr::Set(ConstSet::Comprehension(const_set_comp)))
                 }
@@ -368,10 +378,11 @@ impl Parser {
                 Ok(TypeSpec::SetWithLen(SetWithLenTypeSpec::new(t_spec, len)))
             }
             Set::Comprehension(set) => {
-                if set.op.is(TokenKind::Colon) {
+                if set.layout.is_none() && set.generators.len() == 1 && set.guard.is_some() {
-                    let typ = Self::expr_to_type_spec(*set.iter)?;
+                    let (ident, expr) = set.generators.into_iter().next().unwrap();
-                    let pred = Self::validate_const_expr(*set.pred)?;
+                    let typ = Self::expr_to_type_spec(expr)?;
-                    let refine = RefinementTypeSpec::new(set.var, typ, pred);
+                    let pred = Self::validate_const_expr(*set.guard.unwrap())?;
+                    let refine = RefinementTypeSpec::new(ident.name.into_token(), typ, pred);
                     Ok(TypeSpec::Refinement(refine))
                 } else {
                     Err(ParseError::simple_syntax_error(line!() as usize, set.loc()))

doc/EN/syntax/29_comprehension.md

@@ -31,27 +31,6 @@ assert {i % 2 | i <- 0..9} == {0, 1}
 assert {k: v | k <- ["a", "b"]; v <- [1, 2]} == {"a": 1, "b": 2}
 ```
 
-Erg comprehensions are inspired by Haskell, but with some differences.
-For Haskell list comprehensions, the order of variables makes a difference in the result, but in Erg it doesn't matter.
-
-``` haskell
--- Haskell
-[(i, j) | i <- [1..3], j <- [3..5]] == [(1,3),(1,4),(1,5),(2 ,3),(2,4),(2,5),(3,3),(3,4),(3,5)]
-[(i, j) | j <- [3..5], i <- [1..3]] == [(1,3),(2,3),(3,3),(1 ,4),(2,4),(3,4),(1,5),(2,5),(3,5)]
-```
-
-```python
-# Erg
-assert [(i, j) | i <- 1..<3; j <- 3..<5] == [(i, j) | j <- 3..<5; i <- 1.. <3]
-```
-
-This specification is the same as that of Python.
-
-```python
-# Python
-assert [(i, j) for i in range(1, 3) for j in range(3, 5)] == [(i, j) for j in range(3, 5) for i in range(1, 3)]
-```
-
 ## Refinement type
 
 Similar to comprehensions are refinement types. A refinement type is a type (enumerated type) created in the form `{Name: Type | Predicate}`.

doc/JA/syntax/29_comprehension.md

@@ -33,27 +33,6 @@ assert {i % 2 | i <- 0..9} == {0, 1}
 assert {k: v | k <- ["a", "b"]; v <- [1, 2]} == {"a": 1, "b": 2}
 ```
 
-Ergの内包表記はHaskellに影響を受けていますが、若干の違いがあります。
-Haskellのリスト内包表記の場合、変数の順番は結果に違いをもたらしますが、Ergでは関係がありません。
-
-```haskell
--- Haskell
-[(i, j) | i <- [1..3], j <- [3..5]] == [(1,3),(1,4),(1,5),(2,3),(2,4),(2,5),(3,3),(3,4),(3,5)]
-[(i, j) | j <- [3..5], i <- [1..3]] == [(1,3),(2,3),(3,3),(1,4),(2,4),(3,4),(1,5),(2,5),(3,5)]
-```
-
-```python
-# Erg
-assert [(i, j) | i <- 1..<3; j <- 3..<5] == [(i, j) | j <- 3..<5; i <- 1..<3]
-```
-
-この仕様はPythonのものと同じです。
-
-```python
-# Python
-assert [(i, j) for i in range(1, 3) for j in range(3, 5)] == [(i, j) for j in range(3, 5) for i in range(1, 3)]
-```
-
 ## 篩型
 
 内包表記と似たものに、篩型があります。篩型は`{Name: Type | Predicate}`という形式で作られる型(列挙型)です。

tests/should_ok/collection.er

@@ -2,6 +2,6 @@ arr = [[1, 2], [3, 4]]
 x = arr[0][0]
 print! x
 
-dict = {"a": {"a": 1}, "b": {"b": 2}}
+dic = {"a": {"a": 1}, "b": {"b": 2}}
-y = dict["a"]["a"]
+y = dic["a"]["a"]
 print! y

tests/should_ok/comprehension.er

@@ -0,0 +1,13 @@
+lc = [i * 2 | i <- 1..4]
+assert lc == [2, 4, 6, 8]
+lc2 = [i + 1 | i <- 1..5 | i <= 3]
+assert lc2 == [2, 3, 4]
+lc3 = [i <- 1..10 | i <= 5]
+assert lc3 == [1, 2, 3, 4, 5]
+
+sc = {i * 2 | i <- 1..4}
+assert sc == {2, 4, 6, 8}
+sc2 = {i + 1 | i <- 1..5 | i <= 3}
+assert sc2 == {2, 3, 4}
+sc3 = {i <- 1..10 | i <= 5}
+assert sc3 == {1, 2, 3, 4, 5}

tests/should_ok/mut_dict.er

@@ -1,6 +1,6 @@
 d = {"a": 1}
-dict = !d
+dic = !d
 
-dict.insert! "b", 2
+dic.insert! "b", 2
-assert dict.get("a") == 1
+assert dic.get("a") == 1
-assert dict.get("b") == 2
+assert dic.get("b") == 2

tests/test.rs

@@ -60,6 +60,11 @@ fn exec_comment() -> Result<(), ()> {
     expect_success("tests/should_ok/comment.er", 0)
 }
 
+#[test]
+fn exec_comprehension() -> Result<(), ()> {
+    expect_success("tests/should_ok/comprehension.er", 0)
+}
+
 #[test]
 fn exec_control() -> Result<(), ()> {
     expect_success("examples/control.er", 2)