Skip to main content

chomsky_emit/
lib.rs

1#![warn(missing_docs)]
2
3use chomsky_extract::{Backend, BackendArtifact, IKunTree};
4use chomsky_types::{ChomskyError, ChomskyResult};
5use gaia_assembler::backends::Backend as GaiaBackend;
6pub use gaia_assembler::backends::jvm::JvmBackend;
7pub use gaia_assembler::backends::wasi::WasiBackend;
8pub use gaia_assembler::backends::x86::X86Backend;
9use gaia_assembler::config::GaiaConfig;
10use gaia_assembler::instruction::{CoreInstruction, GaiaInstruction};
11use gaia_assembler::program::{
12    GaiaBlock, GaiaClass, GaiaConstant, GaiaFunction, GaiaModule, GaiaTerminator,
13};
14use gaia_assembler::types::{GaiaSignature, GaiaType};
15use std::collections::HashMap;
16use std::sync::Arc;
17
18/// Unified emitter for Gaia.
19/// Directly generates Gaia IR (GaiaModule) to avoid assembly parsing overhead.
20pub struct GaiaEmitter {
21    pub target_arch: String,
22    pub standalone: bool,
23    backends: HashMap<String, Arc<dyn GaiaBackend>>,
24}
25
26impl GaiaEmitter {
27    pub fn new(target_arch: &str) -> Self {
28        Self {
29            target_arch: target_arch.to_string(),
30            standalone: false,
31            backends: HashMap::new(),
32        }
33    }
34
35    pub fn with_backend(mut self, name: &str, backend: Arc<dyn GaiaBackend>) -> Self {
36        self.backends.insert(name.to_string(), backend);
37        self
38    }
39
40    pub fn emit(&self, tree: &IKunTree) -> ChomskyResult<BackendArtifact> {
41        let module = self.tree_to_ir(tree);
42
43        if let Some(backend) = self.backends.get(&self.target_arch) {
44            let config = GaiaConfig {
45                target: backend.primary_target(),
46                ..Default::default()
47            };
48
49            let result = backend.generate(&module, &config).map_err(|e| {
50                ChomskyError::new(chomsky_types::ChomskyErrorKind::BackendError {
51                    target: self.target_arch.clone(),
52                    stage: "Generation".to_string(),
53                    message: format!("{:?}", e),
54                })
55            })?;
56
57            // Take the first file as the main artifact for now
58            if let Some((_, bytes)) = result.files.into_iter().next() {
59                Ok(BackendArtifact::Binary(bytes))
60            } else {
61                Ok(BackendArtifact::Binary(vec![]))
62            }
63        } else {
64            // Fallback: just return the IR if no backend matches
65            Ok(BackendArtifact::Assembly(format!("{:#?}", module)))
66        }
67    }
68
69    pub fn standalone(mut self) -> Self {
70        self.standalone = true;
71        self
72    }
73
74    fn tree_to_ir(&self, tree: &IKunTree) -> GaiaModule {
75        let mut module = GaiaModule {
76            name: "jit_module".to_string(),
77            functions: Vec::new(),
78            structs: Vec::new(),
79            classes: Vec::new(),
80            constants: Vec::new(),
81            globals: Vec::new(),
82            imports: Vec::new(),
83        };
84
85        self.extract_module_elements(tree, &mut module);
86
87        // If no elements were extracted, wrap the whole tree in a default function
88        if module.functions.is_empty() && module.classes.is_empty() {
89            let mut body = Vec::new();
90            if !self.standalone {
91                self.push_jit_prologue(&mut body);
92            }
93            self.emit_tree_to_instructions(tree, &mut body);
94            if !self.standalone {
95                body.push(GaiaInstruction::Core(CoreInstruction::Ret));
96            }
97
98            module.functions.push(GaiaFunction {
99                name: "main".to_string(),
100                signature: GaiaSignature {
101                    params: Vec::new(),
102                    return_type: GaiaType::Void,
103                },
104                blocks: vec![GaiaBlock {
105                    label: "entry".to_string(),
106                    instructions: body,
107                    terminator: GaiaTerminator::Return,
108                }],
109                is_external: false,
110            });
111        }
112
113        module
114    }
115
116    fn extract_module_elements(&self, tree: &IKunTree, module: &mut GaiaModule) {
117        match tree {
118            IKunTree::Seq(items) => {
119                for item in items {
120                    self.extract_module_elements(item, module);
121                }
122            }
123            IKunTree::Extension(name, args) if name == "class" => {
124                if let (Some(IKunTree::StringConstant(name)), Some(IKunTree::Seq(members))) =
125                    (args.get(0), args.get(1))
126                {
127                    let mut class = GaiaClass {
128                        name: name.clone(),
129                        parent: Some("java/lang/Object".to_string()),
130                        interfaces: Vec::new(),
131                        fields: Vec::new(),
132                        methods: Vec::new(),
133                        attributes: Vec::new(),
134                    };
135
136                    for member in members {
137                        if let IKunTree::Extension(ext_name, ext_args) = member {
138                            if ext_name == "method" {
139                                if let (
140                                    Some(IKunTree::StringConstant(m_name)),
141                                    Some(IKunTree::StringConstant(_m_ret)),
142                                    Some(m_body),
143                                ) = (ext_args.get(0), ext_args.get(1), ext_args.get(2))
144                                {
145                                    let mut body = Vec::new();
146                                    self.emit_tree_to_instructions(m_body, &mut body);
147
148                                    let method = GaiaFunction {
149                                        name: m_name.clone(),
150                                        signature: GaiaSignature {
151                                            params: Vec::new(), // TODO: parse params
152                                            return_type: GaiaType::Void,
153                                        },
154                                        blocks: vec![GaiaBlock {
155                                            label: "entry".to_string(),
156                                            instructions: body,
157                                            terminator: GaiaTerminator::Return,
158                                        }],
159                                        is_external: false,
160                                    };
161                                    class.methods.push(method);
162                                }
163                            }
164                        }
165                    }
166                    module.classes.push(class);
167                }
168            }
169            _ => {}
170        }
171    }
172
173    fn push_jit_prologue(&self, body: &mut Vec<GaiaInstruction>) {
174        // Prologue for NyarVM JIT calling convention (Win64)
175        body.push(GaiaInstruction::Core(CoreInstruction::LoadArg(
176            0,
177            GaiaType::I64,
178        )));
179        body.push(GaiaInstruction::Core(CoreInstruction::StoreLocal(
180            0,
181            GaiaType::I64,
182        )));
183        body.push(GaiaInstruction::Core(CoreInstruction::LoadArg(
184            1,
185            GaiaType::I64,
186        )));
187        body.push(GaiaInstruction::Core(CoreInstruction::StoreLocal(
188            1,
189            GaiaType::I64,
190        )));
191        body.push(GaiaInstruction::Core(CoreInstruction::LoadArg(
192            2,
193            GaiaType::I64,
194        )));
195        body.push(GaiaInstruction::Core(CoreInstruction::StoreLocal(
196            2,
197            GaiaType::I64,
198        )));
199        body.push(GaiaInstruction::Core(CoreInstruction::LoadArg(
200            3,
201            GaiaType::I64,
202        )));
203        body.push(GaiaInstruction::Core(CoreInstruction::StoreLocal(
204            3,
205            GaiaType::I64,
206        )));
207    }
208
209    fn emit_tree_to_instructions(&self, tree: &IKunTree, body: &mut Vec<GaiaInstruction>) {
210        match tree {
211            IKunTree::Seq(items) => {
212                for item in items {
213                    self.emit_tree_to_instructions(item, body);
214                }
215            }
216            IKunTree::Extension(name, args) => match name.as_str() {
217                "ldc" => {
218                    if let Some(IKunTree::StringConstant(s)) = args.get(0) {
219                        body.push(GaiaInstruction::Core(CoreInstruction::PushConstant(
220                            GaiaConstant::String(s.clone()),
221                        )));
222                    }
223                }
224                "getstatic" => {
225                    if let (
226                        Some(IKunTree::StringConstant(cls)),
227                        Some(IKunTree::StringConstant(fld)),
228                    ) = (args.get(0), args.get(1))
229                    {
230                        body.push(GaiaInstruction::Core(CoreInstruction::LoadField(
231                            cls.clone(),
232                            fld.clone(),
233                        )));
234                    }
235                }
236                "invokevirtual" => {
237                    if let (
238                        Some(IKunTree::StringConstant(_cls)),
239                        Some(IKunTree::StringConstant(meth)),
240                    ) = (args.get(0), args.get(1))
241                    {
242                        body.push(GaiaInstruction::Core(CoreInstruction::Call(
243                            meth.clone(),
244                            1,
245                        )));
246                    }
247                }
248                _ => {}
249            },
250            _ => {}
251        }
252    }
253}
254
255impl Backend for GaiaEmitter {
256    fn name(&self) -> &str {
257        &self.target_arch
258    }
259
260    fn generate(&self, tree: &IKunTree) -> ChomskyResult<BackendArtifact> {
261        self.emit(tree)
262    }
263}