cells/src/main.rs

use std::collections::HashMap;
use std::num::NonZeroUsize;
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 run(code: &String) {
    let llvm_options = codegen::CodeGenLLVMOptions {
        opt_level: inkwell::OptimizationLevel::Default,
        target: codegen::CodeGenTargetMachineOptions::from_host(),
    };
    let context = inkwell::context::Context::create();
    let size_t = context
        .ptr_sized_int_type(
            &llvm_options
                .target
                .create_target_machine(llvm_options.opt_level)
                .map(|tm| tm.get_target_data())
                .unwrap(),
            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) => {
                eprintln!("parse error: {}", err);
                return;
            }
        };
    for stmt in parser_result {
        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) => {
                eprintln!("composer error: {}", err);
                return;
            }
        }
    }

    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();
        eprintln!("{error_count} error(s) occurred during top level analysis.");

        for (error_i, error) in errors.iter().enumerate() {
            let error_num = error_i + 1;
            eprintln!("=========== ERROR {error_num}/{error_count} ============");
            eprintln!("{error}");
        }
        eprintln!("==================================");
        return;
    }

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

    let run_id_def = match resolver.get_identifier_def("run".into()) {
        Ok(run_id_def) => run_id_def,
        Err(_) => {
            eprintln!("no run() entry point");
            return;
        }
    };
    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".to_string());
        instance_to_stmt[""].clone()
    };

    let task = codegen::CodeGenTask {
        subst: Vec::default(),
        symbol_name: "run".to_string(),
        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();

    let execution_engine = main
        .create_jit_execution_engine(llvm_options.opt_level)
        .unwrap();
    type Run = unsafe extern "C" fn() -> i32;
    let run: inkwell::execution_engine::JitFunction<Run> =
        unsafe { execution_engine.get_function("run").unwrap() };
    println!("{}", unsafe { run.call() });
}

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

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

    let mut code = String::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(|i| i.consume_shortcut(&submit_key)) {
            run(&code);
        }

        egui::CentralPanel::default().show(ctx, |ui| {
            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.add(
                egui::TextEdit::multiline(&mut code)
                    .code_editor()
                    .desired_rows(4)
                    .lock_focus(true)
                    .desired_width(f32::INFINITY)
                    .layouter(&mut layouter)
                    .font(egui::FontId::monospace(16.0)),
            );
        });
    })
}