regulus 0.0.14

mod storage;

use crate::interned_stdlib::INTERNED_STL;
use crate::no_path;
use crate::optimizations::run_optimizations;
use crate::parsing::{build_program, tokenize};
use crate::prelude::*;
use std::io::{BufRead, BufReader, Read, Write, stderr, stdin, stdout};
use std::path::{Path, PathBuf};
use std::rc::Rc;
use std::{env, fs, io};
pub use storage::Storage;

#[derive(Clone)]
pub(crate) enum Directory {
    Regular(PathBuf),
    FromEval,
    /// Should only be used internally.
    InternedSTL,
}

// TODO: add and update all docs here as well as on `Storage`.
/// The central structure for running a program.
pub struct State {
    /// All values that can be accessed during the program's execution.
    pub storage: Storage,
    /// Handle to the standard input. Defaults to [`std::io::stdin()`], but can be replaced.
    pub stdin: Box<dyn BufRead>,
    /// Handle to the standard output. Defaults to [`std::io::stdout()`], but can be replaced.
    pub stdout: WriteHandle,
    /// Handle to the standard error. Defaults to [`std::io::stderr()`], but can be replaced.
    pub stderr: WriteHandle,
    /// The directory (or pseudo-directory) in which the current program is placed.
    pub(crate) file_directory: Directory,
    pub(crate) current_file_path: Option<PathBuf>,
    pub(crate) exit_unwind_value: Option<Atom>,
    pub(crate) backtrace: Vec<Span>,
    // TODO: consider merging `current_doc_comment` and `current_fn_name`
    pub(crate) current_doc_comment: Option<String>,
    pub(crate) current_fn_name: Option<String>,
    /// Tracks the current stack of nested `import`-calls to emit an error on cyclic imports.
    /// Note that this only operates on user-written code and does not catch cyclic import
    /// errors within the STL (those still cause a rust stack overflow).
    pub(crate) import_stack: Vec<PathBuf>,
    code: Option<String>,
    next_type_id: i64,
    pub(crate) optimizations_enabled: bool,
    // make sure this type can never be constructed from outside
    __private: (),
}

impl Default for State {
    fn default() -> Self {
        Self::new()
    }
}

impl State {
    /// Creates a new state for running a program.
    ///
    /// You must use a method such as [`with_code`](Self::with_code) or [`with_source_file`](Self::with_source_file)
    /// to set the program source code before using [`run`](Self::run) to execute it.
    pub fn new() -> Self {
        Self {
            storage: Storage::initial(),
            stdin: Box::new(BufReader::new(stdin())),
            stdout: WriteHandle::new_write(stdout()),
            stderr: WriteHandle::new_write(stderr()),
            file_directory: Directory::InternedSTL,
            current_file_path: None,
            exit_unwind_value: None,
            backtrace: Vec::new(),
            current_doc_comment: None,
            current_fn_name: None,
            import_stack: Vec::new(),
            code: None,
            next_type_id: Atom::MIN_OBJECT_TY_ID,
            optimizations_enabled: false,
            __private: (),
        }
    }

    /// Sets the code that will be executed.
    #[must_use = "this returns the new state without modifying the original"]
    pub fn with_code(mut self, code: impl AsRef<str>) -> Self {
        self.code = Some(code.as_ref().to_owned());
        self
    }

    /// Sets both the code and the current directory by reading from the given file path.
    ///
    /// # Errors
    /// Returns an error if reading from the file failed.
    #[must_use = "this returns the new state without modifying the original"]
    pub fn with_source_file(mut self, path: impl AsRef<Path>) -> io::Result<Self> {
        self.code = Some(fs::read_to_string(&path)?);

        self.current_file_path = Some(path.as_ref().to_owned());

        let mut current_dir = path
            .as_ref()
            .parent()
            .ok_or_else(|| {
                io::Error::new(io::ErrorKind::InvalidFilename, "file path has no parent")
            })?
            .to_path_buf();

        if current_dir == PathBuf::new() {
            current_dir = PathBuf::from(".");
        }
        self.file_directory = Directory::Regular(current_dir);
        Ok(self)
    }

    /// Sets the source directory for resolving imports to the given directory.
    ///
    /// Note that this does not set or change the program code.
    #[must_use = "this returns the new state without modifying the original"]
    pub fn with_source_directory(mut self, dir_path: impl AsRef<Path>) -> Self {
        self.file_directory = Directory::Regular(dir_path.as_ref().to_path_buf());
        self
    }

    /// Enables optimizations which run prior to execution.
    ///
    /// **WARNING**: Enabling optimizations may cause programs that redefine builtins
    /// (or possibly also STL functions) to change their behavior.
    #[must_use = "this returns the new state without modifying the original"]
    pub fn enable_optimizations(mut self) -> Self {
        self.optimizations_enabled = true;
        self
    }

    /// Sets the current directory to the operating systems current working directory.
    ///
    /// Note that this does not set or change the program code.
    ///
    /// # Panics
    /// Panics if [`env::current_dir`] returned an error.
    #[must_use = "this returns the new state without modifying the original"]
    pub fn with_cwd(self) -> Self {
        self.with_source_directory(env::current_dir().unwrap())
    }

    /// Runs the given program with the details specified by this state.
    ///
    /// Returns the result the program returned.
    ///
    /// # Panics
    /// Panics if `code` was not set.
    pub fn run(&mut self) -> Result<Atom> {
        self.exit_unwind_value = None;

        // newlines are needed to avoid interaction with comments
        // and also help with calculating the actual spans (just do line - 1)
        let code = format!(
            "_(\n{}\n)",
            self.code
                .as_ref()
                .expect("setting the source code is required")
        );

        let file_path = if let Some(path) = &self.current_file_path {
            Rc::new(path.clone())
        } else {
            no_path()
        };

        let tokens = tokenize(&code, file_path)?;

        let mut program = build_program(tokens)?;
        if self.optimizations_enabled {
            run_optimizations(&mut program);
        }

        self.import_prelude();
        let result = program.eval(self)?;

        if let Some(exit_unwind_value) = &self.exit_unwind_value {
            return Ok(exit_unwind_value.clone());
        }

        Ok(result)
    }

    fn import_prelude(&mut self) {
        if matches!(self.file_directory, Directory::InternedSTL) {
            return;
        }
        let mut import_state = Self::new();
        let code = INTERNED_STL
            .get("prelude")
            .expect("`prelude.re` missing from STL");
        import_state = import_state.with_code(code);
        import_state.set_current_file_path("<stl:prelude>");
        import_state.optimizations_enabled = self.optimizations_enabled;
        import_state.run().expect("prelude import failed");

        self.storage.extend_from(import_state.storage, None);
    }

    /// Writes the given string to stdout, without any extra newline.
    pub(crate) fn write_to_stdout(&mut self, msg: &str) {
        self.stdout.as_write().write_all(msg.as_bytes()).unwrap();
    }

    /// Writes the given string to stdout, without any extra newline.
    pub(crate) fn write_to_stderr(&mut self, msg: &str) {
        self.stderr.as_write().write_all(msg.as_bytes()).unwrap();
    }

    /// Only intended to be used by `import` internals for now.
    pub(crate) fn set_current_file_path(&mut self, path: impl AsRef<Path>) {
        self.current_file_path = Some(path.as_ref().to_owned());
    }

    /// Returns a new type id for a `type` call.
    pub fn make_type_id(&mut self) -> i64 {
        let old = self.next_type_id;
        self.next_type_id += 1;
        old
    }

    /// Constructs a new exception with the given error and message at the current point of execution.
    pub fn raise(&self, error: impl Into<String>, msg: impl Into<String>) -> Exception {
        Exception::with_trace(error, msg, &self.backtrace)
    }

    pub(crate) fn get_function(&self, name: &str) -> Result<Function> {
        match self.storage.get(name) {
            Some(atom) => {
                if let Atom::Function(func) = atom {
                    Ok(func.clone())
                } else {
                    raise!(self, "Name", "`{name}` is not a function")
                }
            }
            None => raise!(self, "Name", "no function `{name}` found"),
        }
    }
}

/// Helper trait for types that can both be read from and written to.
pub trait ReadAndWrite: Read + Write {}

impl<T> ReadAndWrite for T where T: Read + Write {}

/// A handle that always allows writing and optionally allows reading. Used for stdout and stderr.
///
/// Usually, writing to stdout / stderr is enough, but sometimes, one may want to capture
/// what is written to them. In this case, the `ReadWrite` variant can be used.
pub enum WriteHandle {
    ReadWrite(Box<dyn ReadAndWrite>),
    Write(Box<dyn Write>),
}

impl WriteHandle {
    /// Access the `Write` part of this handle.
    pub fn as_write(&mut self) -> &mut dyn Write {
        match self {
            Self::Write(w) => w,
            Self::ReadWrite(w) => w,
        }
    }

    /// Constructs a new `WriteHandle` that only allows writing.
    pub fn new_write(write: impl Write + 'static) -> Self {
        Self::Write(Box::new(write))
    }

    /// Access the `Read` part of this handle or `None` if the handle does not allow reading.
    pub fn as_read(&mut self) -> Option<&mut dyn Read> {
        match self {
            Self::ReadWrite(w) => Some(w),
            Self::Write(_) => None,
        }
    }

    /// Constructs a new `WriteHandle` that allows both writing and reading.
    pub fn new_read_write(read_write: impl ReadAndWrite + 'static) -> Self {
        Self::ReadWrite(Box::new(read_write))
    }

    /// Return a string representation of the data in this handle if it allows reading.
    ///
    /// # Panics
    /// Panics if it is does not allow reading or if it does not contain valid UTF-8.
    pub fn read_to_string(&mut self) -> String {
        let Self::ReadWrite(buf) = self else {
            panic!("read_to_string(): cannot read from write only handle")
        };
        let mut vec = vec![0; 1024];
        let mut n = 0;
        let mut last_was_zero = false;
        loop {
            let amount = buf.read(&mut vec[n..]).unwrap();
            n += amount;
            if amount == 0 {
                if last_was_zero {
                    break;
                }
                vec.extend([0; 1024]);
                last_was_zero = true;
            } else {
                last_was_zero = false;
            }
        }
        vec.retain(|&x| x != 0);
        String::from_utf8(vec).unwrap()
    }
}