cells/src/main.rs

use std::collections::HashMap;
use std::num::NonZeroUsize;
use std::path::Path;
use std::process::Command;
use std::sync::Arc;

use parking_lot::Mutex;
use eframe::egui;

use nac3core::codegen;
use nac3core::inkwell;
use nac3core::nac3parser;
use nac3core::toplevel;
use nac3core::toplevel::composer;
use nac3core::typecheck::{type_inferencer, typedef};

mod basic_symbol_resolver;
use basic_symbol_resolver::{Resolver, ResolverInternal};

fn compile(code: &String, run_symbol: &String, output_filename: &Path) -> Result<(), String> {
    let mut target_machine_options = codegen::CodeGenTargetMachineOptions::from_host();
    target_machine_options.reloc_mode = inkwell::targets::RelocMode::PIC;
    let llvm_options = codegen::CodeGenLLVMOptions {
        opt_level: inkwell::OptimizationLevel::Default,
        target: target_machine_options,
    };
    let context = inkwell::context::Context::create();
    let target_machine = llvm_options
        .target
        .create_target_machine(llvm_options.opt_level)
        .expect("couldn't create target machine");
    let size_t = context
        .ptr_sized_int_type(&target_machine.get_target_data(), None)
        .get_bit_width();
    let primitive: type_inferencer::PrimitiveStore =
        composer::TopLevelComposer::make_primitives(size_t).0;
    let (mut composer, builtins_def, builtins_ty) = composer::TopLevelComposer::new(
        vec![],
        vec![],
        composer::ComposerConfig::default(),
        size_t,
    );
    let internal_resolver: Arc<ResolverInternal> = ResolverInternal {
        id_to_type: builtins_ty.into(),
        id_to_def: builtins_def.into(),
        module_globals: Mutex::default(),
        str_store: Mutex::default(),
    }
    .into();
    let resolver = Arc::new(Resolver(internal_resolver.clone()))
        as Arc<dyn nac3core::symbol_resolver::SymbolResolver + Send + Sync>;
    let irrt = codegen::irrt::load_irrt(&context, resolver.as_ref());

    let parser_result =
        match nac3parser::parser::parse_program(code.as_str(), String::from("cell1").into()) {
            Ok(parser_result) => parser_result,
            Err(err) => {
                return Err(format!("parse error: {}", err));
            }
        };
    for mut stmt in parser_result {
        if let nac3parser::ast::StmtKind::FunctionDef { name, .. } = &mut stmt.node {
            if name.to_string() == "run" {
                *name = run_symbol.as_str().into();
            }
        }
        match composer.register_top_level(stmt, Some(resolver.clone()), "__main__", true) {
            Ok((name, def_id, ty)) => {
                internal_resolver.add_id_def(name, def_id);
                if let Some(ty) = ty {
                    internal_resolver.add_id_type(name, ty);
                }
            }
            Err(err) => {
                return Err(format!("composer error: {}", err));
            }
        }
    }

    let signature = typedef::FunSignature {
        args: vec![],
        ret: primitive.int32,
        vars: typedef::VarMap::new(),
    };
    let mut store = codegen::concrete_type::ConcreteTypeStore::new();
    let mut cache = HashMap::new();
    let signature = store.from_signature(&mut composer.unifier, &primitive, &signature, &mut cache);
    let signature = store.add_cty(signature);

    if let Err(errors) = composer.start_analysis(true) {
        let error_count = errors.len();
        let mut msg = format!("{error_count} error(s) occurred during top level analysis.");

        for (error_i, error) in errors.iter().enumerate() {
            let error_num = error_i + 1;
            msg.push_str(&format!(
                "\n=========== ERROR {error_num}/{error_count} ============\n{error}"
            ));
        }
        return Err(msg);
    }

    let top_level = Arc::new(composer.make_top_level_context());

    let run_id_def = match resolver.get_identifier_def(run_symbol.as_str().into()) {
        Ok(run_id_def) => run_id_def,
        Err(_) => {
            return Err(format!("no run() entry point"));
        }
    };
    let instance = {
        let defs = top_level.definitions.read();
        let mut instance = defs[run_id_def.0].write();
        let toplevel::TopLevelDef::Function {
            instance_to_stmt,
            instance_to_symbol,
            ..
        } = &mut *instance
        else {
            unreachable!()
        };
        instance_to_symbol.insert(String::new(), run_symbol.clone());
        instance_to_stmt[""].clone()
    };

    let task = codegen::CodeGenTask {
        subst: Vec::default(),
        symbol_name: run_symbol.clone(),
        body: instance.body,
        signature,
        resolver,
        store,
        unifier_index: instance.unifier_id,
        calls: instance.calls,
        id: 0,
    };

    let nthreads = if inkwell::support::is_multithreaded() {
        std::thread::available_parallelism()
            .map(NonZeroUsize::get)
            .unwrap_or(1usize)
    } else {
        1
    };

    let membuffers: Arc<Mutex<Vec<Vec<u8>>>> = Arc::default();
    let membuffer = membuffers.clone();

    let f = Arc::new(codegen::WithCall::new(Box::new(move |module| {
        let buffer = module.write_bitcode_to_memory();
        let buffer = buffer.as_slice().into();
        membuffer.lock().push(buffer);
    })));
    let threads = (0..nthreads)
        .map(|i| {
            Box::new(codegen::DefaultCodeGenerator::new(
                format!("module{i}"),
                size_t,
            ))
        })
        .collect();
    let (registry, handles) =
        codegen::WorkerRegistry::create_workers(threads, top_level, &llvm_options, &f);
    registry.add_task(task);
    registry.wait_tasks_complete(handles);

    // Link all modules together into `main`
    let buffers = membuffers.lock();
    let main = context
        .create_module_from_ir(
            inkwell::memory_buffer::MemoryBuffer::create_from_memory_range(&buffers[0], "main"),
        )
        .unwrap();
    for buffer in buffers.iter().skip(1) {
        let other = context
            .create_module_from_ir(
                inkwell::memory_buffer::MemoryBuffer::create_from_memory_range(buffer, "main"),
            )
            .unwrap();
        main.link_in_module(other).unwrap();
    }
    main.link_in_module(irrt).unwrap();

    // Private all functions except "run"
    let mut function_iter = main.get_first_function();
    while let Some(func) = function_iter {
        if func.count_basic_blocks() > 0 && func.get_name().to_str().unwrap() != run_symbol {
            func.set_linkage(inkwell::module::Linkage::Private);
        }
        function_iter = func.get_next_function();
    }

    let pass_options = inkwell::passes::PassBuilderOptions::create();
    pass_options.set_merge_functions(true);
    main.run_passes("default<O2>", &target_machine, pass_options)
        .unwrap();

    target_machine
        .write_to_file(&main, inkwell::targets::FileType::Object, output_filename)
        .unwrap();
    Ok(())
}

// The year is 2024, and compiler toolchains are still a trash fire.
fn link_with_lld(elf_filename: &Path, obj_filename: &Path) -> Result<(), String> {
    let linker_args = vec![
        "-shared".to_string(),
        "--eh-frame-hdr".to_string(),
        "-x".to_string(),
        "-o".to_string(),
        elf_filename.to_str().unwrap().to_string(),
        obj_filename.to_str().unwrap().to_string(),
    ];

    if let Ok(linker_status) = Command::new("ld.lld").args(linker_args).status() {
        if !linker_status.success() {
            return Err("failed to start linker".to_string());
        }
    } else {
        return Err("linker returned non-zero status code".to_string());
    }

    Ok(())
}

type RunFn = unsafe extern "C" fn() -> i32;

struct CellBin {
    cell_id: usize,
    // note destructor order
    run_fn: Option<RunFn>,
    library: Option<libloading::Library>,
    directory: tempfile::TempDir,
}

impl CellBin {
    fn new(cell_id: usize) -> Self {
        Self {
            cell_id,
            run_fn: None,
            library: None,
            directory: tempfile::tempdir().unwrap(),
        }
    }

    fn compile_and_load(&mut self, code: &String) -> Result<(), String> {
        assert!(self.run_fn.is_none());
        assert!(self.library.is_none());
        let object = self.directory.path().join("module.o");
        let library = self.directory.path().join("module.so");
        let run_symbol = format!("__cells_run_{}", self.cell_id);
        compile(code, &run_symbol, &object)?;
        link_with_lld(&library, &object)?;
        unsafe {
            self.library = Some(libloading::Library::new(library).or_else(|e| Err(e.to_string()))?);
            let raw_fun_ptr = self
                .library
                .as_ref()
                .unwrap()
                .get::<RunFn>(run_symbol.as_bytes())
                .unwrap()
                .try_as_raw_ptr()
                .unwrap();
            self.run_fn = Some(std::mem::transmute(raw_fun_ptr));
        }
        Ok(())
    }
}

struct Cell {
    code: String,
    result: Result<CellBin, String>,
}

impl Cell {
    fn new() -> Self {
        Self {
            code: "".to_string(),
            result: Err("".to_string()),
        }
    }

    fn update(&mut self) {
        let mut new_bin = CellBin::new(0usize);
        self.result = new_bin
            .compile_and_load(&self.code)
            .and_then(|_| Ok(new_bin));
    }
}

struct Cells {
    cells: Vec<Cell>,
}

impl Cells {
    fn new() -> Self {
        Self {
            cells: vec![Cell::new()],
        }
    }

    fn update(&mut self) {
        for cell in self.cells.iter_mut() {
            cell.update()
        }
    }

    fn ensure_last(&mut self) {
        if self.cells.last().unwrap().result.is_ok() {
            self.cells.push(Cell::new());
        }
    }

    fn ui(&mut self, ui: &mut egui::Ui) {
        for (cell_i, cell) in self.cells.iter_mut().enumerate() {
            let theme = egui_extras::syntax_highlighting::CodeTheme::from_memory(ui.ctx());
            let mut layouter = |ui: &egui::Ui, string: &str, wrap_width: f32| {
                let mut layout_job =
                    egui_extras::syntax_highlighting::highlight(ui.ctx(), &theme, string, "Python");
                layout_job.wrap.max_width = wrap_width;
                layout_job.sections = layout_job
                    .sections
                    .iter()
                    .map(|layout_section| {
                        let mut section = layout_section.clone();
                        section.format.font_id = egui::FontId::monospace(16.0);
                        section
                    })
                    .collect();
                ui.fonts(|f| f.layout_job(layout_job))
            };
            ui.group(|ui| {
                ui.horizontal(|ui| {
                    ui.label(format!("[{}]", cell_i));
                    ui.add(
                        egui::TextEdit::multiline(&mut cell.code)
                            .code_editor()
                            .desired_rows(4)
                            .lock_focus(true)
                            .desired_width(f32::INFINITY)
                            .layouter(&mut layouter)
                            .font(egui::FontId::monospace(16.0)),
                    );
                });
                match &cell.result {
                    Ok(bin) => {
                        ui.label(format!("{}", unsafe { bin.run_fn.unwrap()() }));
                    }
                    Err(msg) => {
                        if !msg.is_empty() {
                            ui.colored_label(egui::Color32::from_rgb(255, 0, 0), msg);
                        }
                    }
                };
            });
        }
    }
}

fn main() -> eframe::Result {
    inkwell::targets::Target::initialize_native(&inkwell::targets::InitializationConfig::default())
        .unwrap();

    let options = eframe::NativeOptions {
        viewport: egui::ViewportBuilder::default().with_inner_size([1024.0, 768.0]),
        ..Default::default()
    };

    let mut cells = Cells::new();

    eframe::run_simple_native("Cells", options, move |ctx, _frame| {
        let submit_key = egui::KeyboardShortcut::new(egui::Modifiers::CTRL, egui::Key::Enter);
        if ctx.input_mut(|input| input.consume_shortcut(&submit_key)) {
            cells.update();
            cells.ensure_last();
        }

        egui::CentralPanel::default().show(ctx, |ui| {
            cells.ui(ui);
        });
    })
}

#[no_mangle]
pub extern "C" fn __nac3_personality(_state: u32, _exception_object: u32, _context: u32) -> u32 {
    unimplemented!()
}