rewrite forkify as single edit per loop

hercules_opt/src/forkify.rs

@@ -298,29 +298,11 @@ pub fn forkify_loop(
     let (_, factors) = function.nodes[fork_id.idx()].try_fork().unwrap();
     let dimension = factors.len() - 1;
 
-    // Create ThreadID
-    editor.edit(|mut edit| {
-        let thread_id = Node::ThreadID {
-            control: fork_id,
-            dimension: dimension,
-        };
-        let thread_id_id = edit.add_node(thread_id);
-
-        // Replace uses that are inside with the thread id
-        edit = edit.replace_all_uses_where(canonical_iv.phi(), thread_id_id, |node| {
-            loop_nodes.contains(node)
-        })?;
+    // Start failable edit:
 
-        // Replace uses that are outside with DC - 1. Or just give up.
-        let bound_dc_node = edit.add_node(Node::DynamicConstant { id: bound_dc_id });
-        edit = edit.replace_all_uses_where(canonical_iv.phi(), bound_dc_node, |node| {
-            !loop_nodes.contains(node)
-        })?;
-
-        edit.delete_node(canonical_iv.phi())
-    });
-
-    for reduction_phi in reductionable_phis {
+    let redcutionable_phis_and_init: Vec<(_, NodeID)> = 
+        reductionable_phis.iter().map(|reduction_phi| {
+        
         let LoopPHI::Reductionable {
             phi,
             data_cycle: _,
@@ -342,12 +324,41 @@ pub fn forkify_loop(
         .unwrap()
         .1;
 
-        editor.edit(|mut edit| {
+        (reduction_phi, init)
+    }).collect();
+
+    editor.edit(|mut edit| {
+        let thread_id = Node::ThreadID {
+            control: fork_id,
+            dimension: dimension,
+        };
+        let thread_id_id = edit.add_node(thread_id);
+
+        // Replace uses that are inside with the thread id
+        edit = edit.replace_all_uses_where(canonical_iv.phi(), thread_id_id, |node| {
+            loop_nodes.contains(node)
+        })?;
+
+        edit = edit.delete_node(canonical_iv.phi())?;
+
+        for (reduction_phi, init) in redcutionable_phis_and_init {
+            let LoopPHI::Reductionable {
+                phi,
+                data_cycle: _,
+                continue_latch,
+                is_associative: _,
+            } = *reduction_phi
+            else {
+                panic!();
+            };
+            
             let reduce = Node::Reduce {
                 control: join_id,
                 init,
                 reduct: continue_latch,
             };
+            
+            
             let reduce_id = edit.add_node(reduce);
 
             if (!edit.get_node(init).is_reduce()
@@ -375,20 +386,14 @@ pub fn forkify_loop(
             edit = edit.replace_all_uses_where(continue_latch, reduce_id, |usee| {
                 !loop_nodes.contains(usee) && *usee != reduce_id
             })?;
-            edit.delete_node(phi)
-        });
-    }
-
-    // Replace all uses of the loop header with the fork
-    editor.edit(|edit| edit.replace_all_uses(l.header, fork_id));
+            edit = edit.delete_node(phi)?
 
-    editor.edit(|edit| edit.replace_all_uses(loop_continue_projection, fork_id));
+        }
 
-    editor.edit(|edit| edit.replace_all_uses(loop_exit_projection, join_id));
+        edit = edit.replace_all_uses(l.header, fork_id)?;
+        edit = edit.replace_all_uses(loop_continue_projection, fork_id)?;
+        edit = edit.replace_all_uses(loop_exit_projection, join_id)?;
 
-    // Get rid of loop condition
-    // DCE should get these, but delete them ourselves because we are nice :)
-    editor.edit(|mut edit| {
         edit = edit.delete_node(loop_continue_projection)?;
         edit = edit.delete_node(condition_node)?; // Might have to get rid of other users of this.
         edit = edit.delete_node(loop_exit_projection)?;
@@ -396,7 +401,7 @@ pub fn forkify_loop(
         edit = edit.delete_node(l.header)?;
         Ok(edit)
     });
-
+    
     return true;
 }