From 7d53706e712b3e037b7fc175e6edbca940957145 Mon Sep 17 00:00:00 2001
From: Jakub Doka <jakub.doka2@gmail.com>
Date: Tue, 22 Oct 2024 12:50:54 +0200
Subject: [PATCH] adding --optimize flag to the compiler

---
 lang/command-help.txt |  1 +
 lang/src/fs.rs        | 16 ++++++++++++++++
 lang/src/son.rs       | 16 +++++++++++++---
 3 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/lang/command-help.txt b/lang/command-help.txt
index 57b4191..bb92479 100644
--- a/lang/command-help.txt
+++ b/lang/command-help.txt
@@ -2,3 +2,4 @@
 --fmt-stdout - dont write the formatted file but print it
 --dump-asm - output assembly instead of raw code, (the assembly is more for debugging the compiler)
 --threads <1...> - number of threads compiler can use [default: 1]
+--optimize - compile with optimizing compiler (dangerose)
diff --git a/lang/src/fs.rs b/lang/src/fs.rs
index a128bec..7b973bc 100644
--- a/lang/src/fs.rs
+++ b/lang/src/fs.rs
@@ -2,6 +2,7 @@ use {
     crate::{
         codegen,
         parser::{self, Ast, FileKind, ParserCtx},
+        son,
     },
     alloc::{string::String, vec::Vec},
     core::{fmt::Write, num::NonZeroUsize},
@@ -38,6 +39,7 @@ pub struct Options {
     pub fmt: bool,
     pub fmt_stdout: bool,
     pub dump_asm: bool,
+    pub optimize: bool,
     pub extra_threads: usize,
 }
 
@@ -51,6 +53,7 @@ impl Options {
 
         Ok(Options {
             fmt: args.contains(&"--fmt"),
+            optimize: args.contains(&"--optimize"),
             fmt_stdout: args.contains(&"--fmt-stdout"),
             dump_asm: args.contains(&"--dump-asm"),
             extra_threads: args
@@ -85,6 +88,19 @@ pub fn run_compiler(root_file: &str, options: Options, out: &mut Vec<u8>) -> std
     } else if options.fmt_stdout {
         let ast = parsed.ast.into_iter().next().unwrap();
         write!(out, "{ast}").unwrap();
+    } else if options.optimize {
+        let mut codegen = son::Codegen::default();
+        codegen.files = &parsed.ast;
+        codegen.push_embeds(parsed.embeds);
+
+        codegen.generate();
+        if options.dump_asm {
+            codegen
+                .disasm(unsafe { std::mem::transmute::<&mut Vec<u8>, &mut String>(out) })
+                .map_err(|e| io::Error::other(e.to_string()))?;
+        } else {
+            codegen.assemble(out);
+        }
     } else {
         let mut codegen = codegen::Codegen::default();
         codegen.files = parsed.ast;
diff --git a/lang/src/son.rs b/lang/src/son.rs
index 423546a..2661d79 100644
--- a/lang/src/son.rs
+++ b/lang/src/son.rs
@@ -22,9 +22,9 @@ use {
         fmt::{self, Debug, Display, Write},
         format_args as fa, mem,
         ops::{self},
-        usize,
     },
     hashbrown::hash_map,
+    hbbytecode::DisasmError,
     regalloc2::VReg,
     std::{borrow::ToOwned, panic},
 };
@@ -1822,12 +1822,22 @@ impl<'a> Codegen<'a> {
         load
     }
 
-    pub fn generate(&mut self) {
-        self.find_type(0, 0, Err("main"), self.files);
+    pub fn generate(&mut self, entry: FileId) {
+        self.find_type(0, entry, Err("main"), self.files);
         self.make_func_reachable(0);
         self.complete_call_graph();
     }
 
+    pub fn assemble(&mut self, buf: &mut Vec<u8>) {
+        self.tys.reassemble(buf);
+    }
+
+    pub fn disasm(&mut self, output: &mut String) -> Result<(), DisasmError> {
+        let mut bin = Vec::new();
+        self.assemble(&mut bin);
+        self.tys.disasm(&bin, self.files, output, |_| {})
+    }
+
     fn make_func_reachable(&mut self, func: ty::Func) {
         let fuc = &mut self.tys.ins.funcs[func as usize];
         if fuc.offset == u32::MAX {