dev/src/natu/scope_graph.rs

use crate::ast::{self, ParsedFile};

#[derive(Debug)]
pub(crate) struct ScopeGraph {
    nodes: Vec<Node>,
    edges: Vec<Edge>,
    lets: Vec<LetScope>,
}

impl ScopeGraph {
    pub(crate) fn from_file(file: &ParsedFile) -> Self {
        let mut graph = Self {
            nodes: Vec::new(),
            edges: Vec::new(),
            lets: Vec::new(),
        };
        let root = graph.add_scope(None, ScopeKind::Root);
        graph.build_expr(&file.root, root);
        graph
    }

    pub(crate) fn let_scopes(&self) -> &[LetScope] {
        &self.lets
    }

    fn build_expr(&mut self, expr: &ast::Expr, scope: NodeId) -> Vec<NodeId> {
        match expr {
            ast::Expr::Ident(ident) => vec![self.add_reference(&ident.name, scope)],
            ast::Expr::Let {
                range,
                bindings,
                body,
            } => self.build_let(*range, bindings, body, scope),
            ast::Expr::AttrSet(bindings) => bindings
                .iter()
                .flat_map(|binding| self.build_expr(&binding.value, scope))
                .collect(),
            ast::Expr::Lambda { parameters, body } => {
                let scope = self.add_scope(Some(scope), ScopeKind::Lambda);
                for parameter in parameters {
                    self.add_declaration(&parameter.name, scope);
                }
                self.build_expr(body, scope)
            }
            ast::Expr::Apply(application) => [
                self.build_expr(&application.lambda, scope),
                self.build_expr(&application.argument, scope),
            ]
            .concat(),
            ast::Expr::Other(expressions) => expressions
                .iter()
                .flat_map(|expression| self.build_expr(expression, scope))
                .collect(),
            ast::Expr::Select(select) => vec![self.add_reference(&select.receiver, scope)],
        }
    }

    fn build_let(
        &mut self,
        range: ast::SourceRange,
        bindings: &[ast::Binding],
        body: &ast::Expr,
        parent: NodeId,
    ) -> Vec<NodeId> {
        let scope = self.add_scope(Some(parent), ScopeKind::Let);
        let mut graph_bindings = bindings
            .iter()
            .map(|binding| Binding {
                declaration: self.add_declaration(&binding.key.name, scope),
                name: binding.key.name.clone(),
                references: Vec::new(),
                is_import: is_import(&binding.value),
                is_lambda: matches!(binding.value, ast::Expr::Lambda { .. }),
                range: binding.range,
            })
            .collect::<Vec<_>>();

        for (binding, graph_binding) in bindings.iter().zip(&mut graph_bindings) {
            graph_binding.references = self.build_expr(&binding.value, scope);
        }

        let entry_references = self.build_expr(body, scope);
        let references = graph_bindings
            .iter()
            .flat_map(|binding| binding.references.iter().copied())
            .chain(entry_references.iter().copied())
            .collect();
        self.lets.push(LetScope {
            range,
            bindings: graph_bindings,
            entry_references,
        });

        references
    }

    pub(crate) fn crosses_lambda_scope(&self, reference: NodeId, declaration: NodeId) -> bool {
        let Some(mut scope) = self.edge_target(reference, Node::is_scope) else {
            return false;
        };
        let Some(declaration_scope) = self
            .edges
            .iter()
            .find(|edge| edge.to == declaration && self.nodes[edge.from.0].is_scope())
            .map(|edge| edge.from)
        else {
            return false;
        };

        loop {
            if scope == declaration_scope {
                return false;
            }
            if matches!(self.nodes[scope.0], Node::Scope(ScopeKind::Lambda)) {
                return true;
            }
            let Some(parent) = self.edge_target(scope, Node::is_scope) else {
                return false;
            };
            scope = parent;
        }
    }

    pub(crate) fn resolve(&self, reference: NodeId) -> Vec<NodeId> {
        let Node::Reference(name) = &self.nodes[reference.0] else {
            return Vec::new();
        };
        let Some(mut scope) = self.edge_target(reference, Node::is_scope) else {
            return Vec::new();
        };

        loop {
            let declarations = self
                .edge_targets(scope, |node| matches!(node, Node::Declaration(declaration) if declaration.name == *name))
                .collect::<Vec<_>>();
            if !declarations.is_empty() {
                return declarations;
            }

            let Some(parent) = self.edge_target(scope, Node::is_scope) else {
                return Vec::new();
            };
            scope = parent;
        }
    }

    fn add_scope(&mut self, parent: Option<NodeId>, kind: ScopeKind) -> NodeId {
        let scope = self.add_node(Node::Scope(kind));
        if let Some(parent) = parent {
            self.add_edge(scope, parent);
        }
        scope
    }

    fn add_declaration(&mut self, name: &str, scope: NodeId) -> NodeId {
        let declaration = self.add_node(Node::Declaration(Declaration {
            name: name.to_owned(),
        }));
        self.add_edge(scope, declaration);
        declaration
    }

    fn add_reference(&mut self, name: &str, scope: NodeId) -> NodeId {
        let reference = self.add_node(Node::Reference(name.to_owned()));
        self.add_edge(reference, scope);
        reference
    }

    fn add_node(&mut self, node: Node) -> NodeId {
        let id = NodeId(self.nodes.len());
        self.nodes.push(node);
        id
    }

    fn add_edge(&mut self, from: NodeId, to: NodeId) {
        self.edges.push(Edge { from, to });
    }

    fn edge_target(&self, from: NodeId, predicate: fn(&Node) -> bool) -> Option<NodeId> {
        self.edge_targets(from, predicate).next()
    }

    fn edge_targets<'graph>(
        &'graph self,
        from: NodeId,
        predicate: impl Fn(&Node) -> bool + 'graph,
    ) -> impl Iterator<Item = NodeId> + 'graph {
        self.edges
            .iter()
            .filter(move |edge| edge.from == from && predicate(&self.nodes[edge.to.0]))
            .map(|edge| edge.to)
    }
}

fn is_import(expr: &ast::Expr) -> bool {
    match unwrap_expr(expr) {
        ast::Expr::Apply(application) => is_import_callee(&application.lambda),
        _ => false,
    }
}

fn is_import_callee(expr: &ast::Expr) -> bool {
    match unwrap_expr(expr) {
        ast::Expr::Apply(application) => is_import_callee(&application.lambda),
        ast::Expr::Ident(ident) => ident.name == "import",
        _ => false,
    }
}

fn unwrap_expr(expr: &ast::Expr) -> &ast::Expr {
    match expr {
        ast::Expr::Other(expressions) if expressions.len() == 1 => unwrap_expr(&expressions[0]),
        _ => expr,
    }
}

#[derive(Debug)]
enum Node {
    Scope(ScopeKind),
    Declaration(Declaration),
    Reference(String),
}

impl Node {
    fn is_scope(&self) -> bool {
        matches!(self, Self::Scope(_))
    }
}

#[derive(Debug)]
enum ScopeKind {
    Root,
    Let,
    Lambda,
}

#[derive(Debug)]
struct Declaration {
    name: String,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct NodeId(usize);

#[derive(Debug)]
struct Edge {
    from: NodeId,
    to: NodeId,
}

#[derive(Debug)]
pub(crate) struct LetScope {
    pub(crate) range: ast::SourceRange,
    pub(crate) bindings: Vec<Binding>,
    pub(crate) entry_references: Vec<NodeId>,
}

#[derive(Debug)]
pub(crate) struct Binding {
    pub(crate) declaration: NodeId,
    pub(crate) name: String,
    pub(crate) references: Vec<NodeId>,
    pub(crate) is_import: bool,
    pub(crate) is_lambda: bool,
    pub(crate) range: ast::SourceRange,
}