From 98037b6c31c4e204e1a9416b891ff88c7f370540 Mon Sep 17 00:00:00 2001
From: Simone <simone.monica@trailofbits.com>
Date: Fri, 4 Aug 2023 14:21:06 +0200
Subject: [PATCH] Support compiler generate loop functions

---
 src/core/compilation_unit.rs | 53 ++++++++++++++++++++++--------------
 src/core/function.rs         | 11 ++++++++
 2 files changed, 44 insertions(+), 20 deletions(-)

diff --git a/src/core/compilation_unit.rs b/src/core/compilation_unit.rs
index 4ec1bb4..c5c1b1f 100644
--- a/src/core/compilation_unit.rs
+++ b/src/core/compilation_unit.rs
@@ -53,7 +53,12 @@ impl CompilationUnit {
         self.functions.iter().filter(|f| {
             matches!(
                 f.ty(),
-                Type::Constructor | Type::External | Type::View | Type::Private | Type::L1Handler
+                Type::Constructor
+                    | Type::External
+                    | Type::View
+                    | Type::Private
+                    | Type::L1Handler
+                    | Type::Loop
             )
         })
     }
@@ -169,22 +174,26 @@ impl CompilationUnit {
                 f.set_ty(Type::AbiCallContract)
             } else {
                 // Event or private function
-                // Could be an event or a private function in the contract's module
-                let possible_event_name = full_name.rsplit_once("::").unwrap().1;
-
-                let mut found = false;
-                for item in self.abi.items.iter() {
-                    if let Event(e) = item {
-                        if e.name == possible_event_name {
-                            f.set_ty(Type::Event);
-                            found = true;
-                            break;
+                // Could be an event a loop function or a private function in the contract's module
+                if full_name.ends_with(']') {
+                    f.set_ty(Type::Loop);
+                } else {
+                    let possible_event_name = full_name.rsplit_once("::").unwrap().1;
+
+                    let mut found = false;
+                    for item in self.abi.items.iter() {
+                        if let Event(e) = item {
+                            if e.name == possible_event_name {
+                                f.set_ty(Type::Event);
+                                found = true;
+                                break;
+                            }
                         }
                     }
-                }
 
-                if !found {
-                    f.set_ty(Type::Private);
+                    if !found {
+                        f.set_ty(Type::Private);
+                    }
                 }
             }
         }
@@ -285,12 +294,16 @@ impl CompilationUnit {
         while changed {
             changed = false;
 
-            for calling_function in self
-                .functions
-                .iter()
-                .filter(|f| matches!(f.ty(), Type::External | Type::L1Handler | Type::Private))
-            {
-                for function_call in calling_function.private_functions_calls() {
+            for calling_function in self.functions.iter().filter(|f| {
+                matches!(
+                    f.ty(),
+                    Type::External | Type::L1Handler | Type::Private | Type::Loop
+                )
+            }) {
+                for function_call in calling_function
+                    .private_functions_calls()
+                    .chain(calling_function.loop_functions_calls())
+                {
                     // It will always be an invocation
                     if let GenStatement::Invocation(invoc) = function_call {
                         // The core lib func instance
diff --git a/src/core/function.rs b/src/core/function.rs
index 802e9a7..54cd029 100644
--- a/src/core/function.rs
+++ b/src/core/function.rs
@@ -47,6 +47,9 @@ pub enum Type {
     AbiLibraryCall,
     /// L1 handler function
     L1Handler,
+    /// Compiler generated function when a loop is present in a user defined function
+    /// Loop generate functions in core library are still Core type
+    Loop,
 }
 
 #[derive(Clone)]
@@ -73,6 +76,8 @@ pub struct Function {
     external_functions_calls: Vec<SierraStatement>,
     /// Library functions called through an ABI trait (NOTE it doesn't have library functions called using the syscall directly)
     library_functions_calls: Vec<SierraStatement>,
+    /// Loop functions called
+    loop_functions_calls: Vec<SierraStatement>,
     /// Analyses results
     analyses: Analyses,
 }
@@ -91,6 +96,7 @@ impl Function {
             events_emitted: Vec::new(),
             external_functions_calls: Vec::new(),
             library_functions_calls: Vec::new(),
+            loop_functions_calls: Vec::new(),
             analyses: Analyses::default(),
         }
     }
@@ -132,6 +138,10 @@ impl Function {
         self.library_functions_calls.iter()
     }
 
+    pub fn loop_functions_calls(&self) -> impl Iterator<Item = &SierraStatement> {
+        self.loop_functions_calls.iter()
+    }
+
     pub fn analyses(&self) -> &Analyses {
         &self.analyses
     }
@@ -226,6 +236,7 @@ impl Function {
                                 Type::AbiLibraryCall => {
                                     self.library_functions_calls.push(s.clone())
                                 }
+                                Type::Loop => self.loop_functions_calls.push(s.clone()),
                                 _ => (),
                             }
                             break;