repugnant_pickle/
eval.rs

1use crate::{ops::*, value::*};
2
3use std::{
4    borrow::Cow,
5    collections::BTreeMap,
6    ops::{Deref, DerefMut},
7};
8
9use anyhow::{anyhow, bail, ensure, Ok, Result};
10
11const MAX_DEPTH: usize = 250;
12const MAX_PROTOCOL: u8 = 5;
13
14#[derive(Debug, Clone, PartialEq, Default)]
15/// Basically just a Vec with some convenience functions.
16pub struct PickleStack<'a>(pub Vec<Value<'a>>);
17
18impl<'a> PickleStack<'a> {
19    pub fn pop(&mut self) -> Result<Value<'a>> {
20        self.0.pop().ok_or_else(|| anyhow!("Stack underrun"))
21    }
22
23    pub fn pop_mark(&mut self) -> Result<Vec<Value<'a>>> {
24        let markidx = self.find_mark()?;
25        let postmark = self.0[markidx + 1..].to_owned();
26        self.truncate(markidx);
27        Ok(postmark)
28    }
29
30    pub fn find_mark(&self) -> Result<usize> {
31        Ok(self
32            .0
33            .iter()
34            .enumerate()
35            .rfind(|(_idx, op)| matches!(op, Value::Raw(Cow::Borrowed(&PickleOp::MARK))))
36            .map(|(idx, _)| idx)
37            .ok_or_else(|| anyhow!("Missing MARK"))?)
38    }
39}
40
41impl<'a> Deref for PickleStack<'a> {
42    type Target = Vec<Value<'a>>;
43
44    fn deref(&self) -> &Self::Target {
45        &self.0
46    }
47}
48
49impl<'a> DerefMut for PickleStack<'a> {
50    fn deref_mut(&mut self) -> &mut Self::Target {
51        &mut self.0
52    }
53}
54
55#[derive(Debug, Clone, PartialEq, Default)]
56/// Basically just a BTreeMap with some convenience functions.
57pub struct PickleMemo<'a>(pub BTreeMap<u32, Value<'a>>);
58
59impl<'a> PickleMemo<'a> {
60    /// Resolve a Value that could be a reference, but it doesn't look instead
61    /// other types. For example if you have `Ref(Ref(Ref(Ref(whatever))))`,
62    /// you'll get `whatever` back. But if you have Something(Ref(whatever))
63    /// it won't do anything.
64    pub fn resolve(&self, mut op: Value<'a>, recursive: bool) -> Result<Value<'a>> {
65        let mut count = 0;
66        while let Value::Ref(ref mid) = op {
67            let val = self.0.get(mid).ok_or_else(|| anyhow!("Bad memo id"))?;
68            if !recursive {
69                return Ok(val.to_owned());
70            }
71            op = val.to_owned();
72            count += 1;
73            if count >= MAX_DEPTH {
74                // It be how it be.
75                break;
76            }
77        }
78
79        Ok(op)
80    }
81
82    pub fn insert(&mut self, mid: u32, val: Value<'a>) {
83        self.0.insert(mid, val);
84    }
85
86    /// Like `resolve` but you get a mutable reference.
87    pub fn resolve_mut<'b, 'c>(
88        &'c mut self,
89        op: &'b mut Value<'a>,
90        recursive: bool,
91    ) -> Result<&'c mut Value<'a>>
92    where
93        'b: 'c,
94        'c: 'b,
95    {
96        let mut lastmid = if let Value::Ref(ref mid) = op {
97            *mid
98        } else {
99            return Ok(op);
100        };
101
102        let mut count = 0;
103        while let Value::Ref(mid) = self.0.get(&lastmid).ok_or_else(|| anyhow!("Bad memo id"))? {
104            lastmid = *mid;
105            if !recursive {
106                break;
107            }
108            count += 1;
109            if count >= MAX_DEPTH {
110                // We did our best but it fell short.
111                break;
112            }
113        }
114        // This unwrap is safe since it's impossible to get here without looking it up
115        // non-mut style first.
116        Ok(self
117            .0
118            .get_mut(&lastmid)
119            .expect("Impossible: Missing memo id"))
120    }
121
122    /// Try to resolve all references in an iterable of
123    /// Value. If `fix_values` is true then it will
124    /// try to fixup the values.
125    pub fn resolve_all_refs_iter(
126        &self,
127        depth: usize,
128        vals: impl IntoIterator<Item = Value<'a>>,
129        fix_values: bool,
130    ) -> Result<Vec<Value<'a>>> {
131        if depth >= MAX_DEPTH {
132            // Sometimes things just don't work out the way you hoped.
133            return Ok(vals.into_iter().collect::<Vec<_>>());
134        }
135        vals.into_iter()
136            .map(|val| self.resolve_all_refs(depth + 1, val, fix_values))
137            .collect::<Result<Vec<_>>>()
138    }
139
140    /// Try to resolve all references.
141    /// If `fix_values` is true, it will try to fixup
142    /// the values.
143    pub fn resolve_all_refs(
144        &self,
145        depth: usize,
146        val: Value<'a>,
147        fix_values: bool,
148    ) -> Result<Value<'a>> {
149        if depth >= MAX_DEPTH {
150            // It be how it be.
151            return Ok(val);
152        }
153        let rar = |v| self.resolve_all_refs(depth + 1, v, fix_values);
154        let rir = |i| self.resolve_all_refs_iter(depth + 1, i, fix_values);
155
156        let output = match val {
157            val @ Value::Ref(_) => rar(self.resolve(val, true)?)?,
158            Value::App(apped, apps) => Value::App(Box::new(rar(*apped)?), rir(apps)?),
159            Value::Object(apped, apps) => Value::Object(Box::new(rar(*apped)?), rir(apps)?),
160            Value::Build(apped, apps) => {
161                Value::Build(Box::new(rar(*apped)?), Box::new(rar(*apps)?))
162            }
163            Value::Global(target, args) => Value::Global(Box::new(rar(*target)?), rir(args)?),
164            Value::Seq(st, args) => Value::Seq(st, rir(args)?),
165            Value::PersId(pid) => Value::PersId(Box::new(rar(*pid)?)),
166            val => val,
167        };
168        if fix_values {
169            fix_value(output)
170        } else {
171            Ok(output)
172        }
173    }
174}
175
176/// Evaluate a slice of pickle ops and try to produce a Vec of
177/// Values. You'll also get the memo map back in case you
178/// need a way to look up references this crate couldn't handle.
179/// You can also pass `resolve_refs` as false and handle
180/// the references yourself.
181pub fn evaluate<'a>(
182    x: &'a [PickleOp],
183    resolve_refs: bool,
184) -> Result<(Vec<Value<'a>>, PickleMemo<'a>)> {
185    let mut stack = PickleStack::default();
186    let mut memo = PickleMemo::default();
187
188    fn make_kvlist(items: Vec<Value<'_>>) -> Result<Vec<Value<'_>>> {
189        ensure!(items.len() & 1 == 0, "Bad value for setitems");
190        let mut kvitems = Vec::with_capacity(items.len());
191        let mut it = items.into_iter();
192        while let Some(k) = it.next() {
193            let v = it.next().expect("Impossible: Missing value item");
194            kvitems.push(Value::Seq(SequenceType::Tuple, vec![k, v]));
195        }
196        Ok(kvitems)
197    }
198
199    for op in x.iter() {
200        let stack = &mut stack;
201
202        match op {
203            PickleOp::MARK => stack.push(Value::Raw(Cow::Borrowed(op))),
204            PickleOp::STOP => break,
205            PickleOp::POP => {
206                let _ = stack.pop()?;
207            }
208            PickleOp::POP_MARK => {
209                let _ = stack.pop_mark()?;
210            }
211            PickleOp::DUP => {
212                let item = stack
213                    .last()
214                    .ok_or_else(|| anyhow!("Cannot DUP with empty stack"))?
215                    .to_owned();
216                stack.push(item);
217            }
218            PickleOp::PERSID(pid) => stack.push(Value::PersId(Box::new(Value::String(pid)))),
219            PickleOp::BINPERSID => {
220                let pid = stack.pop()?;
221                stack.push(Value::PersId(Box::new(pid)));
222            }
223            PickleOp::REDUCE => {
224                let args = memo.resolve(stack.pop()?, true)?;
225                let target = memo.resolve(stack.pop()?, true)?;
226                stack.push(Value::Global(Box::new(target), vec![args]));
227            }
228            PickleOp::BUILD => {
229                let args = Box::new(memo.resolve(stack.pop()?, true)?);
230                let target = Box::new(memo.resolve(stack.pop()?, true)?);
231                stack.push(Value::Build(target, args));
232            }
233            PickleOp::EMPTY_DICT => stack.push(Value::Seq(SequenceType::Dict, Default::default())),
234            PickleOp::GET(mids) => stack.push(Value::Ref(mids.parse()?)),
235            PickleOp::BINGET(mid) => stack.push(Value::Ref(*mid as u32)),
236            PickleOp::LONG_BINGET(mid) => stack.push(Value::Ref(*mid)),
237            PickleOp::EMPTY_LIST => stack.push(Value::Seq(SequenceType::List, Default::default())),
238            PickleOp::BINPUT(mid) => {
239                let mid = *mid as u32;
240                memo.insert(mid, stack.pop()?);
241                stack.push(Value::Ref(mid));
242            }
243            PickleOp::LONG_BINPUT(mid) => {
244                memo.insert(*mid, stack.pop()?);
245                stack.push(Value::Ref(*mid));
246            }
247            PickleOp::TUPLE => {
248                let postmark = stack.pop_mark()?;
249                stack.push(Value::Seq(SequenceType::Tuple, postmark));
250            }
251            PickleOp::EMPTY_TUPLE => {
252                stack.push(Value::Seq(SequenceType::Tuple, Default::default()))
253            }
254            PickleOp::SETITEM => {
255                let v = stack.pop()?;
256                let k = stack.pop()?;
257                let top = stack
258                    .last_mut()
259                    .ok_or_else(|| anyhow!("Unexpected empty stack"))?;
260                let rtop = memo.resolve_mut(top, true)?;
261                match rtop {
262                    Value::Global(_, args) | Value::Seq(_, args) => {
263                        args.push(Value::Seq(SequenceType::Tuple, vec![k, v]));
264                    }
265                    _wut => bail!("Bad stack top for SETITEM!"),
266                }
267            }
268            PickleOp::SETITEMS => {
269                let kvitems = make_kvlist(stack.pop_mark()?)?;
270                let top = stack
271                    .last_mut()
272                    .ok_or_else(|| anyhow!("Unexpected empty stack"))?;
273                let rtop = memo.resolve_mut(top, true)?;
274                match rtop {
275                    Value::Global(_, args) | Value::Seq(_, args) => {
276                        args.push(Value::Seq(SequenceType::Tuple, kvitems));
277                    }
278                    _wut => bail!("Bad stack top for SETITEMS"),
279                }
280            }
281            PickleOp::PROTO(proto) => {
282                ensure!(*proto <= MAX_PROTOCOL, "Unsupported protocol {proto}")
283            }
284            PickleOp::TUPLE1 => {
285                let t1 = stack.pop()?;
286                stack.push(Value::Seq(SequenceType::Tuple, vec![t1]));
287            }
288            PickleOp::TUPLE2 => {
289                let (t2, t1) = (stack.pop()?, stack.pop()?);
290                stack.push(Value::Seq(SequenceType::Tuple, vec![t1, t2]));
291            }
292            PickleOp::TUPLE3 => {
293                let (t3, t2, t1) = (stack.pop()?, stack.pop()?, stack.pop()?);
294                stack.push(Value::Seq(SequenceType::Tuple, vec![t1, t2, t3]));
295            }
296            PickleOp::APPEND => {
297                let v = stack.pop()?;
298                let top = stack
299                    .last_mut()
300                    .ok_or_else(|| anyhow!("Unexpected empty stack"))?;
301                let rtop = memo.resolve_mut(top, true)?;
302                match rtop {
303                    Value::Global(_, args) | Value::Seq(_, args) => {
304                        args.push(v);
305                    }
306                    _wut => bail!("Bad stack top for APPEND!"),
307                }
308            }
309            PickleOp::APPENDS => {
310                let postmark = stack.pop_mark()?;
311                let top = stack
312                    .last_mut()
313                    .ok_or_else(|| anyhow!("Unexpected empty stack"))?;
314                let rtop = memo.resolve_mut(top, true)?;
315
316                match rtop {
317                    Value::Global(_, args) | Value::Seq(_, args) => {
318                        args.extend(postmark);
319                    }
320                    _wut => bail!("Bad stack top for APPENDS"),
321                }
322            }
323            PickleOp::DICT => {
324                let kvitems = make_kvlist(stack.pop_mark()?)?;
325                stack.push(Value::Seq(SequenceType::Dict, kvitems));
326            }
327            PickleOp::LIST => {
328                let items = stack.pop_mark()?;
329                stack.push(Value::Seq(SequenceType::List, items));
330            }
331            PickleOp::INST(mn, cn) => {
332                let args = stack.pop_mark()?;
333                stack.push(Value::Object(
334                    Box::new(Value::Seq(
335                        SequenceType::Tuple,
336                        vec![Value::String(mn), Value::String(cn)],
337                    )),
338                    args,
339                ))
340            }
341            PickleOp::OBJ => {
342                let markidx = stack.find_mark()?;
343                let args = stack.0[markidx + 2..].to_owned();
344                let cls = stack.0[markidx + 1].clone();
345                stack.0.truncate(markidx);
346                stack.push(Value::Object(Box::new(cls), args));
347            }
348            PickleOp::PUT(midstr) => {
349                // Note: This technically incorrect since the memo id could actually be a string, but
350                // it doesn't seem like that happens in practice.
351                let mid = midstr.parse()?;
352                memo.insert(mid, stack.pop()?);
353                stack.push(Value::Ref(mid));
354            }
355            PickleOp::NEWOBJ => {
356                let (args, cls) = (stack.pop()?, stack.pop()?);
357                stack.push(Value::Object(Box::new(cls), vec![args]))
358            }
359            PickleOp::EMPTY_SET => stack.push(Value::Seq(SequenceType::Set, vec![])),
360            PickleOp::ADDITEMS => {
361                let postmark = stack.pop_mark()?;
362                let top = stack
363                    .last_mut()
364                    .ok_or_else(|| anyhow!("Unexpected empty stack"))?;
365                let rtop = memo.resolve_mut(top, true)?;
366
367                match rtop {
368                    Value::Global(_, args) | Value::Seq(_, args) => {
369                        args.extend(postmark);
370                    }
371                    _wut => bail!("Bad stack top for ADDITEMS"),
372                }
373            }
374            PickleOp::FROZENSET => {
375                let items = stack.pop_mark()?;
376                stack.push(Value::Seq(SequenceType::FrozenSet, items));
377            }
378            PickleOp::NEWOBJ_EX => {
379                let (kwargs, args, cls) = (stack.pop()?, stack.pop()?, stack.pop()?);
380                stack.push(Value::Object(
381                    Box::new(cls),
382                    vec![Value::Seq(SequenceType::Tuple, vec![args, kwargs])],
383                ))
384            }
385            PickleOp::STACK_GLOBAL => {
386                let (gn, mn) = (
387                    memo.resolve(stack.pop()?, true)?,
388                    memo.resolve(stack.pop()?, true)?,
389                );
390                stack.push(Value::Global(
391                    Box::new(Value::Seq(SequenceType::Tuple, vec![gn, mn])),
392                    vec![],
393                ));
394            }
395            PickleOp::MEMOIZE => {
396                let item = stack.last().ok_or_else(|| anyhow!("Stack underrun"))?;
397                memo.insert(memo.0.len() as u32, item.to_owned());
398            }
399
400            // Fallthrough case is just to push the op onto the stack as a Value::Raw.
401            op => stack.push(Value::Raw(Cow::Borrowed(op))),
402        }
403    }
404    if !resolve_refs {
405        return Ok((stack.0, memo));
406    }
407    let stack = memo.resolve_all_refs_iter(0, stack.0, true)?;
408
409    Ok((stack, memo))
410}
repugnant_pickle/eval.rs

repugnant_pickle/
eval.rs
