Also cache by return value to improve caching preformance for polymorphic return types

[iFrostizz]

Finding Description

In the result cache, we currently disallow the usage of polymorphic function calls. This is a fix for a compiler crash in the past, where the function call result was cached but since this instruction may have different return types, the SSA couldn't be reconciled when hitting the cache.

Another solution that allows maximum cache hits is to also cache by return types, allowing caching of polymorphic instructions.

Recommendation

diff --git a/compiler/noirc_evaluator/src/ssa/opt/constant_folding/mod.rs b/compiler/noirc_evaluator/src/ssa/opt/constant_folding/mod.rs
index 177669aebc..e116b289da 100644
--- a/compiler/noirc_evaluator/src/ssa/opt/constant_folding/mod.rs
+++ b/compiler/noirc_evaluator/src/ssa/opt/constant_folding/mod.rs
@@ -476,7 +476,7 @@ impl Context {
             let array_get = Instruction::ArrayGet { array: instruction_results[0], index: *index };
 
             // If we encounter an array_get for this address, we know what the result will be.
-            self.cached_instruction_results.cache(dom, array_get, predicate, block, vec![*value]);
+            self.cached_instruction_results.cache(dom, array_get, predicate, block, vec![*value], dfg);
         }
 
         self.cached_instruction_results
@@ -499,6 +499,7 @@ impl Context {
                 predicate,
                 block,
                 instruction_results,
+                dfg
             );
         };

diff --git a/compiler/noirc_evaluator/src/ssa/opt/constant_folding/result_cache.rs b/compiler/noirc_evaluator/src/ssa/opt/constant_folding/result_cache.rs
index aba3950dcf..8b65f8dd39 100644
--- a/compiler/noirc_evaluator/src/ssa/opt/constant_folding/result_cache.rs
+++ b/compiler/noirc_evaluator/src/ssa/opt/constant_folding/result_cache.rs
@@ -9,6 +9,8 @@ use crate::ssa::{
    opt::pure::Purity,
};
use rustc_hash::FxHashMap as HashMap;
+use iter_extended::vecmap;
+use crate::ssa::ir::types::Type;

/// HashMap from `(Instruction, side_effects_enabled_var)` to the results of the instruction.
/// Stored as a two-level map to avoid cloning Instructions during the `.get` call.
@@ -35,29 +37,14 @@ impl InstructionResultCache {
        block: BasicBlockId,
    ) -> Option<CacheResult> {
        let results_for_instruction = self.0.get(instruction)?;
+        let result_types = &vecmap(dfg.instruction_results(id), |&v| dfg.type_of_value(v));

-        let cached_results = results_for_instruction.get(&predicate)?.get(
+        results_for_instruction.get(&predicate)?.get(
            block,
            dom,
+            result_types,
            instruction.has_side_effects(dfg),
-        );
-
-        cached_results.filter(|results| {
-            // This is a hacky solution to https://github.com/noir-lang/noir/issues/9477
-            // We explicitly check that the cached result values are of the same type as expected by the instruction
-            // being checked against the cache and reject if they differ.
-            if let CacheResult::Cached { results, .. } = results {
-                let old_results = dfg.instruction_results(id).to_vec();
-
-                results.len() == old_results.len()
-                    && old_results
-                        .iter()
-                        .zip(results.iter())
-                        .all(|(old, new)| dfg.type_of_value(*old) == dfg.type_of_value(*new))
-            } else {
-                true
-            }
-        })
+        )
    }

    pub(super) fn cache(
@@ -67,13 +54,16 @@ impl InstructionResultCache {
        predicate: Option<ValueId>,
        block: BasicBlockId,
        results: Vec<ValueId>,
+        dfg: &DataFlowGraph,
    ) {
+        let result_types = vecmap(&results, |&v| dfg.type_of_value(v));
+
        self.0
            .entry(instruction)
            .or_default()
            .entry(predicate)
            .or_default()
-            .cache(block, dom, results);
+            .cache(block, dom, results, result_types);
    }

    pub(super) fn remove(
@@ -159,18 +149,19 @@ impl InstructionResultCache {
/// For more information see [`InstructionResultCache`].
#[derive(Default, Debug)]
pub(super) struct ResultCache {
-    result: Option<(BasicBlockId, Vec<ValueId>)>,
+    result: HashMap<Vec<Type>, (BasicBlockId, Vec<ValueId>)>,
}
+
impl ResultCache {
    /// Records that an `Instruction` in block `block` produced the result values `results`.
-    fn cache(&mut self, block: BasicBlockId, dom: &mut DominatorTree, results: Vec<ValueId>) {
-        let overwrite = match self.result {
+    fn cache(&mut self, block: BasicBlockId, dom: &mut DominatorTree, results: Vec<ValueId>, result_types: Vec<Type>) {
+        let overwrite = match self.result.get(&result_types) {
            None => true,
-            Some((origin, _)) => origin != block && dom.dominates(block, origin),
+            Some((origin, _)) => origin != &block && dom.dominates(block, *origin),
        };

        if overwrite {
-            self.result = Some((block, results));
+            self.result.insert(result_types, (block, results));
        }
    }

@@ -184,9 +175,10 @@ impl ResultCache {
        &self,
        block: BasicBlockId,
        dom: &mut DominatorTree,
+        result_types: &Vec<Type>,
        has_side_effects: bool,
    ) -> Option<CacheResult> {
-        self.result.as_ref().and_then(|(origin, results)| {
+        self.result.get(result_types).and_then(|(origin, results)| {
            if dom.dominates(*origin, block) {
                Some(CacheResult::Cached { results })
            } else if !has_side_effects {