serpent-serializer 0.2.0

//! The deserializer.
//!
//! Parses the Lua subset that the serializer emits and rebuilds a [`Value`].
//! Handles both forms: a bare expression from `line` and `block`, and the
//! `do local _ = ...; ... return _ end` block from `dump`, including the
//! self-reference section that wires shared and cyclic references.
//!
//! Safety follows serpent's sandbox. In safe mode any function call in the
//! input is rejected with "cannot call functions", and any global assignment or
//! read resolves against a sandbox, so it cannot touch real globals. Unsafe mode
//! is not modeled because this loader has no real global environment to expose.

use crate::value::{Global, Key, Table, Value};
use std::collections::HashMap;

/// Options for [`load`].
#[derive(Default)]
pub struct LoadOptions {
    /// When `false`, allow function calls in the input. Defaults to safe (true).
    /// The exact serpent gate is `safe == false`, so only an explicit `false`
    /// disables safety.
    pub safe: Option<bool>,
}

/// An error from [`load`].
///
/// Each variant names a distinct failure the parser can hit, so callers can
/// match on the mode instead of parsing a message string.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LoadError {
    /// A function call reached in safe mode. The sandbox rejects calls.
    UnsafeCall,
    /// An unexpected byte at a source position. `byte` is `None` at end of input.
    UnexpectedByte {
        /// The offending byte, or `None` at end of input.
        byte: Option<u8>,
        /// The byte offset in the input.
        pos: usize,
    },
    /// A string or table that never closed. `what` names which one.
    Unterminated {
        /// Either `"string"` or `"table"`.
        what: &'static str,
    },
    /// A `return NAME` in a `do` block that names an undeclared local.
    UnknownVariable(String),
    /// A numeric literal that did not parse.
    BadNumber(String),
    /// Any other syntax error, with a message.
    Syntax(String),
}

impl std::fmt::Display for LoadError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            LoadError::UnsafeCall => write!(f, "cannot call functions"),
            LoadError::UnexpectedByte { byte, pos } => {
                write!(f, "unexpected byte {byte:?} at {pos}")
            }
            LoadError::Unterminated { what } => write!(f, "unterminated {what}"),
            LoadError::UnknownVariable(name) => write!(f, "unknown variable {name}"),
            LoadError::BadNumber(text) => write!(f, "bad number {text}"),
            LoadError::Syntax(msg) => write!(f, "{msg}"),
        }
    }
}

impl std::error::Error for LoadError {}

/// Deserialize a serpent fragment.
///
/// Accepts either a bare expression or a `do ... end` block. Returns `Ok(value)`
/// on success. Returns a [`LoadError`] on a parse error, or in safe mode when
/// the input tries to call a function.
///
/// # Errors
///
/// Returns [`LoadError`] when parsing fails or when a call is rejected in safe
/// mode.
pub fn load(data: &str, opts: &LoadOptions) -> Result<Value, LoadError> {
    let safe = opts.safe != Some(false);
    let mut p = Parser::new(data, safe);
    p.parse_chunk()
}

struct Parser<'a> {
    src: &'a [u8],
    pos: usize,
    safe: bool,
    /// Local variables declared in a `do ... end` block, by name.
    locals: HashMap<String, Value>,
    /// Stable ids for global names, so a global used twice stays one identity.
    global_ids: HashMap<String, usize>,
    next_global: usize,
}

impl<'a> Parser<'a> {
    fn new(src: &'a str, safe: bool) -> Self {
        Parser {
            src: src.as_bytes(),
            pos: 0,
            safe,
            locals: HashMap::new(),
            global_ids: HashMap::new(),
            next_global: 1,
        }
    }

    /// A `Global` value standing in for a name looked up in the sandbox. Two
    /// lookups of the same name share one identity, matching serpent's sandbox
    /// where every global read returns the same table.
    fn global_value(&mut self, name: String) -> Value {
        let id = *self.global_ids.entry(name.clone()).or_insert_with(|| {
            let id = self.next_global;
            self.next_global += 1;
            id
        });
        Value::Global(Global { name, id })
    }

    fn parse_chunk(&mut self) -> Result<Value, LoadError> {
        self.skip_ws();
        let v = if self.looking_at_keyword("do") {
            self.parse_do_block()?
        } else {
            self.parse_expr()?
        };
        self.skip_comment();
        self.skip_ws();
        if let Some(byte) = self.peek() {
            return Err(LoadError::UnexpectedByte {
                byte: Some(byte),
                pos: self.pos,
            });
        }
        Ok(v)
    }

    /// Parse `do local NAME = EXPR; ...statements...; [return NAME;] end`.
    ///
    /// serpent's `dump` always ends with `return NAME`, but the sandbox test
    /// inputs omit it, so a block without a `return` yields nil.
    fn parse_do_block(&mut self) -> Result<Value, LoadError> {
        self.expect_word("do")?;
        self.skip_ws();
        loop {
            self.skip_ws();
            if self.looking_at_keyword("return") || self.at_end_keyword() {
                break;
            }
            if self.looking_at_keyword("local") {
                self.parse_local()?;
            } else {
                self.parse_statement()?;
            }
            self.skip_sep();
        }
        let mut val = Value::Nil;
        if self.looking_at_keyword("return") {
            self.expect_word("return")?;
            self.skip_ws();
            let name = self.parse_name()?;
            val = self
                .locals
                .get(&name)
                .cloned()
                .ok_or(LoadError::UnknownVariable(name))?;
            self.skip_sep();
        }
        self.skip_ws();
        self.expect_word("end")?;
        Ok(val)
    }

    /// Whether the cursor is at the `end` keyword, not just any word starting
    /// with those letters.
    fn at_end_keyword(&self) -> bool {
        self.looking_at_keyword("end")
    }

    /// Parse `local NAME = EXPR` or `local NAME = {}`.
    fn parse_local(&mut self) -> Result<(), LoadError> {
        self.expect_word("local")?;
        self.skip_ws();
        let name = self.parse_name()?;
        self.skip_ws();
        if self.peek() == Some(b'=') {
            self.pos += 1;
            let v = self.parse_expr()?;
            self.locals.insert(name, v);
        } else {
            self.locals.insert(name, Value::Nil);
        }
        Ok(())
    }

    /// Parse a statement inside a `do ... end` block. It is either an assignment
    /// into a self-reference path, a call statement, or an assignment to a global
    /// that the sandbox discards.
    fn parse_statement(&mut self) -> Result<(), LoadError> {
        // Left side: NAME then a chain of [expr] or .name.
        let base = self.parse_name()?;
        let known_local = matches!(self.locals.get(&base), Some(Value::Table(_)));
        if !known_local {
            return self.consume_sandbox_statement();
        }
        let mut table = match self.locals.get(&base) {
            Some(Value::Table(t)) => t.clone(),
            _ => unreachable!("checked known_local"),
        };
        let mut pending_key: Option<Key> = None;
        loop {
            self.skip_ws();
            match self.peek() {
                Some(b'[') => {
                    self.pos += 1;
                    let k = self.parse_expr()?;
                    self.skip_ws();
                    self.expect_byte(b']')?;
                    if let Some(prev) = pending_key.take() {
                        table = self.descend(&table, prev)?;
                    }
                    pending_key = Some(value_to_key(k)?);
                }
                Some(b'.') => {
                    self.pos += 1;
                    let field = self.parse_name()?;
                    if let Some(prev) = pending_key.take() {
                        table = self.descend(&table, prev)?;
                    }
                    pending_key = Some(Key::Str(field.into_bytes()));
                }
                _ => break,
            }
        }
        self.skip_ws();
        self.expect_byte(b'=')?;
        let rhs = self.parse_expr()?;
        let key = pending_key.ok_or(LoadError::Syntax("assignment with no key".to_string()))?;
        table.set(key, rhs);
        Ok(())
    }

    /// Consume a statement whose target is a global or `_G` chain, not a known
    /// local. A call is rejected in safe mode. An assignment resolves against the
    /// sandbox and is discarded, so it never touches real globals.
    fn consume_sandbox_statement(&mut self) -> Result<(), LoadError> {
        // Walk the field/index chain following the base name.
        loop {
            self.skip_ws();
            match self.peek() {
                Some(b'.') => {
                    self.pos += 1;
                    let _ = self.parse_name();
                }
                Some(b'[') => {
                    self.pos += 1;
                    let _ = self.parse_expr()?;
                    self.skip_ws();
                    self.expect_byte(b']')?;
                }
                _ => break,
            }
        }
        self.skip_ws();
        match self.peek() {
            Some(b'(') => {
                if self.safe {
                    return Err(LoadError::UnsafeCall);
                }
                self.skip_balanced_parens()?;
                Ok(())
            }
            Some(b'=') if self.peek_at(1) != Some(b'=') => {
                self.pos += 1;
                // The right side runs in the sandbox and its result is dropped.
                let _ = self.parse_expr()?;
                Ok(())
            }
            _ => Ok(()),
        }
    }

    /// Follow a key into a nested table, for multi-step assignment paths.
    fn descend(&self, table: &Table, key: Key) -> Result<Table, LoadError> {
        match table.get(&key) {
            Some(Value::Table(t)) => Ok(t),
            _ => Err(LoadError::Syntax("path does not reach a table".to_string())),
        }
    }

    fn parse_expr(&mut self) -> Result<Value, LoadError> {
        self.skip_ws();
        match self.peek() {
            Some(b'{') => self.parse_table(),
            Some(b'"') | Some(b'\'') => self.parse_string(),
            Some(b'-') | Some(b'0'..=b'9') => self.parse_number_or_special(),
            Some(c) if c.is_ascii_alphabetic() || c == b'_' => self.parse_word_expr(),
            other => Err(LoadError::UnexpectedByte {
                byte: other,
                pos: self.pos,
            }),
        }
    }

    /// A word-starting expression: a keyword literal, a local variable with an
    /// access chain, a `function() ... end` stub, or a global lookup or call.
    fn parse_word_expr(&mut self) -> Result<Value, LoadError> {
        let word = self.parse_name()?;
        match word.as_str() {
            "nil" => return Ok(Value::Nil),
            "true" => return Ok(Value::Bool(true)),
            "false" => return Ok(Value::Bool(false)),
            "function" => return self.parse_function_stub(),
            _ => {}
        }

        if let Some(base) = self.locals.get(&word).cloned() {
            // Follow field and index access into the local value. serpent uses
            // the helper table this way, for example __._table1 or _[key].
            return self.follow_access(base);
        }

        // A global chain such as math.random, io.stdin, or _G.print, optionally
        // called. Capture the dotted name for a stable sandbox identity.
        let mut name = word;
        self.skip_ws();
        while matches!(self.peek(), Some(b'.') | Some(b'[')) {
            if self.peek() == Some(b'.') {
                self.pos += 1;
                let field = self.parse_name()?;
                name.push('.');
                name.push_str(&field);
            } else {
                self.pos += 1;
                self.parse_expr()?;
                self.skip_ws();
                self.expect_byte(b']')?;
                name.push_str("[]");
            }
            self.skip_ws();
        }
        if self.peek() == Some(b'(') {
            if self.safe {
                return Err(LoadError::UnsafeCall);
            }
            self.skip_balanced_parens()?;
            return Ok(Value::Nil);
        }
        // A global name resolves to a sandbox stand-in with a stable identity.
        Ok(self.global_value(name))
    }

    /// Follow `.field` and `[index]` access starting from `base`.
    fn follow_access(&mut self, mut cur: Value) -> Result<Value, LoadError> {
        loop {
            self.skip_ws();
            let key = match self.peek() {
                Some(b'.') => {
                    self.pos += 1;
                    Key::Str(self.parse_name()?.into_bytes())
                }
                Some(b'[') => {
                    self.pos += 1;
                    let k = self.parse_expr()?;
                    self.skip_ws();
                    self.expect_byte(b']')?;
                    value_to_key(k)?
                }
                _ => return Ok(cur),
            };
            let next = match &cur {
                Value::Table(t) => t.get(&key).unwrap_or(Value::Nil),
                _ => Value::Nil,
            };
            cur = next;
        }
    }

    /// Parse a `function() ... end` stub. The serializer emits either a bytecode
    /// reload wrapped in a call, or under nocode an empty body. This loader
    /// cannot run functions, so it yields nil after consuming the body.
    fn parse_function_stub(&mut self) -> Result<Value, LoadError> {
        // Skip the parameter list.
        self.skip_ws();
        if self.peek() == Some(b'(') {
            self.skip_balanced_parens()?;
        }
        // Skip to the matching `end`, accounting for nested blocks.
        let mut depth = 1;
        while depth > 0 {
            self.skip_ws();
            match self.peek() {
                None => return Err(LoadError::Syntax("expected end".to_string())),
                Some(b'-') if self.peek_at(1) == Some(b'-') => {
                    self.skip_function_body_comment();
                }
                Some(b'[') => {
                    if !self.skip_long_bracket() {
                        self.pos += 1;
                    }
                }
                Some(quote @ (b'"' | b'\'')) => {
                    self.pos += 1;
                    self.skip_quoted_bytes(quote)?;
                }
                Some(c) if is_ident_start(c) => {
                    let word = self.parse_name()?;
                    match word.as_str() {
                        "end" => depth -= 1,
                        "function" | "do" | "if" => depth += 1,
                        _ => {}
                    }
                }
                Some(_) => {
                    self.pos += 1;
                }
            }
        }
        Ok(Value::Nil)
    }

    fn skip_balanced_parens(&mut self) -> Result<(), LoadError> {
        self.expect_byte(b'(')?;
        let mut depth = 1;
        while depth > 0 {
            match self.next_byte() {
                Some(b'"') => self.skip_quoted_bytes(b'"')?,
                Some(b'\'') => self.skip_quoted_bytes(b'\'')?,
                Some(b'(') => depth += 1,
                Some(b')') => depth -= 1,
                Some(_) => {}
                None => return Err(LoadError::Syntax("unbalanced parentheses".to_string())),
            }
        }
        Ok(())
    }

    fn parse_table(&mut self) -> Result<Value, LoadError> {
        self.expect_byte(b'{')?;
        let table = Table::new();
        let mut array_idx = 1u64;
        loop {
            // Skip a comment before the element, so pretty output that leads
            // with `--[[...]]` inside the braces parses.
            self.skip_comment();
            self.skip_ws();
            match self.peek() {
                Some(b'}') => {
                    self.pos += 1;
                    break;
                }
                None => return Err(LoadError::Unterminated { what: "table" }),
                _ => {}
            }
            // [key] = value, name = value, or a positional value.
            if self.peek() == Some(b'[') {
                self.pos += 1;
                let key = self.parse_expr()?;
                self.skip_ws();
                self.expect_byte(b']')?;
                self.skip_ws();
                self.expect_byte(b'=')?;
                let val = self.parse_expr()?;
                table.set(value_to_key(key)?, val);
            } else if let Some(name) = self.try_named_field()? {
                let val = self.parse_expr()?;
                table.set(Key::Str(name.into_bytes()), val);
            } else {
                let val = self.parse_expr()?;
                table.set(Key::Number(array_idx as f64), val);
                array_idx += 1;
            }
            // A value may carry a trailing --[[...]] comment before the
            // separator, for example a nested table in pretty output.
            self.skip_comment();
            self.skip_ws();
            match self.peek() {
                Some(b',') | Some(b';') => {
                    self.pos += 1;
                }
                Some(b'}') => {}
                other => {
                    return Err(LoadError::Syntax(format!(
                        "expected , or }} in table, got {other:?}"
                    )))
                }
            }
        }
        Ok(Value::Table(table))
    }

    /// Try to parse a `name =` field prefix. Returns the name on success and
    /// leaves the cursor after `=`. Restores position when the lookahead is not
    /// a named field.
    fn try_named_field(&mut self) -> Result<Option<String>, LoadError> {
        let save = self.pos;
        match self.peek() {
            Some(c) if c.is_ascii_alphabetic() || c == b'_' => {}
            _ => return Ok(None),
        }
        let name = self.parse_name()?;
        self.skip_ws();
        if self.peek() == Some(b'=') && self.peek_at(1) != Some(b'=') {
            self.pos += 1;
            Ok(Some(name))
        } else {
            self.pos = save;
            Ok(None)
        }
    }

    fn parse_string(&mut self) -> Result<Value, LoadError> {
        let quote = self
            .next_byte()
            .ok_or(LoadError::Unterminated { what: "string" })?;
        let mut bytes = Vec::new();
        loop {
            match self.next_byte() {
                None => return Err(LoadError::Unterminated { what: "string" }),
                Some(b) if b == quote => break,
                Some(b'\\') => {
                    let e = self
                        .next_byte()
                        .ok_or(LoadError::Unterminated { what: "string" })?;
                    match e {
                        b'n' => bytes.push(b'\n'),
                        b'r' => bytes.push(b'\r'),
                        b't' => bytes.push(b'\t'),
                        b'a' => bytes.push(7),
                        b'b' => bytes.push(8),
                        b'f' => bytes.push(12),
                        b'v' => bytes.push(11),
                        b'\\' => bytes.push(b'\\'),
                        b'"' => bytes.push(b'"'),
                        b'\'' => bytes.push(b'\''),
                        b'\n' => bytes.push(b'\n'),
                        b'0'..=b'9' => {
                            // Up to three decimal digits.
                            let mut n = (e - b'0') as u32;
                            for _ in 0..2 {
                                if let Some(d @ b'0'..=b'9') = self.peek() {
                                    n = n * 10 + (d - b'0') as u32;
                                    self.pos += 1;
                                } else {
                                    break;
                                }
                            }
                            bytes.push(n as u8);
                        }
                        other => bytes.push(other),
                    }
                }
                Some(b) => bytes.push(b),
            }
        }
        Ok(Value::Str(bytes))
    }

    /// Parse a number, or a special form: `1/0`, `-1/0`, `0/0`.
    fn parse_number_or_special(&mut self) -> Result<Value, LoadError> {
        let start = self.pos;
        if self.peek() == Some(b'-') {
            self.pos += 1;
        }
        while matches!(
            self.peek(),
            Some(b'0'..=b'9')
                | Some(b'.')
                | Some(b'e')
                | Some(b'E')
                | Some(b'+')
                | Some(b'-')
                | Some(b'x')
                | Some(b'X')
                | Some(b'a'..=b'f')
                | Some(b'A'..=b'F')
                | Some(b'p')
                | Some(b'P')
        ) {
            self.pos += 1;
        }
        let text = std::str::from_utf8(&self.src[start..self.pos])
            .map_err(|_| LoadError::BadNumber("bad bytes".to_string()))?;
        // Handle the division forms serpent emits for special numbers.
        self.skip_ws();
        if self.peek() == Some(b'/') {
            self.pos += 1;
            self.skip_ws();
            let dstart = self.pos;
            if self.peek() == Some(b'-') {
                self.pos += 1;
            }
            while matches!(self.peek(), Some(b'0'..=b'9') | Some(b'.')) {
                self.pos += 1;
            }
            let denom = std::str::from_utf8(&self.src[dstart..self.pos]).unwrap_or("");
            let num: f64 = text
                .trim()
                .parse()
                .map_err(|_| LoadError::BadNumber(text.trim().to_string()))?;
            let den: f64 = denom
                .trim()
                .parse()
                .map_err(|_| LoadError::BadNumber(denom.trim().to_string()))?;
            self.skip_comment();
            return Ok(Value::Number(num / den));
        }
        let n: f64 = if let Some(hex) = text.strip_prefix("0x").or_else(|| text.strip_prefix("0X"))
        {
            i64::from_str_radix(hex, 16)
                .map(|v| v as f64)
                .map_err(|_| LoadError::BadNumber(text.to_string()))?
        } else {
            text.parse()
                .map_err(|_| LoadError::BadNumber(text.to_string()))?
        };
        self.skip_comment();
        Ok(Value::Number(n))
    }

    fn parse_name(&mut self) -> Result<String, LoadError> {
        self.skip_ws();
        let start = self.pos;
        match self.peek() {
            Some(c) if c.is_ascii_alphabetic() || c == b'_' => {}
            _ => return Err(LoadError::Syntax(format!("expected name at {}", self.pos))),
        }
        while matches!(self.peek(), Some(c) if c.is_ascii_alphanumeric() || c == b'_') {
            self.pos += 1;
        }
        Ok(String::from_utf8_lossy(&self.src[start..self.pos]).into_owned())
    }

    fn skip_comment(&mut self) {
        self.skip_ws();
        if self.looking_at("--[[") {
            self.pos += 4;
            while self.pos < self.src.len() && !self.looking_at("]]") {
                self.pos += 1;
            }
            if self.looking_at("]]") {
                self.pos += 2;
            }
        }
    }

    fn skip_function_body_comment(&mut self) {
        self.pos += 2;
        let is_long_comment = self.peek() == Some(b'[') && self.skip_long_bracket();
        if !is_long_comment {
            while !matches!(self.peek(), None | Some(b'\n') | Some(b'\r')) {
                self.pos += 1;
            }
        }
    }

    fn skip_long_bracket(&mut self) -> bool {
        if self.peek() != Some(b'[') {
            return false;
        }
        let mut close_len = 1;
        while self.peek_at(close_len) == Some(b'=') {
            close_len += 1;
        }
        if self.peek_at(close_len) != Some(b'[') {
            return false;
        }
        self.pos += close_len + 1;
        while self.pos < self.src.len() && !self.looking_at_long_bracket_close(close_len) {
            self.pos += 1;
        }
        if self.looking_at_long_bracket_close(close_len) {
            self.pos += close_len + 1;
        }
        true
    }

    fn looking_at_long_bracket_close(&self, len: usize) -> bool {
        self.peek() == Some(b']')
            && (1..len).all(|i| self.peek_at(i) == Some(b'='))
            && self.peek_at(len) == Some(b']')
    }

    fn skip_sep(&mut self) {
        self.skip_ws();
        while matches!(self.peek(), Some(b';') | Some(b'\n')) {
            self.pos += 1;
            self.skip_ws();
        }
    }

    fn skip_ws(&mut self) {
        loop {
            match self.peek() {
                Some(b) if b == b' ' || b == b'\t' || b == b'\r' || b == b'\n' => {
                    self.pos += 1;
                }
                _ => break,
            }
        }
    }

    fn looking_at(&self, s: &str) -> bool {
        self.src[self.pos..].starts_with(s.as_bytes())
    }

    fn looking_at_keyword(&self, s: &str) -> bool {
        self.looking_at(s) && !matches!(self.peek_at(s.len()), Some(b) if is_ident_byte(b))
    }

    fn expect_word(&mut self, w: &str) -> Result<(), LoadError> {
        self.skip_ws();
        if self.looking_at_keyword(w) {
            self.pos += w.len();
            Ok(())
        } else {
            Err(LoadError::Syntax(format!("expected {w}")))
        }
    }

    fn expect_byte(&mut self, b: u8) -> Result<(), LoadError> {
        self.skip_ws();
        if self.peek() == Some(b) {
            self.pos += 1;
            Ok(())
        } else {
            Err(LoadError::Syntax(format!("expected {}", b as char)))
        }
    }

    fn peek(&self) -> Option<u8> {
        self.src.get(self.pos).copied()
    }

    fn peek_at(&self, off: usize) -> Option<u8> {
        self.src.get(self.pos + off).copied()
    }

    fn next_byte(&mut self) -> Option<u8> {
        let b = self.src.get(self.pos).copied();
        if b.is_some() {
            self.pos += 1;
        }
        b
    }

    fn skip_quoted_bytes(&mut self, quote: u8) -> Result<(), LoadError> {
        loop {
            match self.next_byte() {
                Some(b) if b == quote => return Ok(()),
                Some(b'\\') => {
                    if self.next_byte().is_none() {
                        return Err(LoadError::Unterminated { what: "string" });
                    }
                }
                Some(_) => {}
                None => return Err(LoadError::Unterminated { what: "string" }),
            }
        }
    }
}

fn is_ident_byte(b: u8) -> bool {
    b.is_ascii_alphanumeric() || b == b'_'
}

fn is_ident_start(b: u8) -> bool {
    b.is_ascii_alphabetic() || b == b'_'
}

fn value_to_key(v: Value) -> Result<Key, LoadError> {
    match v {
        Value::Bool(b) => Ok(Key::Bool(b)),
        Value::Number(n) => Ok(Key::Number(n)),
        Value::Str(s) => Ok(Key::Str(s)),
        Value::Table(t) => Ok(Key::Table(t)),
        Value::Function(f) => Ok(Key::Function(f)),
        Value::Global(g) => Ok(Key::Global(g)),
        Value::Nil => Err(LoadError::Syntax("nil is not a valid key".to_string())),
    }
}