Fix missing types in compilation with unsafe_panic. (#7995)

Author: ilyalesokhin-starkware

crates/cairo-lang-sierra-generator/src/program_generator.rs

@@ -136,7 +136,7 @@ fn generate_type_declarations_helper(
 fn collect_used_types(
     db: &dyn SierraGenGroup,
     libfunc_declarations: &[program::LibfuncDeclaration],
-    functions: &[Arc<pre_sierra::Function>],
+    functions: &[program::Function],
 ) -> OrderedHashSet<ConcreteTypeId> {
     // Collect types that appear in libfuncs.
     let types_in_libfuncs = libfunc_declarations.iter().flat_map(|libfunc| {
@@ -160,12 +160,16 @@ fn collect_used_types(
         )
     });
 
-    // Collect types that appear in user functions.
-    // This is only relevant for types that are arguments to entry points and are not used in
-    // any libfunc. For example, an empty entry point that gets and returns an empty struct, will
-    // have no libfuncs, but we still need to declare the struct.
-    let types_in_user_functions =
-        functions.iter().flat_map(|func| func.parameters.iter().map(|param| param.ty.clone()));
+    // Gather types used in user-defined functions.
+    // This is necessary for types that are used as entry point arguments but do not appear in any
+    // libfunc. For instance, if an entry point takes and returns an empty struct and no
+    // libfuncs are involved, we still need to declare that struct.
+    // Additionally, we include the return types of functions, since with unsafe panic enabled,
+    // a function that always panics might declare a return type that does not appear in anywhere
+    // else in the program.
+    let types_in_user_functions = functions
+        .iter()
+        .flat_map(|func| chain!(&func.signature.param_types, &func.signature.ret_types).cloned());
 
     chain!(types_in_libfuncs, types_in_user_functions).collect()
 }
@@ -263,31 +267,32 @@ pub fn assemble_program(
     functions: Vec<Arc<pre_sierra::Function>>,
     statements: Vec<pre_sierra::StatementWithLocation>,
 ) -> (program::Program, Vec<Vec<StableLocation>>) {
+    let label_replacer = LabelReplacer::from_statements(&statements);
+    let funcs = functions
+        .into_iter()
+        .map(|function| {
+            let sierra_signature = db.get_function_signature(function.id.clone()).unwrap();
+            program::Function::new(
+                function.id.clone(),
+                function.parameters.clone(),
+                sierra_signature.ret_types.clone(),
+                label_replacer.handle_label_id(function.entry_point),
+            )
+        })
+        .collect_vec();
+
     let libfunc_declarations =
         generate_libfunc_declarations(db, collect_used_libfuncs(&statements).iter());
     let type_declarations =
-        generate_type_declarations(db, collect_used_types(db, &libfunc_declarations, &functions));
+        generate_type_declarations(db, collect_used_types(db, &libfunc_declarations, &funcs));
     // Resolve labels.
-    let label_replacer = LabelReplacer::from_statements(&statements);
     let (resolved_statements, statements_locations) =
         resolve_labels_and_extract_locations(statements, &label_replacer);
-
     let program = program::Program {
         type_declarations,
         libfunc_declarations,
         statements: resolved_statements,
-        funcs: functions
-            .into_iter()
-            .map(|function| {
-                let sierra_signature = db.get_function_signature(function.id.clone()).unwrap();
-                program::Function::new(
-                    function.id.clone(),
-                    function.parameters.clone(),
-                    sierra_signature.ret_types.clone(),
-                    label_replacer.handle_label_id(function.entry_point),
-                )
-            })
-            .collect(),
+        funcs,
     };
     (program, statements_locations)
 }