Address feedback

external-crates/move/crates/move-compiler/src/cfgir/mod.rs

@@ -22,7 +22,6 @@ use crate::{
 use cfg::*;
 use move_ir_types::location::Loc;
 use move_symbol_pool::Symbol;
-use optimize::optimize;
 use std::collections::BTreeSet;
 
 pub struct CFGContext<'a> {

external-crates/move/crates/move-compiler/src/cfgir/optimize/coalesce_locals.rs

@@ -17,8 +17,8 @@ use heuristic_graph_coloring as coloring;
 use move_proc_macros::growing_stack;
 use move_symbol_pool::Symbol;
 
-// DEBUG flag for general printing. Defined as a const so the code will be eliminated by the rustc
-// optimizer when debugging is not required.
+/// DEBUG flag for general printing. Defined as a const so the code will be eliminated by the rustc
+/// optimizer when debugging is not required.
 const DEBUG_COALESCE: bool = false;
 
 //**************************************************************************************************
@@ -31,7 +31,7 @@ pub fn optimize(
     infinite_loop_starts: &BTreeSet<Label>,
     locals: UniqueMap<Var, (Mutability, SingleType)>,
     cfg: &mut MutForwardCFG,
-) -> UniqueMap<Var, (Mutability, SingleType)> {
+) -> Option<UniqueMap<Var, (Mutability, SingleType)>> {
     macro_rules! unique_add_or_error {
         ($step:expr, $map:expr, $var:expr, $mut:expr, $ty:expr) => {{
             if let Err(_) = $map.add($var, ($mut, $ty)) {
@@ -40,7 +40,7 @@ pub fn optimize(
                     $var.value(),
                     $step
                 );
-                return locals;
+                return None;
             }
         }};
     }
@@ -59,7 +59,7 @@ pub fn optimize(
                 conflicts::Conflicts::Set(vars) => {
                     let var_string = vars
                         .iter()
-                        .map(|x| format!("{}", x.value()))
+                        .map(|x| format!("{}", x))
                         .collect::<Vec<_>>()
                         .join(", ");
                     println!("{{{}}}", var_string);
@@ -71,38 +71,20 @@ pub fn optimize(
 
     let mut new_locals = UniqueMap::new();
 
-    // Collect everything that is going to keep its name. This currently includes:
-    // - All parameters.
-    // - Anything that is borrowed.
-    // - Anything that is not used, but whose type does not have `drop`.
-    let param_set = signature
-        .parameters
-        .iter()
-        .map(|(_, var, _)| var)
-        .collect::<BTreeSet<_>>();
-    let all_set = locals
-        .key_cloned_iter()
-        .filter_map(|(var, (_mut, ty))| {
-            if param_set.contains(&var)
-                || (!(conflict_graph.used(&var)
-                    || ty.value.abilities(ty.loc).has_ability_(Ability_::Drop)))
-                || matches!(
-                    conflict_graph.get_conflicts(&var),
-                    Some(conflicts::Conflicts::All)
-                )
-            {
-                Some(var)
-            } else {
-                None
-            }
-        })
-        .collect::<BTreeSet<_>>();
+    // Set of things that cannot be coalesced.
+    let uncoalesced_set = uncoalescable_vars(&conflict_graph, &signature.parameters, &locals);
+    if DEBUG_COALESCE {
+        println!(
+            "\n\nuncoalescable: {}",
+            format_oxford_list!("and", "{}", &uncoalesced_set)
+        );
+    }
 
     let mut locals_by_type = BTreeMap::new();
     for (var, (mut_, ty)) in locals.clone().into_iter() {
         // Add the name-retaining variables to the rename map. This also inlcudes all the
         // parameters we've seen.
-        if all_set.contains(&var) {
+        if uncoalesced_set.contains(&var) {
             unique_add_or_error!("all_conflicting", new_locals, var, mut_, ty.clone());
             continue;
         } else {
@@ -115,7 +97,7 @@ pub fn optimize(
 
     // If there are no locals to coalesce, return.
     if locals_by_type.is_empty() {
-        return locals;
+        return None;
     };
 
     let mut rename_map = BTreeMap::new();
@@ -127,19 +109,20 @@ pub fn optimize(
             .enumerate()
             .map(|(n, v)| (*v, n))
             .collect::<BTreeMap<_, _>>();
-        for var in local_set {
-            match conflict_graph.get_conflicts(var) {
-                None => {
-                    continue;
-                }
-                Some(conflicts::Conflicts::Set(vars)) => {
+        for (var, conflict) in local_set.iter().filter_map(|var| {
+            conflict_graph
+                .get_conflicts(var)
+                .map(|conflict| (var, conflict))
+        }) {
+            match conflict {
+                conflicts::Conflicts::Set(vars) => {
                     let index = var_index_map.get(var).unwrap();
                     let relevant_vars = vars.intersection(local_set).collect::<BTreeSet<_>>();
                     for var in relevant_vars {
                         graph.add_edge(*index, *var_index_map.get(var).unwrap());
                     }
                 }
-                Some(conflicts::Conflicts::All) => unreachable!(),
+                conflicts::Conflicts::All => unreachable!(),
             }
         }
 
@@ -208,31 +191,69 @@ pub fn optimize(
         print!("New locals: ");
         println!(
             "{}",
-            format_oxford_list!(
-                "and",
-                "{}",
-                new_locals
-                    .key_cloned_iter()
-                    .map(|(key, _)| key.value())
-                    .collect::<Vec<_>>()
-            )
+            format_oxford_list!("and", "{}", new_locals.key_cloned().collect::<Vec<_>>())
         );
     }
 
     if rename_map.is_empty() {
         if DEBUG_COALESCE {
             println!("-- false ---------------------------");
         }
-        locals
+        None
     } else {
         coalesce(cfg, &rename_map);
         if DEBUG_COALESCE {
             println!("-- done (new locals: {:000})--------", new_locals.len());
         }
-        new_locals
+        Some(new_locals)
     }
 }
 
+/// Collect everything that _must_ keep its name, and can't be coalesced. This includes:
+/// - All parameters, as they are not locals.
+/// - Any local that was marked as Conflits::All in the conflict graph, indicating it is
+///   borrowed at some point in the function. We don't have enough information (here) to
+///   determine if coalescing is valid, so we simply do not.
+/// - Any local that is not used, but whose type does not have `drop`. These are things that
+///   should have been dropped by the function, but were not for some reason. In some cases this
+///   is acceptable (if the block ends in an abort), but coalescing will result in attempting to
+///   store to a local with a non-droppable value, which is an error.
+fn uncoalescable_vars(
+    conflict_graph: &conflicts::Graph,
+    parameters: &[(Mutability, Var, SingleType)],
+    locals: &UniqueMap<Var, (Mutability, SingleType)>,
+) -> BTreeSet<Var> {
+    let param_set = parameters
+        .iter()
+        .map(|(_, var, _)| var)
+        .collect::<BTreeSet<_>>();
+
+    let is_param = |var: &Var| param_set.contains(var);
+    let is_borrowed = |var: &Var| {
+        matches!(
+            conflict_graph.get_conflicts(var),
+            Some(conflicts::Conflicts::All)
+        )
+    };
+    let is_unused_without_drop = |var: &Var, ty: &SingleType| {
+        let result =
+            !conflict_graph.used(var) && !ty.value.abilities(ty.loc).has_ability_(Ability_::Drop);
+        if DEBUG_COALESCE && result {
+            println!("{} is unused without drop", var);
+        }
+        result
+    };
+
+    // Set of things that cannot be coalesced.
+    locals
+        .key_cloned_iter()
+        .filter(|(var, (_mut, ty))| {
+            is_param(var) || is_borrowed(var) || is_unused_without_drop(var, ty)
+        })
+        .map(|(var, _)| var)
+        .collect::<BTreeSet<_>>()
+}
+
 //**************************************************************************************************
 // Conflict Graph
 //**************************************************************************************************
@@ -284,7 +305,7 @@ mod conflicts {
 
             for (lbl, commands) in cfg.blocks() {
                 let per_command_states = per_command_states.get(lbl).unwrap();
-                assert!(commands.len() == per_command_states.len());
+                assert_eq!(commands.len(), per_command_states.len());
                 for (cmd, lives) in commands.iter().zip(per_command_states) {
                     command(&mut graph, &lives.live_set, cmd);
                 }

external-crates/move/crates/move-compiler/src/cfgir/optimize/mod.rs

@@ -20,6 +20,8 @@ use crate::{
     shared::{unique_map::UniqueMap, CompilationEnv},
 };
 
+use std::collections::BTreeSet;
+
 pub type Optimization = fn(
     &FunctionSignature,
     &UniqueMap<Var, (Mutability, SingleType)>,
@@ -43,16 +45,58 @@ const MOVE_2024_OPTIMIZATIONS: &[Optimization] = &[
 ];
 
 #[growing_stack]
-pub fn optimize(
+/// Attempts to optimize a constant into a value.
+pub fn optimize_constant(
+    env: &mut CompilationEnv,
+    package: Option<Symbol>,
+    signature: &FunctionSignature,
+    locals: &mut UniqueMap<Var, (Mutability, SingleType)>,
+    constants: &UniqueMap<ConstantName, Value>,
+    cfg: &mut MutForwardCFG,
+) {
+    optimization_loop(env, package, signature, locals, constants, cfg);
+}
+
+#[growing_stack]
+/// Attempts to optimize a function body, including the optimizations that constants receive plus
+/// variable coalescing.
+pub fn optimize_function(
+    env: &mut CompilationEnv,
+    package: Option<Symbol>,
+    signature: &FunctionSignature,
+    infinite_loop_starts: &BTreeSet<Label>,
+    locals: &mut UniqueMap<Var, (Mutability, SingleType)>,
+    constants: &UniqueMap<ConstantName, Value>,
+    cfg: &mut MutForwardCFG,
+) {
+    optimization_loop(env, package, signature, locals, constants, cfg);
+
+    if env.supports_feature(package, FeatureGate::Move2024Optimizations) {
+        if let Some(new_locals) =
+            coalesce_locals::optimize(signature, infinite_loop_starts, locals.clone(), cfg)
+        {
+            *locals = new_locals;
+            cfg.recompute();
+            optimization_loop(env, package, signature, locals, constants, cfg);
+        }
+    }
+}
+
+/// Main loop of optimizations
+#[growing_stack]
+fn optimization_loop(
     env: &mut CompilationEnv,
     package: Option<Symbol>,
     signature: &FunctionSignature,
-    locals: &UniqueMap<Var, (Mutability, SingleType)>,
+    locals: &mut UniqueMap<Var, (Mutability, SingleType)>,
     constants: &UniqueMap<ConstantName, Value>,
     cfg: &mut MutForwardCFG,
 ) {
     let mut count = 0;
-    let optimizations = if env.supports_feature(package, FeatureGate::Move2024Optimizations) {
+
+    let optimize_2024 = env.supports_feature(package, FeatureGate::Move2024Optimizations);
+
+    let optimizations = if optimize_2024 {
         MOVE_2024_OPTIMIZATIONS
     } else {
         OPTIMIZATIONS

external-crates/move/crates/move-compiler/src/cfgir/translate.rs

@@ -11,7 +11,6 @@ use crate::{
     },
     diag,
     diagnostics::Diagnostics,
-    editions::FeatureGate,
     expansion::ast::{Attributes, ModuleIdent, Mutability},
     hlir::ast::{self as H, BlockLabel, Label, Value, Value_, Var},
     parser::ast::{ConstantName, FunctionName},
@@ -462,7 +461,7 @@ fn constant_(
     full_loc: Loc,
     attributes: &Attributes,
     signature: H::BaseType,
-    locals: UniqueMap<Var, (Mutability, H::SingleType)>,
+    mut locals: UniqueMap<Var, (Mutability, H::SingleType)>,
     body: H::Block,
 ) -> Option<H::Exp> {
     use H::Command_ as C;
@@ -501,11 +500,11 @@ fn constant_(
         "{}",
         ICE_MSG
     );
-    cfgir::optimize(
+    cfgir::optimize::optimize_constant(
         context.env,
         context.current_package,
         &fake_signature,
-        &locals,
+        &mut locals,
         constant_values,
         &mut cfg,
     );
@@ -665,34 +664,15 @@ fn function_body(
             cfgir::refine_inference_and_verify(context.env, &function_context, &mut cfg);
             // do not optimize if there are errors, warnings are okay
             if !context.env.has_errors() {
-                cfgir::optimize(
+                cfgir::optimize::optimize_function(
                     context.env,
                     context.current_package,
                     signature,
-                    &locals,
+                    &infinite_loop_starts,
+                    &mut locals,
                     &UniqueMap::new(),
                     &mut cfg,
                 );
-                if context
-                    .env
-                    .supports_feature(context.current_package, FeatureGate::Move2024Optimizations)
-                {
-                    locals = cfgir::optimize::coalesce_locals::optimize(
-                        signature,
-                        &infinite_loop_starts,
-                        locals,
-                        &mut cfg,
-                    );
-                    cfg.recompute();
-                    cfgir::optimize(
-                        context.env,
-                        context.current_package,
-                        signature,
-                        &locals,
-                        &UniqueMap::new(),
-                        &mut cfg,
-                    );
-                }
             }
 
             let block_info = block_info