From de34a9135069a559e3a2249b481fe551a5c1a8dd Mon Sep 17 00:00:00 2001
From: Ralf Jung <post@ralfj.de>
Date: Thu, 29 Aug 2024 16:53:07 +0200
Subject: [PATCH] interpret/visitor: make memory order iteration slightly more
 efficient

---
 .../rustc_const_eval/src/interpret/visitor.rs | 19 ++++++++++---------
 src/tools/miri/src/helpers.rs                 | 13 +++++--------
 2 files changed, 15 insertions(+), 17 deletions(-)

diff --git a/compiler/rustc_const_eval/src/interpret/visitor.rs b/compiler/rustc_const_eval/src/interpret/visitor.rs
index fd649d608c691..b02f12e3c7f0b 100644
--- a/compiler/rustc_const_eval/src/interpret/visitor.rs
+++ b/compiler/rustc_const_eval/src/interpret/visitor.rs
@@ -25,14 +25,15 @@ pub trait ValueVisitor<'tcx, M: Machine<'tcx>>: Sized {
     }
 
     /// This function provides the chance to reorder the order in which fields are visited for
-    /// `FieldsShape::Aggregate`: The order of fields will be
-    /// `(0..num_fields).map(aggregate_field_order)`.
+    /// `FieldsShape::Aggregate`.
     ///
-    /// The default means we iterate in source declaration order; alternative this can do an inverse
-    /// lookup in `memory_index` to use memory field order instead.
+    /// The default means we iterate in source declaration order; alternatively this can do some
+    /// work with `memory_index` to iterate in memory order.
     #[inline(always)]
-    fn aggregate_field_order(_memory_index: &IndexVec<FieldIdx, u32>, idx: usize) -> usize {
-        idx
+    fn aggregate_field_iter(
+        memory_index: &IndexVec<FieldIdx, u32>,
+    ) -> impl Iterator<Item = FieldIdx> + 'static {
+        memory_index.indices()
     }
 
     // Recursive actions, ready to be overloaded.
@@ -172,9 +173,9 @@ pub trait ValueVisitor<'tcx, M: Machine<'tcx>>: Sized {
             &FieldsShape::Union(fields) => {
                 self.visit_union(v, fields)?;
             }
-            FieldsShape::Arbitrary { offsets, memory_index } => {
-                for idx in 0..offsets.len() {
-                    let idx = Self::aggregate_field_order(memory_index, idx);
+            FieldsShape::Arbitrary { memory_index, .. } => {
+                for idx in Self::aggregate_field_iter(memory_index) {
+                    let idx = idx.as_usize();
                     let field = self.ecx().project_field(v, idx)?;
                     self.visit_field(v, idx, &field)?;
                 }
diff --git a/src/tools/miri/src/helpers.rs b/src/tools/miri/src/helpers.rs
index 0483745621252..a546ad20fef98 100644
--- a/src/tools/miri/src/helpers.rs
+++ b/src/tools/miri/src/helpers.rs
@@ -630,14 +630,11 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                 self.ecx
             }
 
-            fn aggregate_field_order(memory_index: &IndexVec<FieldIdx, u32>, idx: usize) -> usize {
-                // We need to do an *inverse* lookup: find the field that has position `idx` in memory order.
-                for (src_field, &mem_pos) in memory_index.iter_enumerated() {
-                    if mem_pos as usize == idx {
-                        return src_field.as_usize();
-                    }
-                }
-                panic!("invalid `memory_index`, could not find {}-th field in memory order", idx);
+            fn aggregate_field_iter(
+                memory_index: &IndexVec<FieldIdx, u32>,
+            ) -> impl Iterator<Item = FieldIdx> + 'static {
+                let inverse_memory_index = memory_index.invert_bijective_mapping();
+                inverse_memory_index.into_iter()
             }
 
             // Hook to detect `UnsafeCell`.