Merge pull request #877 from aiken-lang/dependencies-pruning

Only compile modules the project depends on

crates/aiken-project/src/lib.rs

@@ -51,7 +51,7 @@ use pallas::ledger::{
     traverse::ComputeHash,
 };
 use std::{
-    collections::HashMap,
+    collections::{BTreeSet, HashMap},
     fs::{self, File},
     io::BufReader,
     path::{Path, PathBuf},
@@ -185,8 +185,6 @@ where
         destination: Option<PathBuf>,
         include_dependencies: bool,
     ) -> Result<(), Vec<Error>> {
-        self.compile_deps()?;
-
         self.event_listener
             .handle_event(Event::BuildingDocumentation {
                 root: self.root.clone(),
@@ -196,12 +194,12 @@ where
 
         self.read_source_files()?;
 
+        let mut modules = self.parse_sources(self.config.name.clone())?;
+
+        self.type_check(&mut modules, Tracing::silent(), false)?;
+
         let destination = destination.unwrap_or_else(|| self.root.join("docs"));
 
-        let parsed_modules = self.parse_sources(self.config.name.clone())?;
-
-        self.type_check(parsed_modules, Tracing::silent(), false, false)?;
-
         self.event_listener.handle_event(Event::GeneratingDocFiles {
             output_path: destination.clone(),
         });
@@ -283,8 +281,6 @@ where
     }
 
     pub fn compile(&mut self, options: Options) -> Result<(), Vec<Error>> {
-        self.compile_deps()?;
-
         self.event_listener
             .handle_event(Event::StartingCompilation {
                 root: self.root.clone(),
@@ -294,9 +290,9 @@ where
 
         self.read_source_files()?;
 
-        let parsed_modules = self.parse_sources(self.config.name.clone())?;
+        let mut modules = self.parse_sources(self.config.name.clone())?;
 
-        self.type_check(parsed_modules, options.tracing, true, false)?;
+        self.type_check(&mut modules, options.tracing, true)?;
 
         match options.code_gen_mode {
             CodeGenMode::Build(uplc_dump) => {
@@ -537,7 +533,7 @@ where
         Ok(blueprint)
     }
 
-    fn compile_deps(&mut self) -> Result<(), Vec<Error>> {
+    fn with_dependencies(&mut self, parsed_packages: &mut ParsedModules) -> Result<(), Vec<Error>> {
         let manifest = deps::download(&self.event_listener, &self.root, &self.config)?;
 
         for package in manifest.packages {
@@ -565,7 +561,7 @@ where
                         .retain(|def| !matches!(def, Definition::Test { .. }))
                 });
 
-            self.type_check(parsed_modules, Tracing::silent(), true, true)?;
+            parsed_packages.extend(Into::<HashMap<_, _>>::into(parsed_modules));
         }
 
         Ok(())
@@ -680,78 +676,33 @@ where
 
     fn type_check(
         &mut self,
-        mut parsed_modules: ParsedModules,
+        modules: &mut ParsedModules,
         tracing: Tracing,
         validate_module_name: bool,
-        is_dependency: bool,
+    ) -> Result<(), Vec<Error>> {
-    ) -> Result<(), Error> {
+        let our_modules: BTreeSet<String> = modules.keys().cloned().collect();
-        let processing_sequence = parsed_modules.sequence()?;
 
-        for name in processing_sequence {
+        self.with_dependencies(modules)?;
-            if let Some(ParsedModule {
-                name,
-                path,
-                code,
-                kind,
-                extra,
-                package,
-                ast,
-            }) = parsed_modules.remove(&name)
-            {
-                let mut type_warnings = Vec::new();
 
-                let ast = ast
+        for name in modules.sequence(&our_modules)? {
-                    .infer(
+            if let Some(module) = modules.remove(&name) {
+                let (checked_module, warnings) = module.infer(
                     &self.id_gen,
-                        kind,
                     &self.config.name.to_string(),
-                        &self.module_types,
                     tracing,
-                        &mut type_warnings,
+                    validate_module_name,
-                    )
+                    &mut self.module_sources,
-                    .map_err(|error| Error::Type {
+                    &mut self.module_types,
-                        path: path.clone(),
+                    &mut self.functions,
-                        src: code.clone(),
+                    &mut self.data_types,
-                        named: NamedSource::new(path.display().to_string(), code.clone()),
+                )?;
-                        error,
-                    })?;
 
-                if validate_module_name {
+                if our_modules.contains(checked_module.name.as_str()) {
-                    ast.validate_module_name()?;
+                    self.warnings.extend(warnings);
                 }
 
-                // Register any warnings emitted as type warnings
+                self.checked_modules
-                let type_warnings = type_warnings
+                    .insert(checked_module.name.clone(), checked_module);
-                    .into_iter()
-                    .map(|w| Warning::from_type_warning(w, path.clone(), code.clone()));
-
-                if !is_dependency {
-                    self.warnings.extend(type_warnings);
-                }
-
-                // Register module sources for an easier access later.
-                self.module_sources
-                    .insert(name.clone(), (code.clone(), LineNumbers::new(&code)));
-
-                // Register the types from this module so they can be
-                // imported into other modules.
-                self.module_types
-                    .insert(name.clone(), ast.type_info.clone());
-
-                // Register function definitions & data-types for easier access later.
-                ast.register_definitions(&mut self.functions, &mut self.data_types);
-
-                let checked_module = CheckedModule {
-                    kind,
-                    extra,
-                    name: name.clone(),
-                    code,
-                    ast,
-                    package,
-                    input_path: path,
-                };
-
-                self.checked_modules.insert(name, checked_module);
             }
         }
 

crates/aiken-project/src/module.rs

@@ -1,14 +1,20 @@
-use crate::error::Error;
+use crate::{Error, Warning};
 use aiken_lang::{
     ast::{
-        DataType, Definition, Function, Located, ModuleKind, TypedModule, TypedValidator,
+        DataType, DataTypeKey, Definition, Function, FunctionAccessKey, Located, ModuleKind,
-        UntypedModule, Validator,
+        Tracing, TypedDataType, TypedFunction, TypedModule, TypedValidator, UntypedModule,
+        Validator,
     },
+    line_numbers::LineNumbers,
     parser::extra::{comments_before, Comment, ModuleExtra},
+    tipo::TypeInfo,
+    IdGenerator,
 };
+use indexmap::IndexMap;
+use miette::NamedSource;
 use petgraph::{algo, graph::NodeIndex, Direction, Graph};
 use std::{
-    collections::{HashMap, HashSet},
+    collections::{BTreeSet, HashMap},
     io,
     ops::{Deref, DerefMut},
     path::PathBuf,
@@ -38,6 +44,74 @@ impl ParsedModule {
 
         (name, deps)
     }
+
+    #[allow(clippy::too_many_arguments)]
+    pub fn infer(
+        self,
+        id_gen: &IdGenerator,
+        package: &str,
+        tracing: Tracing,
+        validate_module_name: bool,
+        module_sources: &mut HashMap<String, (String, LineNumbers)>,
+        module_types: &mut HashMap<String, TypeInfo>,
+        functions: &mut IndexMap<FunctionAccessKey, TypedFunction>,
+        data_types: &mut IndexMap<DataTypeKey, TypedDataType>,
+    ) -> Result<(CheckedModule, Vec<Warning>), Error> {
+        let mut warnings = Vec::new();
+
+        let ast = self
+            .ast
+            .infer(
+                id_gen,
+                self.kind,
+                package,
+                module_types,
+                tracing,
+                &mut warnings,
+            )
+            .map_err(|error| Error::Type {
+                path: self.path.clone(),
+                src: self.code.clone(),
+                named: NamedSource::new(self.path.display().to_string(), self.code.clone()),
+                error,
+            })?;
+
+        let warnings = warnings
+            .into_iter()
+            .map(|w| Warning::from_type_warning(w, self.path.clone(), self.code.clone()))
+            .collect::<Vec<_>>();
+
+        // Unless we're compiling prelude documentation, prevent keywords in module name
+        if validate_module_name {
+            ast.validate_module_name()?;
+        }
+
+        // Register module sources for an easier access later.
+        module_sources.insert(
+            self.name.clone(),
+            (self.code.clone(), LineNumbers::new(&self.code)),
+        );
+
+        // Register the types from this module so they can be
+        // imported into other modules.
+        module_types.insert(self.name.clone(), ast.type_info.clone());
+
+        // Register function definitions & data-types for easier access later.
+        ast.register_definitions(functions, data_types);
+
+        Ok((
+            CheckedModule {
+                ast,
+                kind: self.kind,
+                extra: self.extra,
+                name: self.name,
+                code: self.code,
+                package: self.package,
+                input_path: self.path,
+            },
+            warnings,
+        ))
+    }
 }
 
 pub struct ParsedModules(HashMap<String, ParsedModule>);
@@ -47,7 +121,7 @@ impl ParsedModules {
         Self(HashMap::new())
     }
 
-    pub fn sequence(&self) -> Result<Vec<String>, Error> {
+    pub fn sequence(&self, our_modules: &BTreeSet<String>) -> Result<Vec<String>, Error> {
         let inputs = self
             .0
             .values()
@@ -56,18 +130,18 @@ impl ParsedModules {
 
         let capacity = inputs.len();
 
-        let mut graph = Graph::<(), ()>::with_capacity(capacity, capacity * 5);
+        let mut graph = Graph::<String, ()>::with_capacity(capacity, capacity * 5);
 
-        // TODO: maybe use a bimap?
         let mut indices = HashMap::with_capacity(capacity);
-        let mut values = HashMap::with_capacity(capacity);
+
+        let mut our_indices = BTreeSet::new();
 
         for (value, _) in &inputs {
-            let index = graph.add_node(());
+            let index = graph.add_node(value.to_string());
-
             indices.insert(value.clone(), index);
-
+            if our_modules.contains(value) {
-            values.insert(index, value.clone());
+                our_indices.insert(index);
+            }
         }
 
         for (value, deps) in inputs {
@@ -80,12 +154,42 @@ impl ParsedModules {
             }
         }
 
+        let mut messed_up_indices = false;
+
+        // Prune the dependency graph to only keep nodes that have a path to one of our (i.e. the
+        // current project) module. This effectively prunes dependencies that are unused from the
+        // graph to ensure that we only compile the modules we actually depend on.
+        graph.retain_nodes(|graph, ix| {
+            // When discarding a node, indices in the graph end up being rewritten. Yet, we need to
+            // know starting indices for our search, so when we remove a dependency, we need find
+            // back what those indices are.
+            if messed_up_indices {
+                our_indices = BTreeSet::new();
+                for j in graph.node_indices() {
+                    if our_modules.contains(graph[j].as_str()) {
+                        our_indices.insert(j);
+                    }
+                }
+            }
+
+            for start in our_indices.iter() {
+                if algo::astar(&*graph, *start, |end| end == ix, |_| 1, |_| 0).is_some() {
+                    messed_up_indices = false;
+                    return true;
+                }
+            }
+
+            messed_up_indices = true;
+            false
+        });
+
         match algo::toposort(&graph, None) {
             Ok(sequence) => {
                 let sequence = sequence
                     .iter()
-                    .filter_map(|i| values.remove(i))
+                    .filter_map(|i| graph.node_weight(*i))
                     .rev()
+                    .cloned()
                     .collect();
 
                 Ok(sequence)
@@ -95,11 +199,12 @@ impl ParsedModules {
 
                 let mut path = vec![];
 
-                find_cycle(origin, origin, &graph, &mut path, &mut HashSet::new());
+                find_cycle(origin, origin, &graph, &mut path, &mut BTreeSet::new());
 
                 let modules = path
                     .iter()
-                    .filter_map(|index| values.remove(index))
+                    .filter_map(|i| graph.node_weight(*i))
+                    .cloned()
                     .collect();
 
                 Err(Error::ImportCycle { modules })
@@ -140,12 +245,12 @@ impl DerefMut for ParsedModules {
     }
 }
 
-fn find_cycle(
+fn find_cycle<W>(
     origin: NodeIndex,
     parent: NodeIndex,
-    graph: &petgraph::Graph<(), ()>,
+    graph: &petgraph::Graph<W, ()>,
     path: &mut Vec<NodeIndex>,
-    seen: &mut HashSet<NodeIndex>,
+    seen: &mut BTreeSet<NodeIndex>,
 ) -> bool {
     seen.insert(parent);