Use groups of mututally recursive functions in typed modules

Author: giacomocavalieri

compiler-cli/src/publish.rs

@@ -151,6 +151,7 @@ fn check_for_name_squatting(package: &Package) -> Result<(), Error> {
         .definitions
         .functions
         .iter()
+        .flatten()
         .find_map(|function| function.main_function())
     else {
         return Ok(());
@@ -173,7 +174,7 @@ fn check_for_default_main(package: &Package) -> Result<(), Error> {
     let has_default_main = package
         .modules
         .iter()
-        .flat_map(|module| module.ast.definitions.functions.iter())
+        .flat_map(|module| module.ast.definitions.functions.iter().flatten())
         .filter_map(|function| function.main_function())
         .any(|main| main.documentation.is_none() && is_default_main(main, package_name));
 

compiler-core/src/analyse.rs

@@ -16,7 +16,7 @@ use crate::{
         UntypedTypeAlias,
     },
     build::{Origin, Outcome, Target},
-    call_graph::{CallGraphNode, into_dependency_order},
+    call_graph::{CallGraphNode, into_dependency_order, into_mutually_recursive_groups},
     config::PackageConfig,
     dep_tree,
     inline::{self, InlinableFunction},
@@ -282,7 +282,7 @@ impl<'a, A> ModuleAnalyzer<'a, A> {
         // Sort functions and constants into dependency order for inference.
         // Definitions that do not depend on other definitions are inferred
         // first, then ones that depend on those, etc.
-        let mut typed_functions = Vec::with_capacity(definitions.functions.len());
+        let mut typed_functions_groups = Vec::with_capacity(definitions.functions.len());
         let mut typed_constants = Vec::with_capacity(definitions.constants.len());
         let definition_groups =
             match into_dependency_order(definitions.functions, definitions.constants) {
@@ -314,21 +314,22 @@ impl<'a, A> ModuleAnalyzer<'a, A> {
                     &self.module_name,
                 ))
             }
-            for inferred_function in working_functions.drain(..) {
-                typed_functions.push(generalise_function(
-                    inferred_function,
-                    &mut env,
-                    &self.module_name,
-                ));
-            }
+
+            let generalised_functions = working_functions
+                .drain(..)
+                .map(|function| generalise_function(function, &mut env, &self.module_name))
+                .collect_vec();
+
+            typed_functions_groups
+                .append(&mut into_mutually_recursive_groups(generalised_functions));
         }
 
         let typed_definitions = TypedDefinitions {
             imports: typed_imports,
             constants: typed_constants,
             custom_types: typed_custom_types,
             type_aliases: typed_type_aliases,
-            functions: typed_functions,
+            functions: typed_functions_groups,
         };
 
         // Generate warnings for unused items

compiler-core/src/ast.rs

@@ -85,7 +85,12 @@ impl TypedModule {
             .or_else(|| (constants.iter()).find_map(|constant| constant.find_node(byte_index)))
             .or_else(|| (custom_types.iter()).find_map(|type_| type_.find_node(byte_index)))
             .or_else(|| (type_aliases.iter()).find_map(|alias| alias.find_node(byte_index)))
-            .or_else(|| (functions.iter()).find_map(|function| function.find_node(byte_index)))
+            .or_else(|| {
+                functions
+                    .iter()
+                    .flatten()
+                    .find_map(|function| function.find_node(byte_index))
+            })
     }
 
     pub fn find_statement(&self, byte_index: u32) -> Option<&TypedStatement> {
@@ -94,6 +99,7 @@ impl TypedModule {
         self.definitions
             .functions
             .iter()
+            .flatten()
             .find_map(|function| function.find_statement(byte_index))
     }
 
@@ -106,7 +112,11 @@ impl TypedModule {
             functions,
         } = &self.definitions;
 
-        imports.len() + constants.len() + custom_types.len() + type_aliases.len() + functions.len()
+        imports.len()
+            + constants.len()
+            + custom_types.len()
+            + type_aliases.len()
+            + functions.iter().map(|group| group.len()).sum::<usize>()
     }
 }
 
@@ -116,7 +126,7 @@ pub struct TypedDefinitions {
     pub constants: Vec<TypedModuleConstant>,
     pub custom_types: Vec<TypedCustomType>,
     pub type_aliases: Vec<TypedTypeAlias>,
-    pub functions: Vec<TypedFunction>,
+    pub functions: Vec<Vec<TypedFunction>>,
 }
 
 /// The `@target(erlang)` and `@target(javascript)` attributes can be used to

compiler-core/src/ast/visit.rs

@@ -635,7 +635,7 @@ where
         v.visit_typed_custom_type(custom_type);
     }
 
-    for function in functions {
+    for function in functions.iter().flatten() {
         v.visit_typed_function(function);
     }
 }

compiler-core/src/build.rs

@@ -329,7 +329,12 @@ impl Module {
             .chain((constants.iter_mut()).map(DocumentableDefinition::Constant))
             .chain((custom_types.iter_mut()).map(DocumentableDefinition::CustomType))
             .chain((type_aliases.iter_mut()).map(DocumentableDefinition::TypeAlias))
-            .chain((functions.iter_mut()).map(DocumentableDefinition::Function))
+            .chain(
+                functions
+                    .iter_mut()
+                    .flatten()
+                    .map(DocumentableDefinition::Function),
+            )
             .sorted_by_key(|definition| definition.location())
             .collect_vec();