ngr/src/repl.rs

use crate::backend::{Backend, BackendError};
use crate::ir::TypeInferenceResult;
use crate::syntax::{ConstantType, Location, ParserError, Statement};
use codespan_reporting::diagnostic::Diagnostic;
use codespan_reporting::files::SimpleFiles;
use codespan_reporting::term::{self, Config};
use cranelift_jit::JITModule;
use cranelift_module::ModuleError;
use pretty::termcolor::{ColorChoice, StandardStream};
use std::collections::HashMap;

/// A high-level REPL helper for NGR.
///
/// This object holds most of the state required to implement some
/// form of interactive compiler for NGR; all you need to do is provide
/// the actual user IO.
///
/// For most console-based used cases, the [`Default`] implementation
/// should be sufficient; it prints any warnings or errors to `stdout`,
/// using a default color scheme that should work based on the terminal
/// type. For more complex interactions, though, you may want to use
/// the `REPL::new` function to provide your own print substrate.
pub struct REPL {
    file_database: SimpleFiles<String, String>,
    jitter: Backend<JITModule>,
    variable_binding_sites: HashMap<String, Location>,
    console: StandardStream,
    console_config: Config,
}

impl Default for REPL {
    fn default() -> Self {
        let console = StandardStream::stdout(ColorChoice::Auto);
        REPL::new(console, Config::default()).unwrap()
    }
}

#[allow(clippy::upper_case_acronyms)]
#[derive(Debug, thiserror::Error)]
enum REPLError {
    #[error("Error parsing statement: {0}")]
    Parser(#[from] ParserError),
    #[error("JIT error: {0}")]
    JIT(#[from] BackendError),
    #[error("Internal cranelift error: {0}")]
    Cranelift(#[from] ModuleError),
    #[error(transparent)]
    Reporting(#[from] codespan_reporting::files::Error),
}

impl From<REPLError> for Diagnostic<usize> {
    fn from(value: REPLError) -> Self {
        match value {
            REPLError::Parser(err) => Diagnostic::from(&err),
            REPLError::JIT(err) => Diagnostic::from(err),
            REPLError::Cranelift(err) => Diagnostic::bug().with_message(format!("{}", err)),
            REPLError::Reporting(err) => Diagnostic::bug().with_message(format!("{}", err)),
        }
    }
}

impl REPL {
    /// Construct a new REPL helper, using the given stream implementation and console configuration.
    ///
    /// For most users, the [`Default::default`] implementation will be sufficient;
    /// it will use `stdout` and a default console configuration. But if you need to
    /// be more specific, this will help you provide more guidance to the REPL as it
    /// evaluates things.
    pub fn new(console: StandardStream, console_config: Config) -> Result<Self, BackendError> {
        Ok(REPL {
            file_database: SimpleFiles::new(),
            jitter: Backend::jit(None)?,
            variable_binding_sites: HashMap::new(),
            console,
            console_config,
        })
    }

    /// Emit a diagnostic to the configured console.
    ///
    /// This is just a convenience function; there's a lot of boilerplate in printing
    /// diagnostics, and it was nice to pull it out into its own function.
    fn emit_diagnostic(
        &mut self,
        diagnostic: Diagnostic<usize>,
    ) -> Result<(), codespan_reporting::files::Error> {
        term::emit(
            &mut self.console,
            &self.console_config,
            &self.file_database,
            &diagnostic,
        )
    }

    /// Process a line of input, printing any problems or the results.
    ///
    /// The line number argument is just for a modicum of source information, to
    /// provide to the user if some parsing or validation step fails. It can be
    /// changed to be any value you like that provides some insight into what
    /// failed, although it is probably a good idea for it to be different for
    /// every invocation of this function. (Not critical, but a good idea.)
    ///
    /// Any warnings or errors generated in processing this command will be
    /// printed to the configured console. If there are no problems, the
    /// command will be compiled and then executed.
    pub fn process_input(&mut self, line_no: usize, command: String) {
        if let Err(err) = self.process(line_no, command) {
            if let Err(e) = self.emit_diagnostic(Diagnostic::from(err)) {
                eprintln!(
                    "WOAH! System having trouble printing error messages. This is very bad. ({})",
                    e
                );
            }
        }
    }

    /// The internal implementation, with a handy `Result` type.
    ///
    /// All information from the documentation of `REPL::process_input` applies here,
    /// as well; this is the internal implementation of that function, which is
    /// differentiated by returning a `Result` type that is hidden from the user
    /// in the case of `REPL::process_input`.
    fn process(&mut self, line_no: usize, command: String) -> Result<(), REPLError> {
        let entry = self.file_database.add("entry".to_string(), command);
        let source = self
            .file_database
            .get(entry)
            .expect("entry exists")
            .source();
        let syntax = Statement::parse(entry, source)?;

        let program = match syntax {
            Statement::Binding(loc, name, expr) => {
                // if this is a variable binding, and we've never defined this variable before,
                // we should tell cranelift about it. this is optimistic; if we fail to compile,
                // then we won't use this definition until someone tries again.
                if !self.variable_binding_sites.contains_key(&name) {
                    self.jitter.define_string(&name)?;
                    self.jitter
                        .define_variable(name.clone(), ConstantType::U64)?;
                }

                crate::syntax::Program {
                    statements: vec![
                        Statement::Binding(loc.clone(), name.clone(), expr),
                        Statement::Print(loc, name),
                    ],
                }
            }

            nonbinding => crate::syntax::Program {
                statements: vec![nonbinding],
            },
        };

        let (mut errors, mut warnings) =
            program.validate_with_bindings(&mut self.variable_binding_sites);
        let stop = !errors.is_empty();
        let messages = errors
            .drain(..)
            .map(Into::into)
            .chain(warnings.drain(..).map(Into::into));

        for message in messages {
            self.emit_diagnostic(message)?;
        }

        if stop {
            return Ok(());
        }

        match program.type_infer() {
            TypeInferenceResult::Failure {
                mut errors,
                mut warnings,
            } => {
                let messages = errors
                    .drain(..)
                    .map(Into::into)
                    .chain(warnings.drain(..).map(Into::into));

                for message in messages {
                    self.emit_diagnostic(message)?;
                }

                Ok(())
            }

            TypeInferenceResult::Success {
                result,
                mut warnings,
            } => {
                for message in warnings.drain(..).map(Into::into) {
                    self.emit_diagnostic(message)?;
                }
                let name = format!("line{}", line_no);
                let function_id = self.jitter.compile_function(&name, result)?;
                self.jitter.module.finalize_definitions()?;
                let compiled_bytes = self.jitter.bytes(function_id);
                let compiled_function =
                    unsafe { std::mem::transmute::<_, fn() -> ()>(compiled_bytes) };
                compiled_function();
                Ok(())
            }
        }
    }
}