rkgk/crates/haku/src/bytecode.rs

use core::fmt::{self, Display};

use alloc::{borrow::ToOwned, string::String, vec::Vec};

use crate::source::Span;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[repr(u8)]
pub enum Opcode {
    // Push literal values onto the stack.
    Nil,
    False,
    True,
    Tag,
    Number, // (float: f32)
    Rgba,   // (r: u8, g: u8, b: u8, a: u8)

    // Duplicate existing values.
    /// Push a value relative to the bottom of the current stack window.
    Local, // (index: u8)
    /// Set the value of a value relative to the bottom of the current stack window.
    SetLocal, // (index: u8)
    /// Push a captured value.
    Capture, // (index: u8)
    /// Get the value of a definition.
    Def, // (index: u16)
    /// Set the value of a definition.
    SetDef, // (index: u16)

    // Create lists.
    List, // (len: u16)

    // Create literal functions.
    Function, // (params: u8, then: u16), at `then`: (local_count: u8, capture_count: u8, captures: [(source: u8, index: u8); capture_count])

    // Control flow.
    Jump,      // (offset: u16)
    JumpIfNot, // (offset: u16)
    Field,     // (count: u8, tags: [u16; count])

    // Function calls.
    Call, // (argc: u8)
    /// This is a fast path for system calls, which are quite common (e.g. basic arithmetic.)
    System, // (index: u8, argc: u8)

    Return,
    // NOTE: There must be no more opcodes after this.
    // They will get treated as invalid.
}

// Constants used by the Function opcode to indicate capture sources.
pub const CAPTURE_LOCAL: u8 = 0;
pub const CAPTURE_CAPTURE: u8 = 1;

#[derive(Debug, Clone)]
pub struct Chunk {
    pub bytecode: Vec<u8>,
    pub span_info: Vec<SpanRun>,
    pub current_span: Span,
}

#[derive(Debug, Clone)]
pub struct SpanRun {
    pub span: Span,
    pub len: u16,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Offset(u16);

impl Chunk {
    pub fn new(capacity: usize) -> Result<Chunk, ChunkSizeError> {
        if capacity <= (1 << 16) {
            Ok(Chunk {
                bytecode: Vec::with_capacity(capacity),
                span_info: Vec::new(),
                current_span: Span::new(0, 0),
            })
        } else {
            Err(ChunkSizeError)
        }
    }

    pub fn offset(&self) -> Offset {
        Offset(self.bytecode.len() as u16)
    }

    fn push_span_info(&mut self, span: Span, len: u16) {
        if let Some(info) = self.span_info.last_mut() {
            if info.span == span {
                info.len += len;
                return;
            }
        }
        self.span_info.push(SpanRun { span, len });
    }

    pub fn emit_bytes(&mut self, bytes: &[u8]) -> Result<Offset, EmitError> {
        if self.bytecode.len() + bytes.len() > self.bytecode.capacity() {
            return Err(EmitError);
        }

        let offset = Offset(self.bytecode.len() as u16);
        self.bytecode.extend_from_slice(bytes);
        self.push_span_info(self.current_span, bytes.len() as u16);

        Ok(offset)
    }

    pub fn emit_opcode(&mut self, opcode: Opcode) -> Result<Offset, EmitError> {
        self.emit_bytes(&[opcode as u8])
    }

    pub fn emit_u8(&mut self, x: u8) -> Result<Offset, EmitError> {
        self.emit_bytes(&[x])
    }

    pub fn emit_u16(&mut self, x: u16) -> Result<Offset, EmitError> {
        self.emit_bytes(&x.to_le_bytes())
    }

    pub fn emit_u32(&mut self, x: u32) -> Result<Offset, EmitError> {
        self.emit_bytes(&x.to_le_bytes())
    }

    pub fn emit_f32(&mut self, x: f32) -> Result<Offset, EmitError> {
        self.emit_bytes(&x.to_le_bytes())
    }

    pub fn patch_u8(&mut self, offset: Offset, x: u8) {
        self.bytecode[offset.0 as usize] = x;
    }

    pub fn patch_u16(&mut self, offset: Offset, x: u16) {
        let b = x.to_le_bytes();
        let i = offset.0 as usize;
        self.bytecode[i] = b[0];
        self.bytecode[i + 1] = b[1];
    }

    pub fn patch_offset(&mut self, offset: Offset, x: Offset) {
        self.patch_u16(offset, x.0);
    }

    pub fn find_span(&self, pc: u16) -> Option<&Span> {
        let mut cur = 0;
        for info in &self.span_info {
            if pc >= cur && pc < cur + info.len {
                return Some(&info.span);
            }
            cur += info.len;
        }
        None
    }
}

impl Offset {
    pub fn to_u16(self) -> u16 {
        self.0
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ChunkSizeError;

impl Display for ChunkSizeError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "chunk size must be less than 64 KiB")
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct EmitError;

impl Display for EmitError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "out of space in chunk")
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
pub struct DefId(u16);

impl DefId {
    pub fn to_u16(self) -> u16 {
        self.0
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
pub struct TagId(u16);

impl TagId {
    pub(crate) fn from_u16(x: u16) -> Self {
        Self(x)
    }

    pub fn to_u16(self) -> u16 {
        self.0
    }
}

#[derive(Debug, Clone)]
pub struct Defs {
    defs: Vec<String>,
    tags: Vec<String>,
}

#[derive(Debug, Clone, Copy)]
pub struct DefsLimits {
    pub max_defs: usize,
    pub max_tags: usize,
}

#[derive(Debug, Clone, Copy)]
pub struct DefsImage {
    defs: usize,
    tags: usize,
}

impl Defs {
    pub fn new(limits: &DefsLimits) -> Self {
        assert!(limits.max_defs < u16::MAX as usize + 1);
        assert!(limits.max_tags < u16::MAX as usize + 1);

        let mut tags = Vec::with_capacity(limits.max_tags);
        add_well_known_tags(&mut tags);

        Self {
            defs: Vec::with_capacity(limits.max_defs),
            tags,
        }
    }

    pub fn len(&self) -> u16 {
        self.defs.len() as u16
    }

    pub fn is_empty(&self) -> bool {
        self.len() != 0
    }

    pub fn get_def(&mut self, name: &str) -> Option<DefId> {
        self.defs
            .iter()
            .position(|n| *n == name)
            .map(|index| DefId(index as u16))
    }

    pub fn add_def(&mut self, name: &str) -> Result<DefId, DefError> {
        if self.defs.iter().any(|n| n == name) {
            Err(DefError::Exists)
        } else {
            if self.defs.len() >= self.defs.capacity() {
                return Err(DefError::OutOfSpace);
            }
            let id = DefId(self.defs.len() as u16);
            self.defs.push(name.to_owned());
            Ok(id)
        }
    }

    fn add_tag(tags: &mut Vec<String>, name: &str) -> Result<TagId, TagError> {
        if tags.len() >= tags.capacity() {
            return Err(TagError::OutOfSpace);
        }
        let id = TagId(tags.len() as u16);
        tags.push(name.to_owned());
        Ok(id)
    }

    pub fn get_or_add_tag(&mut self, name: &str) -> Result<TagId, TagError> {
        if let Some(index) = self.tags.iter().position(|n| n == name) {
            Ok(TagId(index as u16))
        } else {
            Self::add_tag(&mut self.tags, name)
        }
    }

    pub fn image(&self) -> DefsImage {
        DefsImage {
            defs: self.defs.len(),
            tags: self.tags.len(),
        }
    }

    pub fn restore_image(&mut self, image: &DefsImage) {
        self.defs.resize_with(image.defs, || {
            panic!("image must be a subset of the current defs")
        });
        self.tags.resize_with(image.tags, || {
            panic!("image must be a subset of the current defs")
        });
    }

    pub fn serialize_defs(&self) -> String {
        let mut result = String::new();
        for def in &self.defs {
            result.push_str(def);
            result.push('\n');
        }
        result
    }

    pub fn serialize_tags(&self) -> String {
        let mut result = String::new();
        for tag in &self.tags {
            result.push_str(tag);
            result.push('\n');
        }
        result
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DefError {
    Exists,
    OutOfSpace,
}

impl Display for DefError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            DefError::Exists => "definition already exists",
            DefError::OutOfSpace => "too many definitions",
        })
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TagError {
    OutOfSpace,
}

impl Display for TagError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            TagError::OutOfSpace => "too many tags",
        })
    }
}

macro_rules! well_known_tags {
    ($($ident:tt = $value:tt),* $(,)?) => {
        impl TagId {
            $(
                #[allow(non_upper_case_globals)]
                pub const $ident: Self = Self($value);
            )*
        }

        fn add_well_known_tags(tags: &mut Vec<String>) {
            $(
                let id = Defs::add_tag(tags, stringify!($ident)).unwrap();
                assert_eq!(id, TagId::from_u16($value));
            )*
        }
    }
}

well_known_tags! {
    // NOTE: The numbers must be sorted from 0 to N, due to limitations of Rust's macro system.
    // https://github.com/rust-lang/rust/issues/83527

    Nil = 0,

    From = 1,
    To = 2,
    Num = 3,
}