From d7ec2131ef8d7b6db76c291e8418aaa9537a7ba7 Mon Sep 17 00:00:00 2001
From: KtorZ <matthias.benkort@gmail.com>
Date: Tue, 4 Jun 2024 10:35:21 +0200
Subject: [PATCH] Automatically merge import lines from same module.

  I slightly altered the way we parse import definitions to ensure we
  merge imports from the same modules (that aren't aliased) together.

  This prevents an annoying warning with duplicated import lines and
  makes it just more convenient overall.

  As a trade-off, we can no longer interleave import definitions with
  other definitions. This should be a minor setback only since the
  formatter was already ensuring that all import definitions would be
  grouped at the top.

  ---

  Note that, I originally attempted to implement this in the formatter
  instead of the parser. As it felt more appropriate there. However, the
  formatter operates on (unmutable) borrowed definitions, which makes it
  annoyingly hard to perform any AST manipulations. The `Document`
  returns by the format carries a lifetime that prevents the creation of
  intermediate local values.

  So instead, slightly tweaking the parser felt like the right thing to
  do.
---
 CHANGELOG.md                                  |  2 +
 crates/aiken-lang/src/parser.rs               | 56 ++++++++++++++++++-
 .../src/parser/definition/import.rs           | 24 ++++----
 .../aiken-lang/src/parser/definition/mod.rs   |  9 +--
 .../definition/snapshots/import_alias.snap    | 26 ++++-----
 .../definition/snapshots/import_basic.snap    | 22 ++++----
 .../snapshots/import_unqualified.snap         | 52 +++++++++--------
 crates/aiken-lang/src/parser/utils.rs         | 24 +++++++-
 .../src/snapshots/merge_imports.snap          | 51 +++++++++++++++++
 crates/aiken-lang/src/tests/format.rs         | 55 +++++++++++++++++-
 .../tests/snapshots/format_merge_imports.snap |  5 ++
 .../snapshots/format_merge_imports_2.snap     |  6 ++
 .../snapshots/format_merge_imports_alias.snap |  6 ++
 .../format_merge_imports_alias_2.snap         |  5 ++
 .../format_merge_imports_comments.snap        |  7 +++
 15 files changed, 272 insertions(+), 78 deletions(-)
 create mode 100644 crates/aiken-lang/src/snapshots/merge_imports.snap
 create mode 100644 crates/aiken-lang/src/tests/snapshots/format_merge_imports.snap
 create mode 100644 crates/aiken-lang/src/tests/snapshots/format_merge_imports_2.snap
 create mode 100644 crates/aiken-lang/src/tests/snapshots/format_merge_imports_alias.snap
 create mode 100644 crates/aiken-lang/src/tests/snapshots/format_merge_imports_alias_2.snap
 create mode 100644 crates/aiken-lang/src/tests/snapshots/format_merge_imports_comments.snap

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7f9ef41a..f1ba3bbc 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,8 @@
 
 - **aiken-lang**: the keyword `fail` on property-based test semantic has changed and now consider a test to succeed only if **every** execution of the test failed (instead of just one). The previous behavior can be recovered by adding the keyword `once` after `fail`. @KtorZ
 
+- **aiken-lang**: duplicate import lines are now automatically merged instead of raising a warning. However, imports can no longer appear anywhere in the file and must come as the first definitions. @KtorZ
+
 ### Fixed
 
 - **aiken-lang**: fixed the number of 'after x tests' number reported on property test failure, which was off by one. @KtorZ
diff --git a/crates/aiken-lang/src/parser.rs b/crates/aiken-lang/src/parser.rs
index d27aabc0..564045d7 100644
--- a/crates/aiken-lang/src/parser.rs
+++ b/crates/aiken-lang/src/parser.rs
@@ -13,10 +13,11 @@ mod utils;
 use crate::{ast, line_numbers::LineNumbers};
 pub use annotation::parser as annotation;
 use chumsky::prelude::*;
-pub use definition::parser as definition;
+pub use definition::{import::parser as import, parser as definition};
 use error::ParseError;
 pub use expr::parser as expression;
 use extra::ModuleExtra;
+use indexmap::IndexMap;
 pub use pattern::parser as pattern;
 
 pub fn module(
@@ -27,7 +28,48 @@ pub fn module(
 
     let stream = chumsky::Stream::from_iter(ast::Span::create(tokens.len(), 1), tokens.into_iter());
 
-    let definitions = definition().repeated().then_ignore(end()).parse(stream)?;
+    let definitions = import()
+        .repeated()
+        .map(|imports| {
+            let mut store = IndexMap::new();
+
+            for import in imports.into_iter() {
+                let key = (import.module, import.as_name);
+                match store.remove(&key) {
+                    None => {
+                        store.insert(key, (import.location, import.unqualified));
+                    }
+                    Some((location, unqualified)) => {
+                        let mut merged_unqualified = Vec::new();
+                        merged_unqualified.extend(unqualified);
+                        merged_unqualified.extend(import.unqualified);
+                        store.insert(key, (location, merged_unqualified));
+                    }
+                }
+            }
+
+            store
+                .into_iter()
+                .map(|((module, as_name), (location, unqualified))| {
+                    ast::Definition::Use(ast::Use {
+                        module,
+                        as_name,
+                        location,
+                        unqualified,
+                        package: (),
+                    })
+                })
+                .collect::<Vec<ast::UntypedDefinition>>()
+        })
+        .then(definition().repeated())
+        .map(|(imports, others)| {
+            let mut defs = Vec::new();
+            defs.extend(imports);
+            defs.extend(others);
+            defs
+        })
+        .then_ignore(end())
+        .parse(stream)?;
 
     let lines = LineNumbers::new(src);
 
@@ -47,6 +89,16 @@ pub fn module(
 mod tests {
     use crate::assert_module;
 
+    #[test]
+    fn merge_imports() {
+        assert_module!(
+            r#"
+            use aiken/list.{bar, foo}
+            use aiken/list.{baz}
+            "#
+        );
+    }
+
     #[test]
     fn windows_newline() {
         assert_module!("use aiken/list\r\n");
diff --git a/crates/aiken-lang/src/parser/definition/import.rs b/crates/aiken-lang/src/parser/definition/import.rs
index d98366c1..22a84d36 100644
--- a/crates/aiken-lang/src/parser/definition/import.rs
+++ b/crates/aiken-lang/src/parser/definition/import.rs
@@ -4,7 +4,7 @@ use crate::{
 };
 use chumsky::prelude::*;
 
-pub fn parser() -> impl Parser<Token, ast::UntypedDefinition, Error = ParseError> {
+pub fn parser() -> impl Parser<Token, ast::UntypedUse, Error = ParseError> {
     let unqualified_import = choice((
         select! {Token::Name { name } => name}.then(
             just(Token::As)
@@ -42,30 +42,28 @@ pub fn parser() -> impl Parser<Token, ast::UntypedDefinition, Error = ParseError
         .then(as_name);
 
     just(Token::Use).ignore_then(module_path).map_with_span(
-        |((module, unqualified), as_name), span| {
-            ast::UntypedDefinition::Use(ast::Use {
-                module,
-                as_name,
-                unqualified: unqualified.unwrap_or_default(),
-                package: (),
-                location: span,
-            })
+        |((module, unqualified), as_name), span| ast::Use {
+            module,
+            as_name,
+            unqualified: unqualified.unwrap_or_default(),
+            package: (),
+            location: span,
         },
     )
 }
 
 #[cfg(test)]
 mod tests {
-    use crate::assert_definition;
+    use crate::assert_import;
 
     #[test]
     fn import_basic() {
-        assert_definition!("use aiken/list");
+        assert_import!("use aiken/list");
     }
 
     #[test]
     fn import_unqualified() {
-        assert_definition!(
+        assert_import!(
             r#"
             use std/address.{Address as A, thing as w}
             "#
@@ -74,6 +72,6 @@ mod tests {
 
     #[test]
     fn import_alias() {
-        assert_definition!("use aiken/list as foo");
+        assert_import!("use aiken/list as foo");
     }
 }
diff --git a/crates/aiken-lang/src/parser/definition/mod.rs b/crates/aiken-lang/src/parser/definition/mod.rs
index 937e07fa..871c54f8 100644
--- a/crates/aiken-lang/src/parser/definition/mod.rs
+++ b/crates/aiken-lang/src/parser/definition/mod.rs
@@ -3,25 +3,22 @@ use chumsky::prelude::*;
 pub mod constant;
 mod data_type;
 mod function;
-mod import;
+pub mod import;
 mod test;
 mod type_alias;
 mod validator;
 
+use super::{error::ParseError, token::Token};
+use crate::ast;
 pub use constant::parser as constant;
 pub use data_type::parser as data_type;
 pub use function::parser as function;
-pub use import::parser as import;
 pub use test::parser as test;
 pub use type_alias::parser as type_alias;
 pub use validator::parser as validator;
 
-use super::{error::ParseError, token::Token};
-use crate::ast;
-
 pub fn parser() -> impl Parser<Token, ast::UntypedDefinition, Error = ParseError> {
     choice((
-        import(),
         data_type(),
         type_alias(),
         validator(),
diff --git a/crates/aiken-lang/src/parser/definition/snapshots/import_alias.snap b/crates/aiken-lang/src/parser/definition/snapshots/import_alias.snap
index 72220ba3..9e32d734 100644
--- a/crates/aiken-lang/src/parser/definition/snapshots/import_alias.snap
+++ b/crates/aiken-lang/src/parser/definition/snapshots/import_alias.snap
@@ -2,17 +2,15 @@
 source: crates/aiken-lang/src/parser/definition/import.rs
 description: "Code:\n\nuse aiken/list as foo"
 ---
-Use(
-    Use {
-        as_name: Some(
-            "foo",
-        ),
-        location: 0..21,
-        module: [
-            "aiken",
-            "list",
-        ],
-        package: (),
-        unqualified: [],
-    },
-)
+Use {
+    as_name: Some(
+        "foo",
+    ),
+    location: 0..21,
+    module: [
+        "aiken",
+        "list",
+    ],
+    package: (),
+    unqualified: [],
+}
diff --git a/crates/aiken-lang/src/parser/definition/snapshots/import_basic.snap b/crates/aiken-lang/src/parser/definition/snapshots/import_basic.snap
index 49015b8f..cac12581 100644
--- a/crates/aiken-lang/src/parser/definition/snapshots/import_basic.snap
+++ b/crates/aiken-lang/src/parser/definition/snapshots/import_basic.snap
@@ -2,15 +2,13 @@
 source: crates/aiken-lang/src/parser/definition/import.rs
 description: "Code:\n\nuse aiken/list"
 ---
-Use(
-    Use {
-        as_name: None,
-        location: 0..14,
-        module: [
-            "aiken",
-            "list",
-        ],
-        package: (),
-        unqualified: [],
-    },
-)
+Use {
+    as_name: None,
+    location: 0..14,
+    module: [
+        "aiken",
+        "list",
+    ],
+    package: (),
+    unqualified: [],
+}
diff --git a/crates/aiken-lang/src/parser/definition/snapshots/import_unqualified.snap b/crates/aiken-lang/src/parser/definition/snapshots/import_unqualified.snap
index a296b3a2..2310c47f 100644
--- a/crates/aiken-lang/src/parser/definition/snapshots/import_unqualified.snap
+++ b/crates/aiken-lang/src/parser/definition/snapshots/import_unqualified.snap
@@ -2,30 +2,28 @@
 source: crates/aiken-lang/src/parser/definition/import.rs
 description: "Code:\n\nuse std/address.{Address as A, thing as w}\n"
 ---
-Use(
-    Use {
-        as_name: None,
-        location: 0..42,
-        module: [
-            "std",
-            "address",
-        ],
-        package: (),
-        unqualified: [
-            UnqualifiedImport {
-                location: 17..29,
-                name: "Address",
-                as_name: Some(
-                    "A",
-                ),
-            },
-            UnqualifiedImport {
-                location: 31..41,
-                name: "thing",
-                as_name: Some(
-                    "w",
-                ),
-            },
-        ],
-    },
-)
+Use {
+    as_name: None,
+    location: 0..42,
+    module: [
+        "std",
+        "address",
+    ],
+    package: (),
+    unqualified: [
+        UnqualifiedImport {
+            location: 17..29,
+            name: "Address",
+            as_name: Some(
+                "A",
+            ),
+        },
+        UnqualifiedImport {
+            location: 31..41,
+            name: "thing",
+            as_name: Some(
+                "w",
+            ),
+        },
+    ],
+}
diff --git a/crates/aiken-lang/src/parser/utils.rs b/crates/aiken-lang/src/parser/utils.rs
index 08853bd4..49016fd3 100644
--- a/crates/aiken-lang/src/parser/utils.rs
+++ b/crates/aiken-lang/src/parser/utils.rs
@@ -1,6 +1,5 @@
-use chumsky::prelude::*;
-
 use super::{error::ParseError, token::Token};
+use chumsky::prelude::*;
 
 pub fn optional_flag(token: Token) -> impl Parser<Token, bool, Error = ParseError> {
     just(token).ignored().or_not().map(|v| v.is_some())
@@ -132,6 +131,27 @@ macro_rules! assert_definition {
     };
 }
 
+#[macro_export]
+macro_rules! assert_import {
+    ($code:expr) => {
+        use chumsky::Parser;
+
+        let $crate::parser::lexer::LexInfo { tokens, .. } = $crate::parser::lexer::run(indoc::indoc! { $code }).unwrap();
+
+        let stream = chumsky::Stream::from_iter($crate::ast::Span::create(tokens.len(), 1), tokens.into_iter());
+
+        let result = $crate::parser::import().parse(stream).unwrap();
+
+        insta::with_settings!({
+            description => concat!("Code:\n\n", indoc::indoc! { $code }),
+            prepend_module_to_snapshot => false,
+            omit_expression => true
+        }, {
+            insta::assert_debug_snapshot!(result);
+        });
+    };
+}
+
 #[macro_export]
 macro_rules! assert_format {
     ($code:expr) => {
diff --git a/crates/aiken-lang/src/snapshots/merge_imports.snap b/crates/aiken-lang/src/snapshots/merge_imports.snap
new file mode 100644
index 00000000..d8b27c86
--- /dev/null
+++ b/crates/aiken-lang/src/snapshots/merge_imports.snap
@@ -0,0 +1,51 @@
+---
+source: crates/aiken-lang/src/parser.rs
+description: "Code:\n\nuse aiken/list.{bar, foo}\nuse aiken/list.{baz}\n"
+---
+Module {
+    name: "",
+    docs: [],
+    type_info: (),
+    definitions: [
+        Use(
+            Use {
+                as_name: None,
+                location: 0..25,
+                module: [
+                    "aiken",
+                    "list",
+                ],
+                package: (),
+                unqualified: [
+                    UnqualifiedImport {
+                        location: 16..19,
+                        name: "bar",
+                        as_name: None,
+                    },
+                    UnqualifiedImport {
+                        location: 21..24,
+                        name: "foo",
+                        as_name: None,
+                    },
+                    UnqualifiedImport {
+                        location: 42..45,
+                        name: "baz",
+                        as_name: None,
+                    },
+                ],
+            },
+        ),
+    ],
+    lines: LineNumbers {
+        line_starts: [
+            0,
+            26,
+            47,
+        ],
+        length: 47,
+        last: Some(
+            47,
+        ),
+    },
+    kind: Validator,
+}
diff --git a/crates/aiken-lang/src/tests/format.rs b/crates/aiken-lang/src/tests/format.rs
index ce0b7128..c8a5f0c2 100644
--- a/crates/aiken-lang/src/tests/format.rs
+++ b/crates/aiken-lang/src/tests/format.rs
@@ -309,8 +309,6 @@ fn format_else_if() {
 
 #[test]
 fn format_imports() {
-    // TODO: Fix this case, this is behaving weirdly, not keeping the order for the comments and
-    // imports.
     assert_format!(
         r#"
         use aiken/list
@@ -324,6 +322,59 @@ fn format_imports() {
     );
 }
 
+#[test]
+fn format_merge_imports() {
+    assert_format!(
+        r#"
+        use aiken/list.{bar, foo}
+        use aiken/list.{baz}
+    "#
+    );
+}
+
+#[test]
+fn format_merge_imports_2() {
+    assert_format!(
+        r#"
+        use aiken/list.{bar, foo}
+        use aiken/dict
+        use aiken/list
+    "#
+    );
+}
+
+#[test]
+fn format_merge_imports_alias() {
+    assert_format!(
+        r#"
+        use aiken/list.{bar, foo}
+        use aiken/list.{baz} as vector
+    "#
+    );
+}
+
+#[test]
+fn format_merge_imports_alias_2() {
+    assert_format!(
+        r#"
+        use aiken/list.{bar, foo} as vector
+        use aiken/list.{baz} as vector
+    "#
+    );
+}
+
+#[test]
+fn format_merge_imports_comments() {
+    assert_format!(
+        r#"
+        // foo
+        use aiken/list.{bar, foo}
+        // bar
+        use aiken/list.{baz}
+    "#
+    );
+}
+
 #[test]
 fn format_negate() {
     assert_format!(
diff --git a/crates/aiken-lang/src/tests/snapshots/format_merge_imports.snap b/crates/aiken-lang/src/tests/snapshots/format_merge_imports.snap
new file mode 100644
index 00000000..41d9d73b
--- /dev/null
+++ b/crates/aiken-lang/src/tests/snapshots/format_merge_imports.snap
@@ -0,0 +1,5 @@
+---
+source: crates/aiken-lang/src/tests/format.rs
+description: "Code:\n\nuse aiken/list.{bar, foo}\nuse aiken/list.{baz}\n"
+---
+use aiken/list.{bar, baz, foo}
diff --git a/crates/aiken-lang/src/tests/snapshots/format_merge_imports_2.snap b/crates/aiken-lang/src/tests/snapshots/format_merge_imports_2.snap
new file mode 100644
index 00000000..70e57a99
--- /dev/null
+++ b/crates/aiken-lang/src/tests/snapshots/format_merge_imports_2.snap
@@ -0,0 +1,6 @@
+---
+source: crates/aiken-lang/src/tests/format.rs
+description: "Code:\n\nuse aiken/list.{bar, foo}\nuse aiken/dict\nuse aiken/list\n"
+---
+use aiken/dict
+use aiken/list.{bar, foo}
diff --git a/crates/aiken-lang/src/tests/snapshots/format_merge_imports_alias.snap b/crates/aiken-lang/src/tests/snapshots/format_merge_imports_alias.snap
new file mode 100644
index 00000000..a7d66b99
--- /dev/null
+++ b/crates/aiken-lang/src/tests/snapshots/format_merge_imports_alias.snap
@@ -0,0 +1,6 @@
+---
+source: crates/aiken-lang/src/tests/format.rs
+description: "Code:\n\nuse aiken/list.{bar, foo}\nuse aiken/list.{baz} as vector\n"
+---
+use aiken/list.{bar, foo}
+use aiken/list.{baz} as vector
diff --git a/crates/aiken-lang/src/tests/snapshots/format_merge_imports_alias_2.snap b/crates/aiken-lang/src/tests/snapshots/format_merge_imports_alias_2.snap
new file mode 100644
index 00000000..3dd1a37f
--- /dev/null
+++ b/crates/aiken-lang/src/tests/snapshots/format_merge_imports_alias_2.snap
@@ -0,0 +1,5 @@
+---
+source: crates/aiken-lang/src/tests/format.rs
+description: "Code:\n\nuse aiken/list.{bar, foo} as vector\nuse aiken/list.{baz} as vector\n"
+---
+use aiken/list.{bar, baz, foo} as vector
diff --git a/crates/aiken-lang/src/tests/snapshots/format_merge_imports_comments.snap b/crates/aiken-lang/src/tests/snapshots/format_merge_imports_comments.snap
new file mode 100644
index 00000000..528d29b8
--- /dev/null
+++ b/crates/aiken-lang/src/tests/snapshots/format_merge_imports_comments.snap
@@ -0,0 +1,7 @@
+---
+source: crates/aiken-lang/src/tests/format.rs
+description: "Code:\n\n// foo\nuse aiken/list.{bar, foo}\n// bar\nuse aiken/list.{baz}\n"
+---
+// foo
+use aiken/list.{bar, baz, foo}
+// bar