lunar_wave/src/loader.rs

use std::io::Read;

use crate::state::{
	Block,
	Chunk,
	Instruction as Inst,
};

pub fn parse_inst (buf: [u8; 4]) -> Option <Inst>
{
	let opcode = buf [0] & 0x7f;
	
	let a = (buf [0] >> 7) | ((buf [1] & 0x7f) << 1);
	let b = buf [2];
	let c = buf [3];
	let bx = 
		(((buf [1] >> 7) as u32) << 0) | 
		((buf [2] as u32) << 1) |
		((buf [3] as u32) << 9);
	let bx = bx.try_into().ok ()?;
	let sbx = bx - 65535;
	let k = (buf [1] & 0x80) >> 7 == 1;
	let s_j = a as i32 + ((b as i32) << 8) + 1;
	
	Some (match opcode {
		0x00 => Inst::Move (a, b),
		0x01 => Inst::LoadI (a, sbx),
		0x03 => Inst::LoadK (a, bx),
		0x09 => Inst::GetUpVal (a, b),
		0x0b => Inst::GetTabUp (a, b, c),
		0x0d => Inst::GetI (a, b, c),
		0x22 => Inst::Add (a, b, c),
		0x24 => Inst::Mul (a, b, c),
		0x2e => Inst::MmBin (a, b, c),
		0x3c => Inst::EqK (a, b, c),
		0x38 => Inst::Jmp (s_j),
		0x44 => Inst::Call (a, b, c),
		0x46 => Inst::Return (a, b, c, k),
		0x47 => Inst::Return0,
		0x48 => Inst::Return1 (a),
		0x4f => Inst::Closure (a, bx),
		0x51 => Inst::VarArgPrep (a.into ()),
		_ => return None,
	})
}

#[derive (Debug, PartialEq)]
struct Header {
	inst_count: u8,
}

fn parse_file_name <R: Read> (rdr: &mut R) -> Option <String> {
	let file_name_sz = {
		let mut file_name_sz = [0u8; 1];
		rdr.read_exact (&mut file_name_sz).ok ()?;
		usize::try_from (file_name_sz [0] - 0x80 - 1).ok ()?
	};
	
	{
		let mut file_name = vec! [0u8; file_name_sz];
		rdr.read_exact (&mut file_name).ok ()?;
		Some (String::from_utf8 (file_name).ok ()?)
	}
}

fn parse_header (buf: [u8; 6]) -> Option <Header> {
	if buf [0] & 0x80 != 0x80 {
		// Not a function header
		return None;
	}
	
	Some (Header {
		inst_count: buf [5] - 0x80,
	})
}

// I don't know what this really is, so I'm calling it a trailer for now
// It appears in luac files after the string table / constants table
// for each function.

#[derive (Debug, PartialEq)]
struct Trailer {
	upvalue_count: u8,
}

fn parse_trailer (buf: [u8; 6]) -> Option <Trailer> {
	Some (Trailer {
		upvalue_count: buf [0] - 0x80,
	})
}

pub fn parse_block <R: Read> (rdr: &mut R) -> Option <Block> 
{
	let header = {
		let mut buf = [0u8; 6];
		rdr.read_exact (&mut buf).ok ()?;
		parse_header (buf)?
	};
	
	let mut instructions = Vec::with_capacity (header.inst_count as usize);
	
	for _ in 0..header.inst_count {
		let mut buf = [0u8; 4];
		rdr.read_exact (&mut buf).ok ()?;
		instructions.push (parse_inst (buf).expect (&format! ("{buf:?}")));
	}
	
	let constant_count = {
		let mut buf = [0u8; 1];
		rdr.read_exact (&mut buf).ok ()?;
		buf [0] - 0x80
	};
	
	let mut constants = Vec::with_capacity (constant_count as usize);
	
	for _ in 0..constant_count {
		let mut buf = [0u8; 2];
		rdr.read_exact (&mut buf).ok ()?;
		
		let len = ((buf [0] as u32) << 8) + (buf [1] as u32) - 0x0481;
		
		let mut s = vec! [0u8; len.try_into().ok ()?];
		rdr.read_exact (&mut s).ok ()?;
		
		let s = String::from_utf8 (s).ok ()?;
		constants.push (s.into ());
	}
	
	let trailer = {
		let mut buf = [0u8; 6];
		rdr.read_exact (&mut buf).ok ()?;
		
		parse_trailer (buf)?
	};
	
	Some (Block {
		constants,
		instructions,
		upvalue_count: trailer.upvalue_count as usize,
	})
}


pub fn parse_chunk <R: Read> (rdr: &mut R) -> Option <Chunk> {
	// Discard 32 bytes from the start of the file.
	// This is magic number, version number, etc.
	
	let mut hdr = [0u8; 32];
	rdr.read_exact (&mut hdr).ok ()?;
	
	let file_name = parse_file_name (rdr)?;
	
	let mut blocks = vec![];
	
	while let Some (block) = parse_block (rdr) {
		blocks.push (block);
	}
	
	Some (Chunk {
		file_name,
		blocks,
	})
}

#[cfg (test)]
mod tests {
	#[test]
	fn parse_inst () {
		use super::Inst;
		
		for (input, expected) in [
			([0x51, 0x00, 0x00, 0x00], Inst::VarArgPrep (0)),
			([0x4f, 0x00, 0x00, 0x00], Inst::Closure (0, 0)),
			([0xcf, 0x00, 0x00, 0x00], Inst::Closure (1, 0)),
			([0x8b, 0x00, 0x00, 0x00], Inst::GetTabUp (1, 0, 0)),
			([0x03, 0x81, 0x00, 0x00], Inst::LoadK (2, 1)),
			([0xc4, 0x00, 0x02, 0x01], Inst::Call (1, 2, 1)),
			([0x80, 0x00, 0x00, 0x00], Inst::Move (1, 0)),
			([0xc4, 0x00, 0x01, 0x02], Inst::Call (1, 1, 2)),
			([0x0b, 0x01, 0x00, 0x00], Inst::GetTabUp (2, 0, 0)),
			([0x83, 0x01, 0x01, 0x00], Inst::LoadK (3, 2)),
			([0x44, 0x01, 0x02, 0x01], Inst::Call (2, 2, 1)),
			([0x0b, 0x01, 0x00, 0x00], Inst::GetTabUp (2, 0, 0)),
			([0x80, 0x01, 0x01, 0x00], Inst::Move (3, 1)),
			([0xc4, 0x01, 0x01, 0x00], Inst::Call (3, 1, 0)),
			([0x44, 0x01, 0x00, 0x01], Inst::Call (2, 0, 1)),
			([0x0b, 0x01, 0x00, 0x00], Inst::GetTabUp (2, 0, 0)),
			([0x83, 0x81, 0x01, 0x00], Inst::LoadK (3, 3)),
			([0x44, 0x01, 0x02, 0x01], Inst::Call (2, 2, 1)),
			([0x46, 0x01, 0x01, 0x01], Inst::Return (2, 1, 1, false)),
			
			([0x01, 0x00, 0x02, 0x80], Inst::LoadI (0, 5)),
			([0xc6, 0x80, 0x02, 0x00], Inst::Return (1, 2, 0, true)),
			
			([0x09, 0x00, 0x01, 0x00], Inst::GetUpVal (0, 1)),
			([0x48, 0x00, 0x02, 0x00], Inst::Return1 (0)),
			([0x47, 0x00, 0x01, 0x00], Inst::Return0),
			([0x8d, 0x00, 0x01, 0x01], Inst::GetI (1, 1, 1)),
			([0xbc, 0x00, 0x01, 0x00], Inst::EqK (1, 1, 0)),
			([0xb8, 0x02, 0x00, 0x80], Inst::Jmp (6)),
			([0x38, 0x02, 0x00, 0x80], Inst::Jmp (5)),
		] {
			let actual = super::parse_inst (input).unwrap ();
			assert_eq!(actual, expected);
		}
	}
	
	#[test]
	fn parse_header () {
		for (input, expected) in [
			// Bytes 0 and 1 are first line and last line for debugging
			// Byte 2 is numparams
			// Byte 3 is is_vararg
			// Byte 4 is slot count / max stack size
			// Byte 5 is instruction count
			
			([0x80, 0x80, 0x00, 0x01, 0x04, 0x92], (18,)),
			([0x81, 0x89, 0x00, 0x00, 0x03, 0x87], (7,)),
			([0x85, 0x88, 0x00, 0x00, 0x02, 0x86], (6,)),
		] {
			let actual = super::parse_header (input).unwrap ();
			assert_eq! (actual, super::Header {
				inst_count: expected.0,
			});
		}
	}
	
	#[test]
	fn parse_trailer () {
		for (input, expected) in [
			([0x81, 0x01, 0x00, 0x00, 0x81, 0x80], (1,)),
			([0x81, 0x00, 0x00, 0x00, 0x81, 0x80], (1,)),
			([0x82, 0x00, 0x00, 0x00, 0x01, 0x00], (2,)),
		] {
			let actual = super::parse_trailer (input).unwrap ();
			assert_eq! (actual, super::Trailer {
				upvalue_count: expected.0,
			});
		}
	}
}