Use ID types in manifest

hercules_cg/src/common.rs

@@ -258,6 +258,27 @@ impl<'a> FunctionContext<'a> {
         res.sort();
         res
     }
+
+    /*
+     * Determine the array numbers for all the array constants. These are needed
+     * to know which pointer passed to the runtime corresponds to which array
+     * constant. Return a map from constant ID to array number - non-array
+     * constants don't have an array number.
+     */
+    pub(crate) fn array_constant_inputs(&self) -> Vec<Option<u32>> {
+        self.constants
+            .iter()
+            .scan(0, |num, cons| {
+                if cons.try_array_type(self.types).is_some() {
+                    let res = Some(*num);
+                    *num += 1;
+                    Some(res)
+                } else {
+                    Some(None)
+                }
+            })
+            .collect()
+    }
 }
 
 /*
@@ -327,14 +348,12 @@ impl<'a> PartitionContext<'a> {
 
         // Assemble the manifest.
         let mut manifest = PartitionManifest::default();
-        manifest.top_node = top_node.idx() as u32;
+        manifest.top_node = top_node;
 
         // The first inputs are the data inputs, from other partitions.
-        manifest.inputs.extend(
-            data_inputs
-                .iter()
-                .map(|x| PartitionInput::DataInput(x.idx() as u32)),
-        );
+        manifest
+            .inputs
+            .extend(data_inputs.iter().map(|x| PartitionInput::DataInput(*x)));
 
         // The next inputs are the function parameters, all in order.
         manifest.inputs.extend(
@@ -359,11 +378,9 @@ impl<'a> PartitionContext<'a> {
         );
 
         // The outputs are the data outputs of this partition.
-        manifest.outputs.extend(
-            data_outputs
-                .iter()
-                .map(|x| PartitionOutput::DataOutput(x.idx() as u32)),
-        );
+        manifest
+            .outputs
+            .extend(data_outputs.iter().map(|x| PartitionOutput::DataOutput(*x)));
 
         // If there are multiple control returns, also output the node being
         // returned from.
@@ -374,7 +391,7 @@ impl<'a> PartitionContext<'a> {
         // Store the successor partitions.
         manifest.successor_partitions = control_successors
             .into_iter()
-            .map(|(part_id, control_id)| (control_id.idx() as u32, part_id.idx() as u32))
+            .map(|(part_id, control_id)| (control_id, part_id))
             .collect();
 
         PartitionContext {

hercules_cg/src/top.rs

@@ -67,6 +67,7 @@ pub fn codegen<W: Write>(
         types: module.types.clone(),
         // TODO: populate array constants.
         array_constants: vec![],
+        array_cons_ids: vec![],
     })
 }
 
@@ -112,11 +113,11 @@ impl<'a> FunctionContext<'a> {
             .function
             .nodes
             .iter()
-            .filter_map(|node| node.try_return().map(|(_, data)| data.idx() as u32))
-            .collect::<Vec<u32>>();
+            .filter_map(|node| node.try_return().map(|(_, data)| data))
+            .collect();
 
         // Get the partition ID of the start node.
-        let top_partition = self.plan.partitions[0].idx() as u32;
+        let top_partition = self.plan.partitions[0];
 
         // Generate code for each individual partition. This generates a single
         // LLVM function per partition. These functions will be called in async

hercules_ir/src/ir.rs

@@ -1189,6 +1189,7 @@ macro_rules! define_id_type {
     ($x: ident) => {
         #[derive(
             Debug,
+            Default,
             Clone,
             Copy,
             PartialEq,

hercules_ir/src/manifest.rs

@@ -23,6 +23,8 @@ pub struct ModuleManifest {
     // the dimensions of arrays at this point, since the runtime doesn't
     // manipulate or otherwise need the dimensions of constant arrays.
     pub array_constants: Vec<Vec<u8>>,
+    // Store map from array number, in array_constants, to constant ID.
+    pub array_cons_ids: Vec<ConstantID>,
 }
 
 #[derive(Debug, Serialize, Deserialize, Hash)]
@@ -47,10 +49,9 @@ pub struct FunctionManifest {
     // their dynamic constant parameters.
     pub dynamic_constant_rules: Vec<DynamicConstantRule>,
     // The partition containing the start node might not be partition 0.
-    pub top_partition: u32,
+    pub top_partition: PartitionID,
     // Keep track of which node values are returned from the overall function.
-    // Integer is the node ID.
-    pub returned_values: Vec<u32>,
+    pub returned_values: Vec<NodeID>,
 }
 
 /*
@@ -67,9 +68,8 @@ pub enum DynamicConstantRule {
 
 #[derive(Debug, Serialize, Deserialize, Hash, PartialEq, Eq)]
 pub enum PartitionInput {
-    // Data input from another partition within this function. Integer is the
-    // node ID used from the other partition.
-    DataInput(u32),
+    // Data input from another partition within this function.
+    DataInput(NodeID),
     // An argument from the function parameters. Integer is the parameter index.
     FunctionArgument(u32),
     // An array constant used in this function. Integer is the array constant
@@ -83,31 +83,30 @@ pub enum PartitionInput {
 #[derive(Debug, Serialize, Deserialize, Hash, PartialEq, Eq)]
 pub enum PartitionOutput {
     // Data output used by another partition within this function, or to be
-    // returned from this function. Integer is the node ID used in the other
-    // partition or by a return node.
-    DataOutput(u32),
+    // returned from this function.
+    DataOutput(NodeID),
     // Value indicating control flow that the runtime should take.
     ControlIndicator,
 }
 
 #[derive(Debug, Serialize, Deserialize, Default, Hash)]
 pub struct PartitionManifest {
-    // Top node for this partition, as an integer.
-    pub top_node: u32,
+    // Top node for this partition.
+    pub top_node: NodeID,
     pub inputs: Vec<PartitionInput>,
     pub outputs: Vec<PartitionOutput>,
 
     // Store the partitions that come after this partition. The first element of
     // the pair is the control returning node in the current partition that
     // corresponds to a particular successor, and the second element of the pair
-    // is the partition index of the corresponding successor partition.
-    pub successor_partitions: Vec<(u32, u32)>,
+    // is the partition ID of the corresponding successor partition.
+    pub successor_partitions: Vec<(NodeID, PartitionID)>,
 }
 
 /*
  * Get the data inputs of partitions in a function.
  */
-pub fn function_partition_inputs(manifest: &FunctionManifest) -> Vec<u32> {
+pub fn function_partition_inputs(manifest: &FunctionManifest) -> Vec<NodeID> {
     let mut all_data_inputs = manifest
         .partitions
         .iter()
@@ -121,7 +120,7 @@ pub fn function_partition_inputs(manifest: &FunctionManifest) -> Vec<u32> {
             })
         })
         .flatten()
-        .collect::<Vec<u32>>();
+        .collect::<Vec<NodeID>>();
     all_data_inputs.sort();
     all_data_inputs.dedup();
     all_data_inputs
@@ -130,7 +129,7 @@ pub fn function_partition_inputs(manifest: &FunctionManifest) -> Vec<u32> {
 /*
  * Get the data outputs of partitions in a function.
  */
-pub fn function_partition_outputs(manifest: &FunctionManifest) -> Vec<u32> {
+pub fn function_partition_outputs(manifest: &FunctionManifest) -> Vec<NodeID> {
     let mut all_data_outputs = manifest
         .partitions
         .iter()
@@ -144,7 +143,7 @@ pub fn function_partition_outputs(manifest: &FunctionManifest) -> Vec<u32> {
             })
         })
         .flatten()
-        .collect::<Vec<u32>>();
+        .collect::<Vec<NodeID>>();
     all_data_outputs.sort();
     all_data_outputs.dedup();
     all_data_outputs
@@ -166,7 +165,7 @@ pub fn verify_manifest(manifest: &ModuleManifest) -> Result<(), String> {
             // Data outputs include values returned from the overall function.
             // These are obviously not used by another partition, but that's
             // fine. Remove them here explicitly.
-            let all_data_outputs = all_data_outputs.into_iter().filter(|idx| !manifest.returned_values.contains(idx)).collect::<Vec<u32>>();
+            let all_data_outputs = all_data_outputs.into_iter().filter(|idx| !manifest.returned_values.contains(idx)).collect::<Vec<NodeID>>();
             if all_data_inputs != all_data_outputs {
                 return Err(format!("PANIC: Partitions in manifest contain inconsistent data inputs and data outputs of partitions.\nHere are the data input IDs:\n{:?}\nHere are the data output IDs:\n{:?}\n", all_data_inputs, all_data_outputs));
             }

hercules_rt_proc/src/lib.rs

@@ -118,7 +118,7 @@ fn codegen(manifest: &ModuleManifest, elf: &[u8]) -> Result<String, anyhow::Erro
                     .unwrap(),
             );
             let input_types = partition.inputs.iter().map(|input| match input {
-                PartitionInput::DataInput(node_idx) => function.typing[*node_idx as usize],
+                PartitionInput::DataInput(node_id) => function.typing[node_id.idx()],
                 PartitionInput::FunctionArgument(param_idx) => {
                     function.param_types[*param_idx as usize]
                 }
@@ -130,9 +130,7 @@ fn codegen(manifest: &ModuleManifest, elf: &[u8]) -> Result<String, anyhow::Erro
                     .outputs
                     .iter()
                     .map(|input| match input {
-                        PartitionOutput::DataOutput(node_idx) => {
-                            function.typing[*node_idx as usize]
-                        }
+                        PartitionOutput::DataOutput(node_id) => function.typing[node_id.idx()],
                         PartitionOutput::ControlIndicator => u64_ty_id,
                     })
                     .collect(),
@@ -166,8 +164,8 @@ fn codegen(manifest: &ModuleManifest, elf: &[u8]) -> Result<String, anyhow::Erro
             write!(
                 rust_code,
                 "let mut node_{}: Option<{}> = None;",
-                node,
-                &rust_types[function.typing[node as usize].idx()]
+                node.idx(),
+                &rust_types[function.typing[node.idx()].idx()]
             )?;
         }
 
@@ -177,7 +175,7 @@ fn codegen(manifest: &ModuleManifest, elf: &[u8]) -> Result<String, anyhow::Erro
         write!(
             rust_code,
             "let mut control_token: u32 = {};loop {{",
-            function.top_partition
+            function.top_partition.idx()
         )?;
 
         // Match on the control token position to determine which partition to
@@ -194,9 +192,9 @@ fn codegen(manifest: &ModuleManifest, elf: &[u8]) -> Result<String, anyhow::Erro
                     "let input_{} = {};",
                     input_idx,
                     match input {
-                        PartitionInput::DataInput(idx) => format!(
+                        PartitionInput::DataInput(id) => format!(
                             "node_{}.expect(\"PANIC: Encountered use before def in runtime.\")",
-                            idx
+                            id.idx()
                         ),
                         PartitionInput::FunctionArgument(idx) => format!("param_{}", idx),
                         PartitionInput::ArrayConstant(_) => todo!(),
@@ -218,7 +216,7 @@ fn codegen(manifest: &ModuleManifest, elf: &[u8]) -> Result<String, anyhow::Erro
                     rust_code,
                     "{} = Some(output.{});",
                     match output {
-                        PartitionOutput::DataOutput(idx) => format!("node_{}", idx),
+                        PartitionOutput::DataOutput(id) => format!("node_{}", id.idx()),
                         PartitionOutput::ControlIndicator => "todo!()".to_string(),
                     },
                     output_idx
@@ -238,13 +236,13 @@ fn codegen(manifest: &ModuleManifest, elf: &[u8]) -> Result<String, anyhow::Erro
                 write!(
                     rust_code,
                     "return node_{}.expect(\"PANIC: Encountered use before def in runtime.\");",
-                    function.returned_values[0]
+                    function.returned_values[0].idx()
                 )?;
             } else if partition.successor_partitions.len() == 1 {
                 write!(
                     rust_code,
                     "control_token = {};",
-                    partition.successor_partitions[0].1
+                    partition.successor_partitions[0].1.idx()
                 )?;
             } else {
                 todo!();