Group by stack slot earlier

1 files changed, 42 insertions, 107 deletions

diff --git a/src/optimize/stack2reg.rs b/src/optimize/stack2reg.rs
index 314561cb7d7..34f75e7da00 100644
--- a/src/optimize/stack2reg.rs
+++ b/src/optimize/stack2reg.rs
@@ -1,4 +1,5 @@
 use std::collections::{BTreeMap, BTreeSet, HashSet};
+use std::ops::Not;
 
 use cranelift_codegen::cursor::{Cursor, FuncCursor};
 use cranelift_codegen::ir::{Opcode, InstructionData, ValueDef};
@@ -22,6 +23,13 @@ impl Ord for OrdStackSlot {
     }
 }
 
+#[derive(Debug, Default)]
+struct StackSlotUsage {
+    stack_addr: HashSet<Inst>,
+    stack_load: HashSet<Inst>,
+    stack_store: HashSet<Inst>,
+}
+
 pub(super) fn optimize_function(
     func: &mut Function,
     clif_comments: &mut crate::pretty_clif::CommentWriter,
@@ -31,10 +39,8 @@ pub(super) fn optimize_function(
 
     // Record all stack_addr, stack_load and stack_store instructions. Also record all stack_addr
     // and stack_load insts whose result is used.
-    let mut stack_addr_load_insts_users = BTreeMap::<Inst, HashSet<Inst>>::new();
-    let mut stack_addr_insts = BTreeSet::new();
-    let mut stack_load_insts = BTreeSet::new();
-    let mut stack_store_insts = BTreeSet::new();
+    let mut stack_addr_load_insts_users = HashMap::<Inst, HashSet<Inst>>::new();
+    let mut stack_slot_usage_map = BTreeMap::<OrdStackSlot, StackSlotUsage>::new();
 
     let mut cursor = FuncCursor::new(func);
     while let Some(_ebb) = cursor.next_ebb() {
@@ -42,25 +48,25 @@ pub(super) fn optimize_function(
             match cursor.func.dfg[inst] {
                 InstructionData::StackLoad {
                     opcode: Opcode::StackAddr,
-                    stack_slot: _,
+                    stack_slot,
                     offset: _,
                 } => {
-                    stack_addr_insts.insert(inst);
+                    stack_slot_usage_map.entry(OrdStackSlot(stack_slot)).or_insert_with(StackSlotUsage::default).stack_addr.insert(inst);
                 }
                 InstructionData::StackLoad {
                     opcode: Opcode::StackLoad,
-                    stack_slot: _,
+                    stack_slot,
                     offset: _,
                 } => {
-                    stack_load_insts.insert(inst);
+                    stack_slot_usage_map.entry(OrdStackSlot(stack_slot)).or_insert_with(StackSlotUsage::default).stack_load.insert(inst);
                 }
                 InstructionData::StackStore {
                     opcode: Opcode::StackStore,
                     arg: _,
-                    stack_slot: _,
+                    stack_slot,
                     offset: _,
                 } => {
-                    stack_store_insts.insert(inst);
+                    stack_slot_usage_map.entry(OrdStackSlot(stack_slot)).or_insert_with(StackSlotUsage::default).stack_store.insert(inst);
                 }
                 _ => {}
             }
@@ -79,125 +85,54 @@ pub(super) fn optimize_function(
     }
 
     println!(
-        "{}:\nstack_addr/stack_load users: {:?}\nstack_addr: {:?}\nstack_load: {:?}\nstack_store: {:?}",
+        "{}:\nstack_addr/stack_load users: {:?}\nstack slot usage: {:?}",
         name,
         stack_addr_load_insts_users,
-        stack_addr_insts,
-        stack_load_insts,
-        stack_store_insts,
+        stack_slot_usage_map,
     );
 
     for inst in stack_addr_load_insts_users.keys() {
-        assert!(stack_addr_insts.contains(inst) || stack_load_insts.contains(inst));
-    }
-
-    // Replace all unused stack_addr instructions with nop.
-    // FIXME remove clone
-    for &inst in stack_addr_insts.clone().iter() {
-        if stack_addr_load_insts_users.get(&inst).map(|users| users.is_empty()).unwrap_or(true) {
-            println!("Removing unused stack_addr {}", inst);
-            func.dfg.detach_results(inst);
-            func.dfg.replace(inst).nop();
-            stack_addr_insts.remove(&inst);
-        }
-    }
-
-    // Replace all unused stack_load instructions with nop.
-    // FIXME remove clone
-    for &inst in stack_load_insts.clone().iter() {
-        if stack_addr_load_insts_users.get(&inst).map(|users| users.is_empty()).unwrap_or(true) {
-            println!("Removing unused stack_load {}", inst);
-            func.dfg.detach_results(inst);
-            func.dfg.replace(inst).nop();
-            stack_load_insts.remove(&inst);
+        let mut is_recorded_stack_addr_or_stack_load = false;
+        for stack_slot_users in stack_slot_usage_map.values() {
+            is_recorded_stack_addr_or_stack_load |= stack_slot_users.stack_addr.contains(inst) || stack_slot_users.stack_load.contains(inst);
         }
+        assert!(is_recorded_stack_addr_or_stack_load);
     }
 
-
-    //println!("stack_addr (after): [{}]", bool_secondary_map_to_string(&stack_addr_insts));
-
-    let mut stack_slot_usage_map: BTreeMap<OrdStackSlot, HashSet<Inst>> = BTreeMap::new();
-    for &inst in stack_load_insts.iter() {
-        match func.dfg[inst] {
-            InstructionData::StackLoad {
-                opcode: Opcode::StackLoad,
-                stack_slot,
-                offset: _,
-            } => {
-                stack_slot_usage_map.entry(OrdStackSlot(stack_slot)).or_insert_with(HashSet::new).insert(inst);
-            }
-            ref data => unreachable!("{:?}", data),
-        }
-    }
-    for &inst in stack_store_insts.iter() {
-        match func.dfg[inst] {
-            InstructionData::StackStore {
-                opcode: Opcode::StackStore,
-                arg: _,
-                stack_slot,
-                offset: _,
-            } => {
-                stack_slot_usage_map.entry(OrdStackSlot(stack_slot)).or_insert_with(HashSet::new).insert(inst);
+    // Replace all unused stack_addr and stack_load instructions with nop.
+    for stack_slot_users in stack_slot_usage_map.values_mut() {
+        // FIXME remove clone
+        for &inst in stack_slot_users.stack_addr.clone().iter() {
+            if stack_addr_load_insts_users.get(&inst).map(|users| users.is_empty()).unwrap_or(true) {
+                println!("Removing unused stack_addr {}", inst);
+                func.dfg.detach_results(inst);
+                func.dfg.replace(inst).nop();
+                stack_slot_users.stack_addr.remove(&inst);
             }
-            ref data => unreachable!("{:?}", data),
         }
-    }
-    for &inst in stack_addr_insts.iter() {
-        match func.dfg[inst] {
-            InstructionData::StackLoad {
-                opcode: Opcode::StackAddr,
-                stack_slot,
-                offset: _,
-            } => {
-                stack_slot_usage_map.entry(OrdStackSlot(stack_slot)).or_insert_with(HashSet::new).insert(inst);
+
+        for &inst in stack_slot_users.stack_load.clone().iter() {
+            if stack_addr_load_insts_users.get(&inst).map(|users| users.is_empty()).unwrap_or(true) {
+                println!("Removing unused stack_addr {}", inst);
+                func.dfg.detach_results(inst);
+                func.dfg.replace(inst).nop();
+                stack_slot_users.stack_load.remove(&inst);
             }
-            ref data => unreachable!("{:?}", data),
         }
     }
 
-    println!("stack slot usage: {:?}", stack_slot_usage_map);
+    println!("stack slot usage (after): {:?}", stack_slot_usage_map);
 
     for (stack_slot, users) in stack_slot_usage_map.iter_mut() {
-        let mut is_addr_leaked = false;
-        let mut is_loaded = false;
-        let mut is_stored = false;
-        for &user in users.iter() {
-            match func.dfg[user] {
-                InstructionData::StackLoad {
-                    opcode: Opcode::StackAddr,
-                    stack_slot,
-                    offset: _,
-                } => {
-                    is_addr_leaked = true;
-                }
-                InstructionData::StackLoad {
-                    opcode: Opcode::StackLoad,
-                    stack_slot,
-                    offset: _,
-                } => {
-                    is_loaded = true;
-                }
-                InstructionData::StackStore {
-                    opcode: Opcode::StackStore,
-                    arg: _,
-                    stack_slot,
-                    offset: _,
-                } => {
-                    is_stored = true;
-                }
-                ref data => unreachable!("{:?}", data),
-            }
-        }
-
-        if is_addr_leaked || (is_loaded && is_stored) {
+        if users.stack_addr.is_empty().not() || (users.stack_load.is_empty().not() && users.stack_store.is_empty().not()) {
             continue;
         }
 
-        if is_loaded {
+        if users.stack_load.is_empty().not() {
             println!("[{}] [BUG?] Reading uninitialized memory", name);
         } else {
             // Stored value never read; just remove reads.
-            for &user in users.iter() {
+            for user in users.stack_store.drain() {
                 println!("[{}] Remove dead stack store {} of {}", name, user, stack_slot.0);
                 func.dfg.replace(user).nop();
             }