driver.rs

use std::fs;
use std::path::Path;

use crate::diag::{DiagCode, DiagEngine, Diagnostic, Label};
use crate::interp::{
    ExecutionTrace, RuntimeError, RuntimeLimits, StackFrame, render_trace, run_main_with_limits,
    run_main_with_trace,
};
use crate::ir::{lower_program, render_program};
use crate::lex::{lex_file, render_tokens};
use crate::parse::parse;
use crate::parse::render_ast;
use crate::sema::{analyze, render_analysis};
use crate::source::SourceManager;

const EXIT_COMPILE_FAILURE: i32 = 2;
const EXIT_RUNTIME_FAILURE: i32 = 3;
const EXIT_INTERNAL_FAILURE: i32 = 4;
const INTERNAL_ERROR: DiagCode = DiagCode::new(9001);

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DriverOutput {
    pub exit_code: i32,
    pub stdout: String,
    pub stderr: String,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DebugMode {
    DumpTokens,
    DumpAst,
    DumpSema,
    DumpIr,
    TraceExec,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct DriverOptions {
    pub debug_mode: Option<DebugMode>,
    pub limits: RuntimeLimits,
    pub no_color: bool,
}

impl Default for DriverOptions {
    fn default() -> Self {
        Self {
            debug_mode: None,
            limits: RuntimeLimits::default(),
            no_color: false,
        }
    }
}

pub fn run_path(path: &Path, limits: RuntimeLimits) -> DriverOutput {
    run_path_with_options(
        path,
        DriverOptions {
            limits,
            ..DriverOptions::default()
        },
    )
}

pub fn run_path_with_options(path: &Path, options: DriverOptions) -> DriverOutput {
    match fs::read_to_string(path) {
        Ok(source) => run_source_with_options(path.to_string_lossy().into_owned(), source, options),
        Err(err) => DriverOutput {
            exit_code: EXIT_INTERNAL_FAILURE,
            stdout: String::new(),
            stderr: format!(
                "error[{INTERNAL_ERROR}]: failed to read '{}': {err}\n",
                path.display()
            ),
        },
    }
}

pub fn run_source(name: String, source: String, limits: RuntimeLimits) -> DriverOutput {
    run_source_with_options(
        name,
        source,
        DriverOptions {
            limits,
            ..DriverOptions::default()
        },
    )
}

pub fn run_source_with_options(
    name: String,
    source: String,
    options: DriverOptions,
) -> DriverOutput {
    let mut sm = SourceManager::new();
    let file = sm.add_file(name, source);
    let mut diags = DiagEngine::new();

    if options.debug_mode == Some(DebugMode::DumpTokens) {
        let tokens = lex_file(&sm, file, &mut diags);
        if diags.has_errors() {
            return compile_failure(&sm, &diags, options.no_color);
        }
        return DriverOutput {
            exit_code: 0,
            stdout: render_tokens(&tokens),
            stderr: String::new(),
        };
    }

    let parsed = parse(&sm, file, &mut diags);

    if diags.has_errors() {
        return compile_failure(&sm, &diags, options.no_color);
    }

    if options.debug_mode == Some(DebugMode::DumpAst) {
        return DriverOutput {
            exit_code: 0,
            stdout: render_ast(&parsed),
            stderr: String::new(),
        };
    }

    let analyzed = analyze(&parsed, &mut diags).expect("analysis should produce a program");
    if diags.has_errors() {
        return compile_failure(&sm, &diags, options.no_color);
    }

    if options.debug_mode == Some(DebugMode::DumpSema) {
        return DriverOutput {
            exit_code: 0,
            stdout: format!("analysis:\n{}", render_analysis(&analyzed)),
            stderr: String::new(),
        };
    }

    let program = lower_program(&analyzed);
    if options.debug_mode == Some(DebugMode::DumpIr) {
        return DriverOutput {
            exit_code: 0,
            stdout: format!("ir:\n{}", render_program(&program)),
            stderr: String::new(),
        };
    }

    let mut stdout = String::new();

    let mut trace = ExecutionTrace::default();
    let execution = if options.debug_mode == Some(DebugMode::TraceExec) {
        run_main_with_trace(&program, &mut stdout, options.limits, &mut trace)
    } else {
        run_main_with_limits(&program, &mut stdout, options.limits)
    };

    match execution {
        Ok(result) => DriverOutput {
            exit_code: result.exit_code,
            stdout,
            stderr: if options.debug_mode == Some(DebugMode::TraceExec) {
                render_trace(&sm, &trace)
            } else {
                String::new()
            },
        },
        Err(error) => DriverOutput {
            exit_code: classify_runtime_error(&error),
            stdout,
            stderr: render_runtime_failure(
                &sm,
                &error,
                options.debug_mode,
                &trace,
                options.no_color,
            ),
        },
    }
}

fn classify_runtime_error(error: &RuntimeError) -> i32 {
    if error.code.0 >= 9000 {
        EXIT_INTERNAL_FAILURE
    } else {
        EXIT_RUNTIME_FAILURE
    }
}

fn render_runtime_error(sm: &SourceManager, error: &RuntimeError, no_color: bool) -> String {
    let mut engine = DiagEngine::new();
    engine.emit(
        Diagnostic::error(error.code, &error.message)
            .with_label(Label::primary(error.span, "here")),
    );

    let mut rendered = if no_color {
        engine.render_to_string(sm)
    } else {
        engine.render_to_ansi_string(sm)
    };
    if !error.trace.is_empty() {
        rendered.push_str("stack trace:\n");
        rendered.push_str(&render_stack_trace(sm, &error.trace));
    }
    rendered
}

fn render_runtime_failure(
    sm: &SourceManager,
    error: &RuntimeError,
    debug_mode: Option<DebugMode>,
    trace: &ExecutionTrace,
    no_color: bool,
) -> String {
    let mut rendered = render_runtime_error(sm, error, no_color);
    if debug_mode == Some(DebugMode::TraceExec) {
        if !rendered.ends_with('\n') {
            rendered.push('\n');
        }
        rendered.push_str(&render_trace(sm, trace));
    }
    rendered
}

fn compile_failure(sm: &SourceManager, diags: &DiagEngine, no_color: bool) -> DriverOutput {
    let stderr = if no_color {
        diags.render_to_string(sm)
    } else {
        diags.render_to_ansi_string(sm)
    };
    DriverOutput {
        exit_code: EXIT_COMPILE_FAILURE,
        stdout: String::new(),
        stderr,
    }
}

fn render_stack_trace(sm: &SourceManager, trace: &[StackFrame]) -> String {
    const HEAD_FRAMES: usize = 6;
    const TAIL_FRAMES: usize = 1;

    let mut rendered = String::new();
    if trace.len() <= HEAD_FRAMES + TAIL_FRAMES + 1 {
        for (index, frame) in trace.iter().enumerate() {
            rendered.push_str(&format_stack_frame(sm, index, frame));
        }
        return rendered;
    }

    for (index, frame) in trace.iter().take(HEAD_FRAMES).enumerate() {
        rendered.push_str(&format_stack_frame(sm, index, frame));
    }

    let omitted = trace.len() - HEAD_FRAMES - TAIL_FRAMES;
    rendered.push_str(&format!("  ... {omitted} more frame(s) omitted ...\n"));

    let last_index = trace.len() - 1;
    rendered.push_str(&format_stack_frame(sm, last_index, &trace[last_index]));
    rendered
}

fn format_stack_frame(sm: &SourceManager, index: usize, frame: &StackFrame) -> String {
    match frame.call_span {
        Some(span) => {
            let file = sm.file(span.file);
            let (line, column) = file.location(span.start);
            format!(
                "  {index}: {} at {}:{}:{}\n",
                frame.function_name,
                file.name(),
                line,
                column
            )
        }
        None => format!("  {index}: {}\n", frame.function_name),
    }
}