fix: perserve purities after SSA normalization

compiler/noirc_evaluator/src/ssa/ir/dfg/simplify.rs

@@ -305,7 +305,7 @@ pub(crate) fn simplify(
 }
 
 /// Given an array access on a length 1 array such as:
-/// ```
+/// ```ssa
 /// v2 = make_array [v0] : [Field; 1]
 /// v3 = array_get v2, index v1 -> Field
 /// ```
@@ -316,7 +316,7 @@ pub(crate) fn simplify(
 /// We then inject an explicit assertion that the index variable has the value zero while replacing the value
 /// being used in the `array_get` instruction with a constant value of zero. This then results in the SSA:
 ///
-/// ```
+/// ```ssa
 /// v2 = make_array [v0] : [Field; 1]
 /// constrain v1 == u32 0, "Index out of bounds"
 /// v4 = array_get v2, index u32 0 -> Field

compiler/noirc_evaluator/src/ssa/opt/defunctionalize.rs

@@ -48,6 +48,7 @@ use crate::ssa::{
         types::{NumericType, Type},
         value::{Value, ValueId},
     },
+    opt::pure::Purity,
     ssa_gen::Ssa,
 };
 use fxhash::FxHashMap as HashMap;
@@ -97,14 +98,19 @@ impl Ssa {
         let variants = find_variants(&self);
 
         // Generate the apply functions for the provided variants
-        let apply_functions = create_apply_functions(&mut self, variants);
+        let (apply_functions, purities) = create_apply_functions(&mut self, variants);
 
         // Setup the pass context
         let context = DefunctionalizationContext { apply_functions };
 
         // Run defunctionalization over all functions in the SSA
         context.defunctionalize_all(&mut self);
 
+        let purities = Arc::new(purities);
+        for function in self.functions.values_mut() {
+            function.dfg.set_function_purities(purities.clone());
+        }
+
         // Check that we have established the properties expected from this pass.
         #[cfg(debug_assertions)]
         self.functions.values().for_each(defunctionalize_post_check);
@@ -415,11 +421,21 @@ fn find_dynamic_dispatches(func: &Function) -> BTreeSet<Signature> {
 /// - `variants_map`:  [Variants]
 ///
 /// # Returns
-/// [ApplyFunctions] keyed by each function's signature _before_ functions are changed
-/// into field types. The inner apply function itself will have its defunctionalized type,
-/// with function values represented as field values.
-fn create_apply_functions(ssa: &mut Ssa, variants_map: Variants) -> ApplyFunctions {
+/// - [ApplyFunctions] keyed by each function's signature _before_ functions are changed
+///   into field types. The inner apply function itself will have its defunctionalized type,
+///   with function values represented as field values.
+/// - [HashMap<FunctionId, Purity>] with purities that must be set to all functions in the SSA,
+///   as this function might have created dummy pure functions.
+fn create_apply_functions(
+    ssa: &mut Ssa,
+    variants_map: Variants,
+) -> (ApplyFunctions, HashMap<FunctionId, Purity>) {
     let mut apply_functions = HashMap::default();
+    let mut purities = if ssa.functions.is_empty() {
+        HashMap::default()
+    } else {
+        (*ssa.functions.iter().next().unwrap().1.dfg.function_purities).clone()
+    };
 
     for ((signature, runtime), variants) in variants_map.into_iter() {
         let dispatches_to_multiple_functions = variants.len() > 1;
@@ -451,13 +467,13 @@ fn create_apply_functions(ssa: &mut Ssa, variants_map: Variants) -> ApplyFunctio
             // If no variants exist for a dynamic call we leave removing those dead calls and parameters to DIE.
             // However, we have to construct a dummy function for these dead calls as to keep a well formed SSA
             // and to not break the semantics of other SSA passes before DIE is reached.
-            create_dummy_function(ssa, defunctionalized_signature, runtime)
+            create_dummy_function(ssa, defunctionalized_signature, runtime, &mut purities)
         };
         apply_functions
             .insert((signature, runtime), ApplyFunction { id, dispatches_to_multiple_functions });
     }
 
-    apply_functions
+    (apply_functions, purities)
 }
 
 /// Transforms a [FunctionId] into a [FieldElement]
@@ -500,11 +516,9 @@ fn create_apply_function(
     // We will be borrowing `ssa` mutably so we need to fetch this shared information
     // before attempting to add a new function to the SSA.
     let globals = ssa.main().dfg.globals.clone();
-    let purities = ssa.main().dfg.function_purities.clone();
     ssa.add_fn(|id| {
         let mut function_builder = FunctionBuilder::new("apply".to_string(), id);
         function_builder.set_globals(globals);
-        function_builder.set_purities(purities);
 
         // We want to push for apply functions to be inlined more aggressively;
         // they are expected to be optimized away by constants visible at the call site.
@@ -637,6 +651,7 @@ fn create_dummy_function(
     ssa: &mut Ssa,
     signature: Signature,
     caller_runtime: RuntimeType,
+    purities: &mut HashMap<FunctionId, Purity>,
 ) -> FunctionId {
     ssa.add_fn(|id| {
         let mut function_builder = FunctionBuilder::new("apply_dummy".to_string(), id);
@@ -658,9 +673,7 @@ fn create_dummy_function(
         // As the dummy function is just meant to be a placeholder for any calls to
         // higher-order functions without variants, we want the function to be marked pure
         // so that dead instruction elimination can remove any calls to it.
-        let mut purities = HashMap::default();
-        purities.insert(id, super::pure::Purity::Pure);
-        function_builder.set_purities(Arc::new(purities));
+        purities.insert(id, Purity::Pure);
 
         let results =
             vecmap(signature.returns, |typ| make_dummy_return_data(&mut function_builder, &typ));
@@ -1639,7 +1652,7 @@ mod tests {
         let variants = find_variants(&ssa);
         assert_eq!(variants.len(), 2);
 
-        let apply_functions = create_apply_functions(&mut ssa, variants);
+        let (apply_functions, _purities) = create_apply_functions(&mut ssa, variants);
         // This was 1 before this bug was fixed.
         assert_eq!(apply_functions.len(), 2);
     }

compiler/noirc_evaluator/src/ssa/opt/normalize_value_ids.rs

@@ -1,4 +1,4 @@
-use std::collections::BTreeMap;
+use std::{collections::BTreeMap, sync::Arc};
 
 use crate::ssa::{
     ir::{
@@ -57,12 +57,28 @@ struct IdMaps {
 
 impl Context {
     fn populate_functions(&mut self, functions: &BTreeMap<FunctionId, Function>) {
+        let Some(old_purities) = &functions.iter().next().map(|f| &f.1.dfg.function_purities)
+        else {
+            return;
+        };
+        let mut new_purities = HashMap::default();
+
         for (id, function) in functions {
             self.functions.insert_with_id(|new_id| {
                 self.new_ids.function_ids.insert(*id, new_id);
+
+                if let Some(purity) = old_purities.get(id) {
+                    new_purities.insert(new_id, *purity);
+                }
+
                 Function::clone_signature(new_id, function)
             });
         }
+
+        let new_purities = Arc::new(new_purities);
+        for new_id in self.new_ids.function_ids.values() {
+            self.functions[*new_id].dfg.set_function_purities(new_purities.clone());
+        }
     }
 
     fn normalize_ids(&mut self, old_function: &mut Function) {