From be5a79a4b0dfedf8cdac1904081d3a453ff99f47 Mon Sep 17 00:00:00 2001
From: Yuna Tomida <ytomida.mmm@gmail.com>
Date: Tue, 8 Nov 2022 01:54:10 +0900
Subject: [PATCH] add FromIterator for Block/Module and make code succinct

---
 compiler/erg_parser/ast.rs     | 12 ++++++++++++
 compiler/erg_parser/desugar.rs | 28 ++++++++++------------------
 2 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/compiler/erg_parser/ast.rs b/compiler/erg_parser/ast.rs
index d9cd0414..ec879cf4 100644
--- a/compiler/erg_parser/ast.rs
+++ b/compiler/erg_parser/ast.rs
@@ -1124,6 +1124,12 @@ impl Locational for Block {
     }
 }
 
+impl FromIterator<Expr> for Block {
+    fn from_iter<T: IntoIterator<Item = Expr>>(iter: T) -> Self {
+        Self(iter.into_iter().collect())
+    }
+}
+
 impl_stream_for_wrapper!(Block, Expr);
 
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
@@ -3438,6 +3444,12 @@ impl Stream<Expr> for Module {
     }
 }
 
+impl FromIterator<Expr> for Module {
+    fn from_iter<T: IntoIterator<Item = Expr>>(iter: T) -> Self {
+        Self(iter.into_iter().collect())
+    }
+}
+
 impl Module {
     pub const fn empty() -> Self {
         Self(Block::empty())
diff --git a/compiler/erg_parser/desugar.rs b/compiler/erg_parser/desugar.rs
index 300612c3..ece1fd36 100644
--- a/compiler/erg_parser/desugar.rs
+++ b/compiler/erg_parser/desugar.rs
@@ -59,6 +59,10 @@ impl Desugarer {
         module
     }
 
+    fn desugar_all_chunks(module: Module, desugar: impl Fn(Expr) -> Expr) -> Module {
+        module.into_iter().map(desugar).collect()
+    }
+
     fn perform_desugar(desugar: impl Fn(Expr) -> Expr, expr: Expr) -> Expr {
         match expr {
             Expr::Record(record) => match record {
@@ -541,12 +545,8 @@ impl Desugarer {
     }
 
     /// `{x; y}` -> `{x = x; y = y}`
-    fn desugar_shortened_record(mut module: Module) -> Module {
-        let mut new = Module::with_capacity(module.len());
-        while let Some(chunk) = module.lpop() {
-            new.push(Self::rec_desugar_shortened_record(chunk));
-        }
-        new
+    fn desugar_shortened_record(module: Module) -> Module {
+        Self::desugar_all_chunks(module, Self::rec_desugar_shortened_record)
     }
 
     fn rec_desugar_shortened_record(expr: Expr) -> Expr {
@@ -859,12 +859,8 @@ impl Desugarer {
         }
     }
 
-    fn desugar_self(mut module: Module) -> Module {
-        let mut new = Module::with_capacity(module.len());
-        while let Some(chunk) = module.lpop() {
-            new.push(Self::desugar_self_inner(chunk));
-        }
-        new
+    fn desugar_self(module: Module) -> Module {
+        Self::desugar_all_chunks(module, Self::desugar_self_inner)
     }
 
     fn desugar_self_inner(_expr: Expr) -> Expr {
@@ -878,12 +874,8 @@ impl Desugarer {
 
     /// x[y] => x.__getitem__(y)
     /// x.0 => x.__Tuple_getitem__(0)
-    fn desugar_acc(mut module: Module) -> Module {
-        let mut new = Module::with_capacity(module.len());
-        while let Some(chunk) = module.lpop() {
-            new.push(Self::rec_desugar_acc(chunk));
-        }
-        new
+    fn desugar_acc(module: Module) -> Module {
+        Self::desugar_all_chunks(module, Self::rec_desugar_acc)
     }
 
     fn rec_desugar_acc(expr: Expr) -> Expr {