From a4eec067149d137f2cab07fe741ffc7bab73771a Mon Sep 17 00:00:00 2001
From: Lucas Franceschino <lucas.franceschino@inria.fr>
Date: Tue, 13 Jun 2023 14:05:35 +0200
Subject: [PATCH] Steal consts items before anything else

---
 cli/driver/src/exporter.rs | 41 +++++++++++++++++++++++---------------
 1 file changed, 25 insertions(+), 16 deletions(-)

diff --git a/cli/driver/src/exporter.rs b/cli/driver/src/exporter.rs
index c386477cc..3fce22755 100644
--- a/cli/driver/src/exporter.rs
+++ b/cli/driver/src/exporter.rs
@@ -25,29 +25,38 @@ fn convert_thir<'tcx>(
     tcx: TyCtxt<'tcx>,
 ) -> (Vec<rustc_span::Span>, Vec<hax_frontend_exporter::Item>) {
     let hir = tcx.hir();
-    let bodies = hir.body_owners();
 
-    let mut bodies = bodies.collect::<Vec<_>>();
-    // we first visit `AnonConst`s, otherwise the thir body might be stolen
-    bodies.sort_by(|a, b| {
-        use std::cmp::Ordering::*;
-        let is_anon_const = |x: &rustc_span::def_id::LocalDefId| {
-            matches!(hir.get_by_def_id(x.clone()), rustc_hir::Node::AnonConst(_))
+    let bodies = {
+        // Here, we partition the bodies so that constant items appear
+        // first.
+        let mut is_const = |x: &rustc_span::def_id::LocalDefId| {
+            matches!(
+                hir.get_by_def_id(x.clone()),
+                rustc_hir::Node::AnonConst(_)
+                    | rustc_hir::Node::Item(rustc_hir::Item {
+                        kind: rustc_hir::ItemKind::Const(..),
+                        ..
+                    })
+                    | rustc_hir::Node::TraitItem(rustc_hir::TraitItem {
+                        kind: rustc_hir::TraitItemKind::Const(..),
+                        ..
+                    })
+                    | rustc_hir::Node::ImplItem(rustc_hir::ImplItem {
+                        kind: rustc_hir::ImplItemKind::Const(..),
+                        ..
+                    })
+            )
         };
-        if is_anon_const(a) {
-            Less
-        } else if is_anon_const(b) {
-            Equal
-        } else {
-            Greater
-        }
-    });
+
+        let (consts, others): (Vec<rustc_span::def_id::LocalDefId>, _) =
+            hir.body_owners().partition(is_const);
+        consts.into_iter().chain(others.into_iter())
+    };
 
     let thirs: std::collections::HashMap<
         rustc_span::def_id::LocalDefId,
         (rustc_middle::thir::Thir<'tcx>, ExprId),
     > = bodies
-        .into_iter()
         .map(|did| {
             let (thir, expr) = tcx
                 .thir_body(rustc_middle::ty::WithOptConstParam {