From 1752e784eb3b89993c42dbb226b532f123f625df Mon Sep 17 00:00:00 2001
From: ethan-000 <s2026080@ed.ac.uk>
Date: Tue, 29 Aug 2023 12:52:07 +0100
Subject: [PATCH] optimize new locations

---
 acvm/src/compiler/mod.rs | 40 ++++++++++++----------------------------
 1 file changed, 12 insertions(+), 28 deletions(-)

diff --git a/acvm/src/compiler/mod.rs b/acvm/src/compiler/mod.rs
index fc240d3f..281b66ad 100644
--- a/acvm/src/compiler/mod.rs
+++ b/acvm/src/compiler/mod.rs
@@ -36,41 +36,25 @@ pub struct AcirTransformationMap {
 }
 
 impl AcirTransformationMap {
-    pub fn new_locations(&self, old_location: OpcodeLocation) -> Vec<OpcodeLocation> {
+    pub fn new_locations(
+        &self,
+        old_location: OpcodeLocation,
+    ) -> impl Iterator<Item = OpcodeLocation> + '_ {
         let old_acir_index = match old_location {
             OpcodeLocation::Acir(index) => index,
             OpcodeLocation::Brillig { acir_index, .. } => acir_index,
         };
 
-        let new_opcode_indexes: Vec<usize> = self
-            .acir_opcode_positions
+        self.acir_opcode_positions
             .iter()
             .enumerate()
-            .filter_map(
-                |(new_index, &old_index)| {
-                    if old_index == old_acir_index {
-                        Some(new_index)
-                    } else {
-                        None
-                    }
-                },
-            )
-            .collect();
-
-        match old_location {
-            OpcodeLocation::Acir(_) => new_opcode_indexes
-                .iter()
-                .map(|&new_index| OpcodeLocation::Acir(new_index))
-                .collect(),
-            OpcodeLocation::Brillig { brillig_index, .. } => {
-                assert!(
-                    new_opcode_indexes.len() == 1,
-                    "The transformation must not decompose or remove brillig opcodes"
-                );
-
-                vec![OpcodeLocation::Brillig { acir_index: new_opcode_indexes[0], brillig_index }]
-            }
-        }
+            .filter(move |(_, &old_index)| old_index == old_acir_index)
+            .map(move |(new_index, _)| match old_location {
+                OpcodeLocation::Acir(_) => OpcodeLocation::Acir(new_index),
+                OpcodeLocation::Brillig { brillig_index, .. } => {
+                    OpcodeLocation::Brillig { acir_index: new_index, brillig_index }
+                }
+            })
     }
 }