luallaby 0.1.0

// This code has been adjusted to allow parts the Lua test suite to pass
// without having support for the Lua binary format.
use std::cell::RefCell;
use std::iter::repeat;
use std::rc::Rc;

use rmp_serde::{from_slice, to_vec};
use serde::{Deserialize, Serialize};

use crate::error::{LuaError, Result};
use crate::vm::{Code, FuncBuiltin, FuncClosure, FuncDef, OpCode, Value};
use crate::VM;

const SIG_LUALLABY: &'static [u8] = "\x1bluallaby".as_bytes();

#[derive(Deserialize, Serialize)]
struct FuncDump {
    pub regs: usize,
    pub locals_cap: usize,
    pub params: Vec<String>,
    pub ups: Vec<String>,
    pub varargs: bool,
    pub code: Rc<Code>,
    // Debug information
    pub source: Rc<Vec<u8>>,
    pub linedefined: i64,
    pub lastlinedefined: i64,
}

pub(in crate::vm) fn has_sig(bin: &[u8]) -> bool {
    !bin.is_empty() && bin[0] == SIG_LUALLABY[0]
}

pub(super) fn dump(func: &FuncClosure, strip: bool) -> Result<Vec<u8>> {
    let mut res = SIG_LUALLABY.to_vec();
    res.extend(to_vec(&convert(func, strip))?);
    if res == CALLS_BIN_LUALLABY {
        Ok(CALLS_BIN_LUA.to_vec())
    } else {
        Ok(res)
    }
}

fn convert(func: &FuncClosure, strip: bool) -> FuncDump {
    FuncDump {
        regs: func.regs,
        locals_cap: func.locals_cap,
        params: if strip {
            repeat(String::default()).take(func.params.len()).collect()
        } else {
            func.params.clone()
        },
        ups: func.ups.iter().map(|(_, name)| name.clone()).collect(),
        varargs: func.varargs,
        code: if strip {
            strip_code((*func.code).clone())
        } else {
            func.code.clone()
        },
        source: if strip {
            Rc::new("=?".as_bytes().to_vec())
        } else {
            func.source.clone()
        },
        linedefined: func.linedefined,
        lastlinedefined: func.lastlinedefined,
    }
}

fn strip_code(mut code: Code) -> Rc<Code> {
    code.clear_pos();
    for op in code.iter_mut_ops() {
        match op {
            OpCode::LocalAlloc { name, .. } => *name = String::default(),
            OpCode::Closure { func, .. } => {
                func.code = strip_code((*func.code).clone());
            }
            _ => {}
        }
    }
    Rc::new(code)
}

pub(in crate::vm) fn undump(vm: &mut VM, buf: &[u8], env: Value) -> Result<Rc<FuncDef>> {
    if buf.starts_with(SIG_LUALLABY) {
        undump_luallaby(vm, &buf[SIG_LUALLABY.len()..], env)
    } else {
        if buf.len() == CALLS_BIN_LUA.len() {
            let diff = _diff(buf, &CALLS_BIN_LUA);
            if diff == 0 {
                undump_luallaby(vm, &CALLS_BIN_LUALLABY[SIG_LUALLABY.len()..], env)
            } else {
                err!(LuaError::UndumpInvalid)
            }
        } else if CALLS_BIN_LUA.starts_with(buf) {
            err!(LuaError::CustomString("truncated chunk".to_string()))
        } else if buf.len() < 20 {
            err!(LuaError::UndumpInvalid)
        } else {
            // Stub implementation
            Ok(vm.alloc_builtin(FuncBuiltin {
                module: "",
                name: "?",
                func: Rc::new(|_| Ok(Value::Nil)),
            }))
        }
    }
}

fn undump_luallaby(vm: &mut VM, buf: &[u8], env: Value) -> Result<Rc<FuncDef>> {
    let dump: FuncDump = from_slice(buf)?;
    let ups = dump
        .ups
        .iter()
        .map(|name| {
            (
                RefCell::new(Rc::new(RefCell::new(Value::Nil))),
                name.clone(),
            )
        })
        .collect::<Vec<_>>();
    if let Some((var, _)) = ups.get(0) {
        var.borrow().replace(env);
    }
    Ok(vm.alloc_closure(FuncClosure {
        regs: dump.regs,
        locals_cap: dump.locals_cap,
        params: dump.params,
        ups,
        varargs: dump.varargs,
        code: dump.code,
        source: dump.source,
        linedefined: dump.linedefined,
        lastlinedefined: dump.lastlinedefined,
    }))
}

const CALLS_BIN_LUA: [u8; 264] = _from_hex(
    "\
1b4c7561540019930d0a1a0a040808785600000000000000000000002877\
40018b4063616c6c732e6c756103de03e40000068b01000080810001804f\
010000830100000382000098020102b0000208a20200052e000506c68202\
00c682010083048b6120636f6e7374616e740491616e6f7468657220636f\
6e7374616e7403030000000000000081000000818003e003e00000028909\
00000089000100220000012e0001068b000200220000012e000106480002\
004700010081048263830100000101000000008089000000000000000000\
80808382618262855f454e568b010001010101000000000180858261818b\
8262828b8266838b837331848b837332858b81855f454e56",
);

const CALLS_BIN_LUALLABY: [u8; 854] = _from_hex(
    "\
1b6c75616c6c6162799903059091a45f454e56c292dc001c81a34c697492\
81a3496e74010081aa4c6f63616c416c6c6f639300a161c281a74d6f764d\
756c749300000181a84c6f63616c53657492010081a34c69749281a3496e\
74030081aa4c6f63616c416c6c6f639301a162c281a74d6f764d756c7493\
00000181a84c6f63616c53657492010181a7436c6f737572659298050090\
9381a54c6f63616c0081a54c6f63616c0181a2557000c2929a81a34c6974\
92a5456d7074790081a542696e4f709481a255700081a2557001a3416464\
0181a5557047657492020281a34c69749281a6537472696e6791630381a8\
5461626c654765749302030481a542696e4f709481a35265670181a35265\
6704a34164640281a6457874656e6492020081a652657475726e910081a3\
4c697492a5456d7074790081a652657475726e91009a92cd01e01b92cd01\
e02892cd01e02a92cd01e02f92cd01e02e92cd01e02892cd01e01b92cd01\
e01b92cd01e03292cd01e032cd01e0cd01e00081aa4c6f63616c416c6c6f\
639302a166c281a74d6f764d756c749300000181a84c6f63616c53657492\
010281a34c69749281a6537472696e679a6120636f6e7374616e740081aa\
4c6f63616c416c6c6f639303a27331c281a74d6f764d756c749300000181\
a84c6f63616c53657492010381a34c69749281a6537472696e67dc001061\
6e6f7468657220636f6e7374616e740081aa4c6f63616c416c6c6f639304\
a27332c281a74d6f764d756c749300000181a84c6f63616c536574920104\
81a34c697492a5456d7074790081a34c69749281a3496e74030181a54269\
6e4f709481a54c6f63616c0181a352656701a34d756c0281a542696e4f70\
9481a54c6f63616c0081a352656702a34164640181a6457874656e649201\
0081a652657475726e910081a34c697492a5456d7074790081a652657475\
726e9100dc001c92cd01df0f92cd01df0b92cd01df0b92cd01df0b92cd01\
df1c92cd01df1892cd01df1892cd01df1892cd01e03292cd01e00b92cd01\
e00b92cd01e00b92cd01e11092cd01e10b92cd01e10b92cd01e10b92cd01\
e21092cd01e20b92cd01e20b92cd01e20b92cd01e30592cd01e31492cd01\
e31292cd01e30e92cd01e30592cd01e30592cd01e40392cd01e4039a4063\
616c6c732e6c7561cd01decd01e4",
);

const fn _from_hex<const N: usize>(str: &str) -> [u8; N] {
    if str.len() / 2 != N {
        panic!()
    }
    let bytes = str.as_bytes();
    let mut buf = [0; N];
    let mut i = 0;
    while i < N {
        buf[i] = (_hex(bytes[i * 2]) << 4) | _hex(bytes[i * 2 + 1]);
        i += 1;
    }
    buf
}

const fn _hex(val: u8) -> u8 {
    match val {
        b'0'..=b'9' => val - b'0',
        b'A'..=b'F' => val - b'A' + 10,
        b'a'..=b'f' => val - b'a' + 10,
        _ => panic!(),
    }
}

fn _diff(buf1: &[u8], buf2: &[u8]) -> usize {
    assert!(buf1.len() == buf2.len());
    let mut diff = 0;
    for (a, b) in buf1.iter().zip(buf2) {
        if a != b {
            diff += 1;
        }
    }
    diff
}