1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
use super::*;
use complete::le_u8;

fn load_unsigned<'a>(mut limit: usize) -> impl Parser<&'a [u8], usize, ErrorTree<&'a [u8]>> {
    move |mut input| -> IResult<&'a [u8], usize> {
        let mut x = 0;
        limit >>= 7;
        loop {
            let (rest, b) = le_u8(input)?;
            input = rest;
            if x >= limit {
                context("integer overflow", fail::<&[u8], &str, _>)(input)?;
            }
            x = (x << 7) | (b as usize & 0x7f);
            if b & 0x80 != 0 {
                break;
            }
        }
        Ok((input, x))
    }
}

fn load_size(input: &[u8]) -> IResult<&[u8], usize> {
    load_unsigned(!0).parse(input)
}

fn lua_int(input: &[u8]) -> IResult<&[u8], u64> {
    map(load_unsigned(i32::MAX as _), |x| x as u64)(input)
}

fn load_string(input: &[u8]) -> IResult<&[u8], &[u8]> {
    let (input, n) = load_size(input)?;
    if n == 0 {
        return Ok((input, &[]));
    }
    context("string", take(n - 1))(input)
}

fn load_upvalue(input: &[u8]) -> IResult<&[u8], UpVal> {
    map(tuple((le_u8, le_u8, le_u8)), |(on_stack, id, kind)| UpVal {
        on_stack: on_stack != 0,
        id,
        kind,
    })(input)
}

fn lua_local<'a>(_header: &LuaHeader) -> impl Parser<&'a [u8], LuaLocal, ErrorTree<&'a [u8]>> {
    tuple((load_string, lua_int, lua_int))
        .map(|(name, start_pc, end_pc)| LuaLocal {
            name: String::from_utf8_lossy(name).into(),
            start_pc,
            end_pc,
        })
        .context("local")
}

pub fn lua_chunk<'h, 'a: 'h>(
    header: &'h LuaHeader,
) -> impl Parser<&'a [u8], LuaChunk, ErrorTree<&'a [u8]>> + 'h {
    |input| {
        let (input, (name, line_defined, last_line_defined, num_params, is_vararg, max_stack)) =
            context(
                "chunk header",
                tuple((load_string, lua_int, lua_int, be_u8, be_u8, be_u8)),
            )(input)?;
        log::trace!(
            "chunk: {}, line: {line_defined}-{last_line_defined}",
            String::from_utf8_lossy(name)
        );

        map(
            tuple((
                length_count(lua_int.map(|x| x as usize), complete::u32(header.endian()))
                    .context("count instruction"),
                length_count(
                    lua_int.map(|x| x as usize),
                    alt((
                        take_lv_nil,
                        take_lv_false,
                        take_lv_true,
                        take_lv_float,
                        take_lv_str,
                        take_lv_u64,
                    )),
                )
                .context("count constants"),
                length_count(lua_int.map(|x| x as usize), load_upvalue).context("count upvalues"),
                |i| {
                    length_count(lua_int.map(|x| x as usize), lua_chunk(header))
                        .context("count prototypes")
                        .parse(i)
                },
                length_count(lua_int.map(|x| x as usize), le_u8).context("count line info"),
                length_count(
                    lua_int.map(|x| x as usize),
                    tuple((lua_int, lua_int)).map(|(a, b)| (a as u32, b as u32)),
                )
                .context("count source lines"),
                length_count(lua_int.map(|x| x as usize), lua_local(header))
                    .context("count locals"),
                length_count(lua_int.map(|x| x as usize), load_string.map(|v| v.to_vec()))
                    .context("count upval names"),
            )),
            move |(
                instructions,
                constants,
                upvalues,
                prototypes,
                _line_info,
                source_lines,
                locals,
                upvalue_names,
            )| {
                LuaChunk {
                    name: name.to_vec(),
                    line_defined,
                    last_line_defined,
                    num_upvalues: upvalues.len() as _,
                    flags: 0,
                    num_params,
                    is_vararg: if (is_vararg & 2) != 0 {
                        Some(LuaVarArgInfo {
                            has_arg: (is_vararg & 1) != 0,
                            needs_arg: (is_vararg & 4) != 0,
                        })
                    } else {
                        None
                    },
                    max_stack,
                    instructions,
                    constants,
                    prototypes,
                    source_lines,
                    locals,
                    upvalue_names,
                    num_constants: vec![],
                    upvalue_infos: upvalues,
                }
            },
        )
        .context("chunk")
        .parse(input)
    }
}

fn take_lv_nil(input: &[u8]) -> IResult<&[u8], LuaConstant> {
    let (input, _) = tag(b"\0")(input)?;
    Ok((input, LuaConstant::Null))
}

fn take_lv_false(input: &[u8]) -> IResult<&[u8], LuaConstant> {
    let (input, _) = tag(b"\x01")(input)?;
    Ok((input, LuaConstant::Bool(false)))
}

fn take_lv_true(input: &[u8]) -> IResult<&[u8], LuaConstant> {
    let (input, _) = tag(b"\x11")(input)?;
    Ok((input, LuaConstant::Bool(true)))
}

fn take_lv_float(input: &[u8]) -> IResult<&[u8], LuaConstant> {
    let (input, (_, f)) = tuple((tag(b"\x13"), complete::le_f64))(input)?;
    Ok((input, LuaConstant::Number(LuaNumber::Float(f))))
}

fn take_lv_str(input: &[u8]) -> IResult<&[u8], LuaConstant> {
    let (input, (_, data)) = tuple((alt((tag(b"\x04"), tag("\x14"))), load_string))(input)?;
    Ok((input, LuaConstant::String(data.to_vec().into())))
}

fn take_lv_u64(input: &[u8]) -> IResult<&[u8], LuaConstant> {
    let (input, (_, val)) = tuple((tag(b"\x03"), complete::le_u64))(input)?;
    Ok((input, LuaConstant::Number(LuaNumber::Integer(val as _))))
}