aiken-project 1.1.17

use crate::{Error, Warning};
use aiken_lang::{
    IdGenerator,
    ast::{
        DataType, DataTypeKey, Definition, Function, FunctionAccessKey, Located, ModuleKind,
        Tracing, TypedDataType, TypedFunction, TypedModule, TypedValidator, UntypedModule,
        Validator,
    },
    expr::TypedExpr,
    line_numbers::LineNumbers,
    parser::extra::{Comment, ModuleExtra, comments_before},
    tipo::TypeInfo,
};
use indexmap::IndexMap;
use miette::NamedSource;
use petgraph::{Direction, Graph, algo, graph::NodeIndex};
use std::{
    collections::{BTreeSet, HashMap},
    io,
    ops::{Deref, DerefMut},
    path::PathBuf,
};

#[derive(Debug)]
pub struct ParsedModule {
    pub path: PathBuf,
    pub name: String,
    pub code: String,
    pub kind: ModuleKind,
    pub package: String,
    pub ast: UntypedModule,
    pub extra: ModuleExtra,
}

impl ParsedModule {
    pub fn deps_for_graph(&self, env_modules: &[String]) -> (String, Vec<String>) {
        let name = self.name.clone();
        let deps: Vec<_> = self.ast.dependencies(env_modules);
        (name, deps)
    }

    #[allow(clippy::too_many_arguments)]
    pub fn infer(
        self,
        id_gen: &IdGenerator,
        package: &str,
        tracing: Tracing,
        env: Option<&str>,
        validate_module_name: bool,
        module_sources: &mut HashMap<String, (String, LineNumbers)>,
        module_types: &mut HashMap<String, TypeInfo>,
        functions: &mut IndexMap<FunctionAccessKey, TypedFunction>,
        constants: &mut IndexMap<FunctionAccessKey, TypedExpr>,
        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,
                env,
            )
            .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, constants, 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>);

impl ParsedModules {
    pub fn new() -> Self {
        Self(HashMap::new())
    }

    pub fn sequence(&self, our_modules: &BTreeSet<String>) -> Result<Vec<String>, Error> {
        let env_modules = self
            .0
            .values()
            .filter_map(|m| match m.kind {
                ModuleKind::Env => Some(m.name.clone()),
                ModuleKind::Lib | ModuleKind::Validator | ModuleKind::Config => None,
            })
            .collect::<Vec<String>>();

        let inputs = self
            .0
            .values()
            .map(|m| m.deps_for_graph(&env_modules))
            .collect::<Vec<(String, Vec<String>)>>();

        let capacity = inputs.len();

        let mut graph = Graph::<String, ()>::with_capacity(capacity, capacity * 5);

        let mut indices = HashMap::with_capacity(capacity);

        let mut our_indices = BTreeSet::new();

        for (value, _) in &inputs {
            let index = graph.add_node(value.to_string());
            indices.insert(value.clone(), index);
            if our_modules.contains(value) {
                our_indices.insert(index);
            }
        }

        for (value, deps) in inputs {
            if let Some(from_index) = indices.get(&value) {
                let deps = deps.into_iter().filter_map(|dep| indices.get(&dep));

                for to_index in deps {
                    graph.add_edge(*from_index, *to_index, ());
                }
            }
        }

        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| graph.node_weight(*i))
                    .rev()
                    .cloned()
                    .collect();

                Ok(sequence)
            }
            Err(cycle) => {
                let origin = cycle.node_id();

                let mut path = vec![];

                find_cycle(origin, origin, &graph, &mut path, &mut BTreeSet::new());

                let modules = path
                    .iter()
                    .filter_map(|i| graph.node_weight(*i))
                    .cloned()
                    .collect();

                Err(Error::ImportCycle { modules })
            }
        }
    }
}

impl Default for ParsedModules {
    fn default() -> Self {
        Self::new()
    }
}

impl From<HashMap<String, ParsedModule>> for ParsedModules {
    fn from(parsed_modules: HashMap<String, ParsedModule>) -> Self {
        ParsedModules(parsed_modules)
    }
}

impl From<ParsedModules> for HashMap<String, ParsedModule> {
    fn from(parsed_modules: ParsedModules) -> Self {
        parsed_modules.0
    }
}

impl Deref for ParsedModules {
    type Target = HashMap<String, ParsedModule>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for ParsedModules {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

fn find_cycle<W>(
    origin: NodeIndex,
    parent: NodeIndex,
    graph: &petgraph::Graph<W, ()>,
    path: &mut Vec<NodeIndex>,
    seen: &mut BTreeSet<NodeIndex>,
) -> bool {
    seen.insert(parent);

    for node in graph.neighbors_directed(parent, Direction::Outgoing) {
        if node == origin {
            path.push(node);

            return true;
        }

        if seen.contains(&node) {
            continue;
        }

        if find_cycle(origin, node, graph, path, seen) {
            path.push(node);

            return true;
        }
    }

    false
}

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct CheckedModule {
    pub name: String,
    pub code: String,
    pub input_path: PathBuf,
    pub kind: ModuleKind,
    pub package: String,
    pub ast: TypedModule,
    pub extra: ModuleExtra,
}

impl CheckedModule {
    pub fn skip_doc_generation(&self) -> bool {
        self.ast
            .docs
            .first()
            .map(|s| s.as_str().trim())
            .unwrap_or_default()
            == "@hidden"
    }

    pub fn to_cbor(&self) -> Vec<u8> {
        let mut module_bytes = vec![];

        ciborium::into_writer(&self, &mut module_bytes)
            .expect("modules should not fail to serialize");

        module_bytes
    }

    pub fn from_cbor(bytes: &[u8]) -> Result<Self, ciborium::de::Error<io::Error>> {
        ciborium::from_reader(bytes)
    }

    pub fn to_cbor_hex(&self) -> (String, Vec<u8>) {
        let module_bytes = self.to_cbor();
        let hex_str = hex::encode(&module_bytes);

        (hex_str, module_bytes)
    }

    pub fn find_node(&self, byte_index: usize) -> Option<Located<'_>> {
        self.ast.find_node(byte_index)
    }

    pub fn attach_doc_and_module_comments(&mut self) {
        // Module Comments
        self.ast.docs = self
            .extra
            .module_comments
            .iter()
            .map(|span| {
                Comment::from((span, self.code.as_str()))
                    .content
                    .to_string()
            })
            .collect();

        // Order definitions to avoid dissociating doc comments from them
        let mut definitions: Vec<_> = self.ast.definitions.iter_mut().collect();
        definitions.sort_by(|a, b| a.location().start.cmp(&b.location().start));

        // Doc Comments
        let mut doc_comments = self.extra.doc_comments.iter().peekable();
        for def in &mut definitions {
            let docs: Vec<&str> =
                comments_before(&mut doc_comments, def.location().start, &self.code);
            if !docs.is_empty() {
                let doc = docs.join("\n");
                def.put_doc(doc);
            }

            match def {
                Definition::DataType(DataType { constructors, .. }) => {
                    for constructor in constructors {
                        let docs: Vec<&str> = comments_before(
                            &mut doc_comments,
                            constructor.location.start,
                            &self.code,
                        );
                        if !docs.is_empty() {
                            let doc = docs.join("\n");
                            constructor.put_doc(doc);
                        }

                        for argument in constructor.arguments.iter_mut() {
                            let docs: Vec<&str> = comments_before(
                                &mut doc_comments,
                                argument.location.start,
                                &self.code,
                            );
                            if !docs.is_empty() {
                                let doc = docs.join("\n");
                                argument.put_doc(doc);
                            }
                        }
                    }
                }
                Definition::Fn(Function { arguments, .. }) => {
                    for argument in arguments {
                        let docs: Vec<&str> =
                            comments_before(&mut doc_comments, argument.location.start, &self.code);

                        if !docs.is_empty() {
                            let doc = docs.join("\n");
                            argument.put_doc(doc);
                        }
                    }
                }
                Definition::Validator(Validator {
                    params,
                    handlers,
                    fallback,
                    ..
                }) => {
                    for param in params {
                        let docs: Vec<&str> =
                            comments_before(&mut doc_comments, param.location.start, &self.code);

                        if !docs.is_empty() {
                            let doc = docs.join("\n");
                            param.put_doc(doc);
                        }
                    }

                    for handler in handlers.iter_mut() {
                        for argument in handler.arguments.iter_mut() {
                            let docs: Vec<&str> = comments_before(
                                &mut doc_comments,
                                argument.location.start,
                                &self.code,
                            );

                            if !docs.is_empty() {
                                let doc = docs.join("\n");
                                argument.put_doc(doc);
                            }
                        }
                    }

                    for argument in fallback.arguments.iter_mut() {
                        let docs: Vec<&str> =
                            comments_before(&mut doc_comments, argument.location.start, &self.code);

                        if !docs.is_empty() {
                            let doc = docs.join("\n");
                            argument.put_doc(doc);
                        }
                    }
                }
                _ => (),
            }
        }
    }
}

#[derive(Default, Debug, Clone)]
pub struct CheckedModules(HashMap<String, CheckedModule>);

impl From<HashMap<String, CheckedModule>> for CheckedModules {
    fn from(checked_modules: HashMap<String, CheckedModule>) -> Self {
        CheckedModules(checked_modules)
    }
}

impl From<CheckedModules> for HashMap<String, CheckedModule> {
    fn from(checked_modules: CheckedModules) -> Self {
        checked_modules.0
    }
}

impl<'a> From<&'a CheckedModules> for &'a HashMap<String, CheckedModule> {
    fn from(checked_modules: &'a CheckedModules) -> Self {
        &checked_modules.0
    }
}

impl CheckedModules {
    pub fn singleton(module: CheckedModule) -> Self {
        let mut modules = Self::default();
        modules.insert(module.name.clone(), module);
        modules
    }

    // todo: this might need fixing
    pub fn validators(&self) -> impl Iterator<Item = (&CheckedModule, &TypedValidator)> {
        let mut items = vec![];

        for validator_module in self.0.values().filter(|module| module.kind.is_validator()) {
            for some_definition in validator_module.ast.definitions() {
                if let Definition::Validator(def) = some_definition {
                    items.push((validator_module, def));
                }
            }
        }

        items.sort_by(|left, right| {
            (
                left.0.package.to_string(),
                left.0.name.to_string(),
                left.1.name.to_string(),
            )
                .cmp(&(
                    right.0.package.to_string(),
                    right.0.name.to_string(),
                    right.1.name.to_string(),
                ))
        });

        items.into_iter()
    }

    pub fn functions(&self) -> impl Iterator<Item = (&CheckedModule, &TypedFunction)> {
        let mut items = vec![];

        for module in self.0.values() {
            for some_definition in module.ast.definitions() {
                if let Definition::Fn(def) = some_definition {
                    items.push((module, def));
                }
            }
        }

        items.into_iter()
    }

    pub fn into_validators(self) -> impl Iterator<Item = CheckedModule> {
        self.0
            .into_values()
            .filter(|module| module.kind.is_validator())
    }
}

impl Deref for CheckedModules {
    type Target = HashMap<String, CheckedModule>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for CheckedModules {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}