IP SROA Parameters

• Break up products in parameters to functions, not just return values.
• TODO: Fix CFD test to not use products in Rust interface.

hercules_opt/src/interprocedural_sroa.rs

-use std::collections::HashMap;
-use std::iter::zip;
-
 use hercules_ir::ir::*;
 
 use crate::*;
@@ -39,15 +36,30 @@ pub fn interprocedural_sroa(
         }
 
         let editor: &mut FunctionEditor = &mut editors[func_id.idx()];
+        let param_types = &editor.func().param_types.to_vec();
         let return_types = &editor.func().return_types.to_vec();
 
-        // We determine the new return types of the function and track a map
-        // that tells us how the old return values are constructed from the
-        // new ones
+        // We determine the new param/return types of the function and track a
+        // map that tells us how the old param/return values are constructed
+        // from the new ones.
+        let mut new_param_types = vec![];
+        let mut old_param_type_map = vec![];
         let mut new_return_types = vec![];
         let mut old_return_type_map = vec![];
         let mut changed = false;
 
+        for par_typ in param_types.iter() {
+            if !can_sroa_type(editor, *par_typ) {
+                old_param_type_map.push(IndexTree::Leaf(new_param_types.len()));
+                new_param_types.push(*par_typ);
+            } else {
+                let (types, index) = sroa_type(editor, *par_typ, new_param_types.len());
+                old_param_type_map.push(index);
+                new_param_types.extend(types);
+                changed = true;
+            }
+        }
+
         for ret_typ in return_types.iter() {
             if !can_sroa_type(editor, *ret_typ) {
                 old_return_type_map.push(IndexTree::Leaf(new_return_types.len()));
@@ -60,25 +72,55 @@ pub fn interprocedural_sroa(
             }
         }
 
-        // If the return type is not changed by IP SROA, skip to the next function
+        // If the param/return types aren't changed by IP SROA, skip to the next
+        // function.
         if !changed {
             continue;
         }
 
-        // Now, modify each return in the current function and the return type
-        let return_nodes = editor
-            .func()
-            .nodes
-            .iter()
-            .enumerate()
-            .filter_map(|(idx, node)| {
-                if node.try_return().is_some() {
-                    Some(NodeID::new(idx))
+        // Modify each parameter in the current function and the param types.
+        let mut param_nodes: Vec<_> = vec![vec![]; param_types.len()];
+        for id in editor.node_ids() {
+            if let Some(idx) = editor.func().nodes[id.idx()].try_parameter() {
+                param_nodes[idx].push(id);
+            }
+        }
+        println!("{}", editor.func().name);
+        let success = editor.edit(|mut edit| {
+            for (idx, ids) in param_nodes.into_iter().enumerate() {
+                let new_indices = &old_param_type_map[idx];
+                let built = if let IndexTree::Leaf(new_idx) = new_indices {
+                    edit.add_node(Node::Parameter { index: *new_idx })
                 } else {
-                    None
+                    let prod_ty = param_types[idx];
+                    let cons = edit.add_zero_constant(prod_ty);
+                    let mut cons = edit.add_node(Node::Constant { id: cons });
+                    new_indices.for_each(|idx: &Vec<Index>, param_idx: &usize| {
+                        let param = edit.add_node(Node::Parameter { index: *param_idx });
+                        cons = edit.add_node(Node::Write {
+                            collect: cons,
+                            data: param,
+                            indices: idx.clone().into_boxed_slice(),
+                        });
+                    });
+                    cons
+                };
+                for id in ids {
+                    edit = edit.replace_all_uses(id, built)?;
+                    edit = edit.delete_node(id)?;
                 }
-            })
-            .collect::<Vec<_>>();
+            }
+
+            edit.set_param_types(new_param_types);
+            Ok(edit)
+        });
+        assert!(success, "IP SROA expects to be able to edit everything, specify what functions to IP SROA via the func_selection argument");
+
+        // Modify each return in the current function and the return types.
+        let return_nodes: Vec<_> = editor
+            .node_ids()
+            .filter(|id| editor.func().nodes[id.idx()].is_return())
+            .collect();
         let success = editor.edit(|mut edit| {
             for node in return_nodes {
                 let Node::Return { control, data } = edit.get_node(node) else {
@@ -114,17 +156,15 @@ pub fn interprocedural_sroa(
             }
 
             edit.set_return_types(new_return_types);
-
             Ok(edit)
         });
         assert!(success, "IP SROA expects to be able to edit everything, specify what functions to IP SROA via the func_selection argument");
 
-        // Finally, update calls of this function
-        // In particular, we actually don't have to update the call node at all but have to update
-        // its DataProjection users
+        // Finally, update calls of this function.
         for (caller, callsite) in callsites {
             let editor = &mut editors[caller.idx()];
             assert!(editor.func_id() == caller);
+
             let projs = editor.get_users(callsite).collect::<Vec<_>>();
             for proj_id in projs {
                 let Node::DataProjection { data: _, selection } = editor.node(proj_id) else {
@@ -134,6 +174,40 @@ pub fn interprocedural_sroa(
                 let typ = types[caller.idx()][proj_id.idx()];
                 replace_returned_value(editor, proj_id, typ, new_return_info, callsite);
             }
+
+            let (control, callee, dc_args, args) =
+                editor.func().nodes[callsite.idx()].try_call().unwrap();
+            let dc_args = dc_args.clone();
+            let args = args.clone();
+            let success = editor.edit(|mut edit| {
+                let mut new_args = vec![];
+                for (idx, (data_id, update_info)) in
+                    args.iter().zip(old_param_type_map.iter()).enumerate()
+                {
+                    if let IndexTree::Leaf(new_idx) = update_info {
+                        // Unchanged parameter value
+                        assert!(new_args.len() == *new_idx);
+                        new_args.push(*data_id);
+                    } else {
+                        // SROA'd parameter value
+                        let reads = generate_reads_edit(&mut edit, param_types[idx], *data_id);
+                        reads.zip(update_info).for_each(|_, (read_id, ret_idx)| {
+                            assert!(new_args.len() == **ret_idx);
+                            new_args.push(*read_id);
+                        });
+                    }
+                }
+                let new_call = edit.add_node(Node::Call {
+                    control,
+                    function: callee,
+                    dynamic_constants: dc_args,
+                    args: new_args.into_boxed_slice(),
+                });
+                edit = edit.replace_all_uses(callsite, new_call)?;
+                edit = edit.delete_node(callsite)?;
+                Ok(edit)
+            });
+            assert!(success);
         }
     }
 }