1use std::collections::BTreeMap;
2use std::mem::size_of;
3
4use anyhow::{bail, Result};
5
6use crate::opcode::Opcode;
7use crate::parser::Parser;
8
9#[derive(Debug, Clone)]
10pub enum Op {
11 Nop,
12 InitStack { args: u8, locals: u8 },
13 Call { target: u32 },
14 Syscall { id: u16, name: String, args: u8 },
15 Ret,
16 RetV,
17 Jmp { target: u32 },
18 Jz { target: u32 },
19 PushNil,
20 PushTrue,
21 PushI32(i32),
22 PushI16(i16),
23 PushI8(i8),
24 PushF32(f32),
25 PushString(String),
26 PushGlobal(u16),
27 PushStack(i8),
28 PushGlobalTable(u16),
29 PushLocalTable(i8),
30 PushTop,
31 PushReturn,
32 PopGlobal(u16),
33 PopStack(i8),
34 PopGlobalTable(u16),
35 PopLocalTable(i8),
36 Neg,
37 Add,
38 Sub,
39 Mul,
40 Div,
41 Mod,
42 BitTest,
43 And,
44 Or,
45 SetE,
46 SetNE,
47 SetG,
48 SetLE,
49 SetL,
50 SetGE,
51 Unknown(u8),
52}
53
54#[derive(Debug, Clone)]
55pub struct Instruction {
56 pub addr: u32,
57 pub op: Op,
58}
59
60#[derive(Debug, Clone)]
61pub struct Function {
62 pub start_addr: u32,
63 pub args: u8,
64 pub locals: u8,
65 pub insts: Vec<Instruction>,
66}
67
68#[derive(Debug, Clone)]
69pub struct Program {
70 pub functions: BTreeMap<u32, Function>,
71}
72
73fn decode_one(parser: &Parser, pc: usize) -> Result<(Instruction, usize)> {
74 let addr = pc as u32;
75 let opcode_u8 = parser.read_u8(pc)?;
76
77 let mut cur = pc + 1;
78 let op = match opcode_u8.try_into() {
79 Ok(Opcode::Nop) => Op::Nop,
80 Ok(Opcode::InitStack) => {
81 let args = parser.read_i8(cur)? as u8;
82 cur += size_of::<i8>();
83 let locals = parser.read_i8(cur)? as u8;
84 cur += size_of::<i8>();
85 Op::InitStack { args, locals }
86 }
87 Ok(Opcode::Call) => {
88 let target = parser.read_u32(cur)?;
89 cur += size_of::<u32>();
90 Op::Call { target }
91 }
92 Ok(Opcode::Syscall) => {
93 let id = parser.read_u16(cur)?;
94 cur += size_of::<u16>();
95
96 let (name, args) = match parser.get_syscall(id) {
97 Some(s) => (s.name.clone(), s.args),
98 None => (format!("syscall_{}", id), 0),
99 };
100 Op::Syscall { id, name, args }
101 }
102 Ok(Opcode::Ret) => Op::Ret,
103 Ok(Opcode::RetV) => Op::RetV,
104 Ok(Opcode::Jmp) => {
105 let target = parser.read_u32(cur)?;
106 cur += size_of::<u32>();
107 Op::Jmp { target }
108 }
109 Ok(Opcode::Jz) => {
110 let target = parser.read_u32(cur)?;
111 cur += size_of::<u32>();
112 Op::Jz { target }
113 }
114 Ok(Opcode::PushNil) => Op::PushNil,
115 Ok(Opcode::PushTrue) => Op::PushTrue,
116 Ok(Opcode::PushI32) => {
117 let v = parser.read_i32(cur)?;
118 cur += size_of::<i32>();
119 Op::PushI32(v)
120 }
121 Ok(Opcode::PushI16) => {
122 let v = parser.read_i16(cur)?;
123 cur += size_of::<i16>();
124 Op::PushI16(v)
125 }
126 Ok(Opcode::PushI8) => {
127 let v = parser.read_i8(cur)?;
128 cur += size_of::<i8>();
129 Op::PushI8(v)
130 }
131 Ok(Opcode::PushF32) => {
132 let v = parser.read_f32(cur)?;
133 cur += size_of::<f32>();
134 Op::PushF32(v)
135 }
136 Ok(Opcode::PushString) => {
137 let len = parser.read_u8(cur)? as usize;
138 cur += size_of::<u8>();
139 let s = parser.read_cstring(cur, len)?;
140 cur += len;
141 Op::PushString(s)
142 }
143 Ok(Opcode::PushGlobal) => {
144 let idx = parser.read_u16(cur)?;
145 cur += size_of::<u16>();
146 Op::PushGlobal(idx)
147 }
148 Ok(Opcode::PushStack) => {
149 let off = parser.read_i8(cur)?;
150 cur += size_of::<i8>();
151 Op::PushStack(off)
152 }
153 Ok(Opcode::PushGlobalTable) => {
154 let idx = parser.read_u16(cur)?;
155 cur += size_of::<u16>();
156 Op::PushGlobalTable(idx)
157 }
158 Ok(Opcode::PushLocalTable) => {
159 let idx = parser.read_i8(cur)?;
160 cur += size_of::<i8>();
161 Op::PushLocalTable(idx)
162 }
163 Ok(Opcode::PushTop) => Op::PushTop,
164 Ok(Opcode::PushReturn) => Op::PushReturn,
165 Ok(Opcode::PopGlobal) => {
166 let idx = parser.read_u16(cur)?;
167 cur += size_of::<u16>();
168 Op::PopGlobal(idx)
169 }
170 Ok(Opcode::PopStack) => {
171 let off = parser.read_i8(cur)?;
172 cur += size_of::<i8>();
173 Op::PopStack(off)
174 }
175 Ok(Opcode::PopGlobalTable) => {
176 let idx = parser.read_u16(cur)?;
177 cur += size_of::<u16>();
178 Op::PopGlobalTable(idx)
179 }
180 Ok(Opcode::PopLocalTable) => {
181 let idx = parser.read_i8(cur)?;
182 cur += size_of::<i8>();
183 Op::PopLocalTable(idx)
184 }
185 Ok(Opcode::Neg) => Op::Neg,
186 Ok(Opcode::Add) => Op::Add,
187 Ok(Opcode::Sub) => Op::Sub,
188 Ok(Opcode::Mul) => Op::Mul,
189 Ok(Opcode::Div) => Op::Div,
190 Ok(Opcode::Mod) => Op::Mod,
191 Ok(Opcode::BitTest) => Op::BitTest,
192 Ok(Opcode::And) => Op::And,
193 Ok(Opcode::Or) => Op::Or,
194 Ok(Opcode::SetE) => Op::SetE,
195 Ok(Opcode::SetNE) => Op::SetNE,
196 Ok(Opcode::SetG) => Op::SetG,
197 Ok(Opcode::SetLE) => Op::SetLE,
198 Ok(Opcode::SetL) => Op::SetL,
199 Ok(Opcode::SetGE) => Op::SetGE,
200 Err(_) => Op::Unknown(opcode_u8),
201 };
202
203 Ok((Instruction { addr, op }, cur))
204}
205
206pub fn decode_program(parser: &Parser) -> Result<Program> {
207 let end = parser.get_sys_desc_offset() as usize;
208 let mut pc = 4usize;
209
210 let mut cur_fn: Option<Function> = None;
211 let mut functions: BTreeMap<u32, Function> = BTreeMap::new();
212
213 while pc < end {
214 let (inst, next_pc) = decode_one(parser, pc)?;
215
216 match &inst.op {
217 Op::InitStack { args, locals } => {
218 if let Some(f) = cur_fn.take() {
219 functions.insert(f.start_addr, f);
220 }
221 cur_fn = Some(Function {
222 start_addr: inst.addr,
223 args: *args,
224 locals: *locals,
225 insts: vec![inst],
226 });
227 }
228 _ => {
229 if let Some(f) = cur_fn.as_mut() {
230 f.insts.push(inst);
231 } else {
232 cur_fn = Some(Function {
234 start_addr: inst.addr,
235 args: 0,
236 locals: 0,
237 insts: vec![inst],
238 });
239 }
240 }
241 }
242
243 if next_pc <= pc {
244 bail!("decoder did not advance at pc=0x{:08X}", pc);
245 }
246 pc = next_pc;
247 }
248
249 if let Some(f) = cur_fn.take() {
250 functions.insert(f.start_addr, f);
251 }
252
253 Ok(Program { functions })
254}
255
256pub fn all_syscalls(program: &Program) -> BTreeMap<String, u8> {
257 let mut m = BTreeMap::new();
258 for f in program.functions.values() {
259 for inst in &f.insts {
260 if let Op::Syscall { name, args, .. } = &inst.op {
261 m.entry(name.clone()).or_insert(*args);
262 }
263 }
264 }
265 m
266}
267
268pub fn all_call_targets(program: &Program) -> BTreeMap<u32, ()> {
269 let mut m = BTreeMap::new();
270 for f in program.functions.values() {
271 for inst in &f.insts {
272 if let Op::Call { target } = &inst.op {
273 m.insert(*target, ());
274 }
275 }
276 }
277 m
278}
279