chomsky_linker/
lib.rs

1#![warn(missing_docs)]
2
3use chomsky_uir::{Analysis, EGraph, IKun, Id, Language};
4use std::collections::HashMap;
5
6/// Linker for merging multiple IKun E-Graphs and resolving cross-module references.
7pub struct IKunLinker<A: Analysis<IKun>> {
8    /// The merged global E-Graph.
9    pub global_graph: EGraph<IKun, A>,
10    /// Mapping from (ModuleName, ExportName) to Id in the global graph.
11    exports: HashMap<(String, String), Id>,
12}
13
14impl<A: Analysis<IKun> + Default> IKunLinker<A> {
15    pub fn new() -> Self {
16        Self {
17            global_graph: EGraph::new(),
18            exports: HashMap::new(),
19        }
20    }
21
22    /// Add a module's E-Graph to the global graph.
23    /// This will copy all nodes and record exports.
24    pub fn add_module(&mut self, module_name: &str, graph: &EGraph<IKun, A>) {
25        // Step 1: Record exports in the source graph
26        for entry in graph.classes.iter() {
27            let class = entry.value();
28            for node in &class.nodes {
29                if let IKun::Export(_name, _body_id) = node {
30                    // Record that this module has this export.
31                }
32            }
33        }
34
35        self.merge_graph(module_name, graph);
36    }
37
38    fn merge_graph(&mut self, module_name: &str, source: &EGraph<IKun, A>) {
39        let mut id_map = HashMap::new();
40
41        // Step 1: Add all nodes and build id_map
42        for entry in source.classes.iter() {
43            let class = entry.value();
44            for node in &class.nodes {
45                let new_id = self.add_node_recursive(node, source, &mut id_map);
46                id_map.insert(class.id, new_id);
47
48                // Record global export ID
49                if let IKun::Export(name, _) = node {
50                    self.exports
51                        .insert((module_name.to_string(), name.clone()), new_id);
52                }
53            }
54        }
55    }
56
57    fn add_node_recursive(
58        &mut self,
59        node: &IKun,
60        source: &EGraph<IKun, A>,
61        id_map: &mut HashMap<Id, Id>,
62    ) -> Id {
63        let remapped = node.map_children(|child_id| {
64            if let Some(&new_id) = id_map.get(&child_id) {
65                new_id
66            } else {
67                // If child not in map yet, we must find it in source and add it
68                let child_class = source.get_class(child_id);
69                // For simplicity, pick the first node in the class
70                let child_node = &child_class.nodes[0];
71                let new_id = self.add_node_recursive(child_node, source, id_map);
72                id_map.insert(child_id, new_id);
73                new_id
74            }
75        });
76        self.global_graph.add(remapped)
77    }
78
79    /// Link all imports to their corresponding exports.
80    pub fn link(&mut self) {
81        let mut links_to_make = Vec::new();
82
83        for entry in self.global_graph.classes.iter() {
84            let id = *entry.key();
85            let class = entry.value();
86            for node in &class.nodes {
87                if let IKun::Import(mod_name, sym_name) = node {
88                    if let Some(&export_id) =
89                        self.exports.get(&(mod_name.clone(), sym_name.clone()))
90                    {
91                        links_to_make.push((id, export_id));
92                    }
93                }
94            }
95        }
96
97        // Merge e-classes in the global graph
98        for (import_id, export_id) in links_to_make {
99            self.global_graph.union(import_id, export_id);
100        }
101
102        self.global_graph.rebuild();
103    }
104}
chomsky_linker/lib.rs

chomsky_linker/
lib.rs
