fixed sort (use & mod) method in formatter (#6438)

crates/cairo-lang-formatter/Cargo.toml

@@ -18,7 +18,6 @@ ignore.workspace = true
 itertools = { workspace = true, default-features = true }
 salsa.workspace = true
 serde = { workspace = true, default-features = true }
-smol_str.workspace = true
 thiserror.workspace = true
 
 [dev-dependencies]

crates/cairo-lang-formatter/src/formatter_impl.rs

@@ -7,8 +7,6 @@ use cairo_lang_syntax::attribute::consts::FMT_SKIP_ATTR;
 use cairo_lang_syntax::node::db::SyntaxGroup;
 use cairo_lang_syntax::node::{ast, SyntaxNode, Terminal, TypedSyntaxNode};
 use itertools::Itertools;
-use smol_str::SmolStr;
-use syntax::node::ast::MaybeModuleBody;
 use syntax::node::helpers::QueryAttrs;
 use syntax::node::kind::SyntaxKind;
 
@@ -758,6 +756,9 @@ pub trait SyntaxNodeFormat {
     /// Otherwise, returns None.
     fn get_protected_zone_precedence(&self, db: &dyn SyntaxGroup) -> Option<usize>;
     fn should_skip_terminal(&self, db: &dyn SyntaxGroup) -> bool;
+    /// Returns the sorting kind of the syntax node. This method will be used to sections in the
+    /// syntax tree.
+    fn as_sort_kind(&self, db: &dyn SyntaxGroup) -> SortKind;
 }
 
 pub struct FormatterImpl<'a> {
@@ -825,9 +826,43 @@ impl<'a> FormatterImpl<'a> {
         // TODO(ilya): consider not copying here.
         let mut children = self.db.get_children(syntax_node.clone()).to_vec();
         let n_children = children.len();
+
         if self.config.sort_module_level_items {
-            children.sort_by_key(|c| MovableNode::new(self.db, c));
-        };
+            let mut start_idx = 0;
+            while start_idx < children.len() {
+                let kind = children[start_idx].as_sort_kind(self.db);
+                let mut end_idx = start_idx + 1;
+                // Find the end of the current section.
+                while end_idx < children.len() {
+                    if kind != children[end_idx].as_sort_kind(self.db) {
+                        break;
+                    }
+                    end_idx += 1;
+                }
+                // Sort within this section if it's `Module` or `UseItem`.
+                match kind {
+                    SortKind::Module => {
+                        children[start_idx..end_idx].sort_by_key(|node| {
+                            ast::ItemModule::from_syntax_node(self.db, node.clone())
+                                .name(self.db)
+                                .text(self.db)
+                        });
+                    }
+                    SortKind::UseItem => {
+                        children[start_idx..end_idx].sort_by_key(|node| {
+                            ast::ItemUse::from_syntax_node(self.db, node.clone())
+                                .use_path(self.db)
+                                .as_syntax_node()
+                                .get_text_without_trivia(self.db)
+                        });
+                    }
+                    SortKind::Immovable => {}
+                }
+
+                // Move past the sorted section.
+                start_idx = end_idx;
+            }
+        }
         for (i, child) in children.iter().enumerate() {
             if child.width(self.db) == TextWidth::default() {
                 continue;
@@ -922,52 +957,16 @@ impl<'a> FormatterImpl<'a> {
     }
 }
 
-/// Represents a sortable SyntaxNode.
+/// Represents the kind of sections in the syntax tree that can be sorted.
+/// Classify consecutive nodes into sections that are eligible for sorting.
 #[derive(PartialEq, Eq)]
-enum MovableNode {
-    ItemModule(SmolStr),
-    ItemUse(SmolStr),
-    ItemHeaderDoc,
-    Immovable,
-}
-impl MovableNode {
-    fn new(db: &dyn SyntaxGroup, node: &SyntaxNode) -> Self {
-        match node.kind(db) {
-            SyntaxKind::ItemModule => {
-                let item = ast::ItemModule::from_syntax_node(db, node.clone());
-                if matches!(item.body(db), MaybeModuleBody::None(_)) {
-                    Self::ItemModule(item.name(db).text(db))
-                } else {
-                    Self::Immovable
-                }
-            }
-            SyntaxKind::ItemUse => Self::ItemUse(node.clone().get_text_without_trivia(db).into()),
-            SyntaxKind::ItemHeaderDoc => Self::ItemHeaderDoc,
-            _ => Self::Immovable,
-        }
-    }
-}
+pub enum SortKind {
+    /// Module items without body, e.g. `mod a;`.
+    Module,
 
-impl Ord for MovableNode {
-    fn cmp(&self, other: &Self) -> Ordering {
-        match (self, other) {
-            (MovableNode::ItemHeaderDoc, _) => Ordering::Less,
-            (_, MovableNode::ItemHeaderDoc) => Ordering::Greater,
-            (MovableNode::Immovable, MovableNode::Immovable) => Ordering::Equal,
-            (MovableNode::ItemModule(a), MovableNode::ItemModule(b))
-            | (MovableNode::ItemUse(a), MovableNode::ItemUse(b)) => a.cmp(b),
-            (_, MovableNode::Immovable) | (MovableNode::ItemModule(_), MovableNode::ItemUse(_)) => {
-                Ordering::Less
-            }
-            (MovableNode::Immovable, _) | (MovableNode::ItemUse(_), MovableNode::ItemModule(_)) => {
-                Ordering::Greater
-            }
-        }
-    }
-}
+    /// Use items, e.g. `use a::b;` or `use c::{d, e as f};`.
+    UseItem,
 
-impl PartialOrd for MovableNode {
-    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
-        Some(self.cmp(other))
-    }
+    /// Items that cannot be moved - would be skipped and not included in any sorted segment.
+    Immovable,
 }

crates/cairo-lang-formatter/src/node_properties.rs

@@ -1,10 +1,12 @@
+use cairo_lang_syntax::node::ast::MaybeModuleBody;
 use cairo_lang_syntax::node::db::SyntaxGroup;
 use cairo_lang_syntax::node::kind::SyntaxKind;
 use cairo_lang_syntax::node::utils::{grandparent_kind, parent_kind};
 use cairo_lang_syntax::node::{ast, SyntaxNode, TypedSyntaxNode};
 
 use crate::formatter_impl::{
-    BreakLinePointIndentation, BreakLinePointProperties, BreakLinePointsPositions, SyntaxNodeFormat,
+    BreakLinePointIndentation, BreakLinePointProperties, BreakLinePointsPositions, SortKind,
+    SyntaxNodeFormat,
 };
 
 impl SyntaxNodeFormat for SyntaxNode {
@@ -720,6 +722,21 @@ impl SyntaxNodeFormat for SyntaxNode {
             false
         }
     }
+    // Merge the `as_sort_kind` method here
+    fn as_sort_kind(&self, db: &dyn SyntaxGroup) -> SortKind {
+        match self.kind(db) {
+            SyntaxKind::ItemModule => {
+                let item_module = ast::ItemModule::from_syntax_node(db, self.clone());
+                if matches!(item_module.body(db), MaybeModuleBody::None(_)) {
+                    SortKind::Module
+                } else {
+                    SortKind::Immovable
+                }
+            }
+            SyntaxKind::ItemUse => SortKind::UseItem,
+            _ => SortKind::Immovable,
+        }
+    }
 }
 
 /// For statement lists, returns if we want these as a single line.

crates/cairo-lang-formatter/src/test.rs

@@ -41,6 +41,11 @@ impl Upcast<dyn FilesGroup> for DatabaseImpl {
     "test_data/expected_results/fmt_skip.cairo",
     false
 )]
+#[test_case(
+    "test_data/cairo_files/sorted_mod_use.cairo",
+    "test_data/expected_results/sorted_mod_use.cairo",
+    true
+)]
 fn format_and_compare_file(unformatted_filename: &str, expected_filename: &str, use_sorting: bool) {
     let db_val = SimpleParserDatabase::default();
     let db = &db_val;

crates/cairo-lang-formatter/test_data/cairo_files/sorted_mod_use.cairo

@@ -0,0 +1,16 @@
+mod zeta;
+mod alpha;
+mod gamma;
+
+use std::fs;
+use std::io::Write;
+use std::path::PathBuf;
+use std::collections::HashMap;
+use std::env;
+
+mod beta;
+
+fn main() {
+    // Example function
+    println!("Hello, world!");
+}

crates/cairo-lang-formatter/test_data/expected_results/sorted_mod_use.cairo

@@ -0,0 +1,16 @@
+mod alpha;
+mod gamma;
+mod zeta;
+use std::collections::HashMap;
+use std::env;
+
+use std::fs;
+use std::io::Write;
+use std::path::PathBuf;
+
+mod beta;
+
+fn main() {
+    // Example function
+    println!("Hello, world!");
+}

crates/cairo-lang-formatter/test_data/expected_results/use_sorting.cairo

@@ -6,23 +6,20 @@ use openzeppelin::introspection::interface;
 #[starknet::contract]
 mod SRC5 {
     //! Header comment, should not be moved by the formatter.
-    mod F;
-    mod G;
-
-    use A;
-
-    use openzeppelin::introspection::first;
     /// Doc comment, should be moved by the formatter.
     use openzeppelin::introspection::interface;
     use openzeppelin::introspection::{interface, AB};
-    use starknet::ArrayTrait;
 
     #[storage]
     struct Storage {
         supported_interfaces: LegacyMap<felt252, bool>
     }
 
+    use openzeppelin::introspection::first;
+
     mod A {}
+    mod F;
+    mod G;
 
     #[abi(embed_v0)]
     impl SRC5Impl of interface::ISRC5<ContractState> {
@@ -31,6 +28,9 @@ mod SRC5 {
         }
     }
 
+    use A;
+    use starknet::ArrayTrait;
+
     mod Inner {
         use B;
         use C;