Fix inlining and editor handling of gravestone nodes

hercules_opt/src/editor.rs

@@ -634,7 +634,17 @@ pub fn repair_plan(plan: &mut Plan, new_function: &Function, edits: &[Edit]) {
         new_partitions[data_id.idx()] = new_partitions[bbs[data_id.idx()].idx()];
     }
 
-    // Step 6: wrap everything up.
+    // Step 6: create a solitary gravestone partition. This will get removed
+    // when gravestone nodes are removed.
+    let gravestone_partition = PartitionID::new(plan.num_partitions);
+    plan.num_partitions += 1;
+    for (idx, node) in new_function.nodes.iter().enumerate() {
+        if idx > 0 && node.is_start() {
+            new_partitions[idx] = Some(gravestone_partition);
+        }
+    }
+
+    // Step 7: wrap everything up.
     plan.partitions = new_partitions.into_iter().map(|id| id.unwrap()).collect();
     plan.partition_devices
         .resize(plan.num_partitions, Device::CPU);

hercules_opt/src/inline.rs

@@ -197,7 +197,14 @@ fn inline_func(
 
             // Stich the control use of the original call node's region with
             // the predecessor control of the inlined function's return.
-            edit = edit.replace_all_uses(control, old_id_to_new_id(called_return_pred))?;
+            edit = edit.replace_all_uses(
+                control,
+                if called_return_pred == NodeID::new(0) {
+                    call_pred
+                } else {
+                    old_id_to_new_id(called_return_pred)
+                },
+            )?;
 
             // Stitch uses of parameter nodes in the inlined function to the IDs
             // of arguments provided to the call node.