char/arbrisseau

use wasmi to download and execute arborium grammars

git clone https://git.t4t.associates/char/arbrisseau

Charlotte Somheavy simplify; throw out initialized flag + runtime heap impls + gc59ec5cd

main
branchesmainall refs →
13.7 KiB382 linesraw
1use std::collections::BTreeMap;
2use std::sync::OnceLock;
3
4use arborium_highlight::{Injection, ParseResult, Span};
5use wasmi::{
6    Caller, CompilationMode, Config, Engine, ExternType, Instance, Linker, Memory, Module, Store,
7    StoreLimits, StoreLimitsBuilder, WasmParams, WasmResults,
8};
9
10const MAX_MEMORY: usize = 128 << 20;
11const WASM_PAGE_SIZE: usize = 1 << 16;
12/// wasm-bindgen's heap keeps its first slots permanently populated: 1024..1028
13/// are `undefined`, `null`, `true` and `false`, and everything below is padding.
14const RESERVED_HANDLES: usize = 1028;
15
16type Result<T> = std::result::Result<T, wasmi::Error>;
17
18fn wasm_error(message: impl Into<String>) -> wasmi::Error {
19    wasmi::Error::new(message)
20}
21
22/// The JS values serde-wasm-bindgen builds while serializing a parse result.
23#[derive(Clone, Debug)]
24enum Value {
25    Undefined,
26    Null,
27    Bool(bool),
28    Number(f64),
29    String(String),
30    Array(Vec<Value>),
31    Object(BTreeMap<String, Value>),
32}
33
34impl Value {
35    fn field(&self, name: &str) -> Result<&Value> {
36        let Value::Object(fields) = self else {
37            return Err(wasm_error("expected object from grammar plugin"));
38        };
39        fields
40            .get(name)
41            .ok_or_else(|| wasm_error(format!("grammar result is missing {name:?}")))
42    }
43
44    fn array(&self) -> Result<&[Value]> {
45        let Value::Array(values) = self else {
46            return Err(wasm_error("expected array from grammar plugin"));
47        };
48        Ok(values)
49    }
50
51    fn text(&self) -> Result<String> {
52        let Value::String(value) = self else {
53            return Err(wasm_error("expected string from grammar plugin"));
54        };
55        Ok(value.clone())
56    }
57
58    fn integer(&self) -> Result<u32> {
59        let Value::Number(value) = *self else {
60            return Err(wasm_error("expected number from grammar plugin"));
61        };
62        if !value.is_finite() || value.fract() != 0.0 || !(0.0..=u32::MAX as f64).contains(&value) {
63            return Err(wasm_error("invalid integer from grammar plugin"));
64        }
65        Ok(value as u32)
66    }
67
68    fn boolean(&self) -> Result<bool> {
69        let Value::Bool(value) = *self else {
70            return Err(wasm_error("expected boolean from grammar plugin"));
71        };
72        Ok(value)
73    }
74
75    fn parse_result(&self) -> Result<ParseResult> {
76        let spans = self
77            .field("spans")?
78            .array()?
79            .iter()
80            .map(|span| {
81                Ok(Span {
82                    start: span.field("start")?.integer()?,
83                    end: span.field("end")?.integer()?,
84                    capture: span.field("capture")?.text()?,
85                    pattern_index: span.field("pattern_index")?.integer()?,
86                })
87            })
88            .collect::<Result<_>>()?;
89        let injections = self
90            .field("injections")?
91            .array()?
92            .iter()
93            .map(|injection| {
94                Ok(Injection {
95                    start: injection.field("start")?.integer()?,
96                    end: injection.field("end")?.integer()?,
97                    language: injection.field("language")?.text()?,
98                    include_children: injection.field("include_children")?.boolean()?,
99                })
100            })
101            .collect::<Result<_>>()?;
102        Ok(ParseResult { spans, injections })
103    }
104}
105
106/// Stands in for wasm-bindgen's JS-side object heap. Serialization builds each
107/// value bottom-up and moves it into its parent, so a slot holds its value
108/// until the guest drops or consumes the handle. The guest caches field-name
109/// strings across calls, so handles must survive between parses.
110struct Host {
111    slots: Vec<Option<Value>>,
112    free: Vec<usize>,
113    limits: StoreLimits,
114}
115
116impl Host {
117    fn new() -> Self {
118        let mut slots = vec![Some(Value::Undefined); RESERVED_HANDLES];
119        slots[1025] = Some(Value::Null);
120        slots[1026] = Some(Value::Bool(true));
121        slots[1027] = Some(Value::Bool(false));
122        Self {
123            slots,
124            free: Vec::new(),
125            limits: StoreLimitsBuilder::new()
126                .memory_size(MAX_MEMORY)
127                .instances(1)
128                .memories(1)
129                .tables(1)
130                .trap_on_grow_failure(true)
131                .build(),
132        }
133    }
134
135    fn add(&mut self, value: Value) -> i32 {
136        match self.free.pop() {
137            Some(index) => {
138                self.slots[index] = Some(value);
139                index as i32
140            }
141            None => {
142                self.slots.push(Some(value));
143                (self.slots.len() - 1) as i32
144            }
145        }
146    }
147
148    fn get_mut(&mut self, handle: i32) -> Result<&mut Value> {
149        self.slots
150            .get_mut(handle as usize)
151            .and_then(Option::as_mut)
152            .ok_or_else(|| wasm_error("invalid wasm-bindgen object handle"))
153    }
154
155    fn take(&mut self, handle: i32) -> Result<Value> {
156        let index = handle as usize;
157        if index < RESERVED_HANDLES {
158            return self.get_mut(handle).cloned();
159        }
160        let value = self
161            .slots
162            .get_mut(index)
163            .and_then(Option::take)
164            .ok_or_else(|| wasm_error("invalid wasm-bindgen object handle"))?;
165        self.free.push(index);
166        Ok(value)
167    }
168}
169
170fn engine() -> &'static Engine {
171    static ENGINE: OnceLock<Engine> = OnceLock::new();
172    ENGINE.get_or_init(|| {
173        let mut config = Config::default();
174        config
175            .ignore_custom_sections(true)
176            .compilation_mode(CompilationMode::Lazy);
177        Engine::new(&config)
178    })
179}
180
181fn guest_string(caller: &Caller<'_, Host>, ptr: i32, len: i32) -> Result<String> {
182    let memory = caller
183        .get_export("memory")
184        .and_then(wasmi::Extern::into_memory)
185        .ok_or_else(|| wasm_error("grammar plugin has no memory export"))?;
186    let mut bytes = vec![0; len as usize];
187    memory.read(caller, ptr as usize, &mut bytes)?;
188    String::from_utf8(bytes).map_err(|_| wasm_error("grammar plugin produced invalid UTF-8"))
189}
190
191/// Provides the imports serde-wasm-bindgen's serializer needs. The plugin also
192/// imports a deserializer's worth of JS introspection for `apply_edit`, which
193/// we never call, so anything else traps if reached. Import names carry a hash
194/// of their JS shim and so are pinned to the wasm-bindgen release that built
195/// the grammars.
196fn link(linker: &mut Linker<Host>, module: &Module) -> Result<()> {
197    for import in module.imports() {
198        let ExternType::Func(ty) = import.ty() else {
199            continue;
200        };
201        let (namespace, name) = (import.module(), import.name());
202        match name {
203            "__wbg_new_ab79df5bd7c26067" => {
204                linker.func_wrap(namespace, name, |mut c: Caller<'_, Host>| {
205                    c.data_mut().add(Value::Object(BTreeMap::new()))
206                })?
207            }
208            "__wbg_new_a70fbab9066b301f" => {
209                linker.func_wrap(namespace, name, |mut c: Caller<'_, Host>| {
210                    c.data_mut().add(Value::Array(Vec::new()))
211                })?
212            }
213            "__wbindgen_cast_0000000000000001" => {
214                linker.func_wrap(namespace, name, |mut c: Caller<'_, Host>, value: f64| {
215                    c.data_mut().add(Value::Number(value))
216                })?
217            }
218            "__wbindgen_cast_0000000000000002" => linker.func_wrap(
219                namespace,
220                name,
221                |mut c: Caller<'_, Host>, ptr: i32, len: i32| -> Result<i32> {
222                    let value = guest_string(&c, ptr, len)?;
223                    Ok(c.data_mut().add(Value::String(value)))
224                },
225            )?,
226            "__wbindgen_object_clone_ref" => linker.func_wrap(
227                namespace,
228                name,
229                |mut c: Caller<'_, Host>, handle: i32| -> Result<i32> {
230                    // Only ever used on immutable strings, so a copy is indistinguishable
231                    // from a shared reference.
232                    let host = c.data_mut();
233                    let value = host.get_mut(handle)?.clone();
234                    Ok(host.add(value))
235                },
236            )?,
237            "__wbindgen_object_drop_ref" => {
238                linker.func_wrap(namespace, name, |mut c: Caller<'_, Host>, handle: i32| {
239                    c.data_mut().take(handle).map(drop)
240                })?
241            }
242            "__wbg_set_282384002438957f" => linker.func_wrap(
243                namespace,
244                name,
245                |mut c: Caller<'_, Host>, target: i32, index: i32, value: i32| -> Result<()> {
246                    let host = c.data_mut();
247                    let value = host.take(value)?;
248                    let Value::Array(values) = host.get_mut(target)? else {
249                        return Err(wasm_error("indexed set on non-array"));
250                    };
251                    let index = index as usize;
252                    if values.len() <= index {
253                        values.resize(index + 1, Value::Undefined);
254                    }
255                    values[index] = value;
256                    Ok(())
257                },
258            )?,
259            "__wbg_set_6be42768c690e380" => linker.func_wrap(
260                namespace,
261                name,
262                |mut c: Caller<'_, Host>, target: i32, key: i32, value: i32| -> Result<()> {
263                    let host = c.data_mut();
264                    let key = host.take(key)?.text()?;
265                    let value = host.take(value)?;
266                    let Value::Object(fields) = host.get_mut(target)? else {
267                        return Err(wasm_error("property set on non-object"));
268                    };
269                    fields.insert(key, value);
270                    Ok(())
271                },
272            )?,
273            "__wbg___wbindgen_throw_6ddd609b62940d55" => linker.func_wrap(
274                namespace,
275                name,
276                |c: Caller<'_, Host>, ptr: i32, len: i32| -> Result<()> {
277                    Err(wasm_error(guest_string(&c, ptr, len)?))
278                },
279            )?,
280            _ => {
281                let message = format!("grammar plugin called unsupported import {name}");
282                linker.func_new(namespace, name, ty.clone(), move |_, _, _| {
283                    Err(wasm_error(message.clone()))
284                })?
285            }
286        };
287    }
288    Ok(())
289}
290
291fn instantiate(module: &Module) -> Result<(Store<Host>, Instance, Memory)> {
292    let mut store = Store::new(engine(), Host::new());
293    store.limiter(|host| &mut host.limits);
294    let mut linker = Linker::new(engine());
295    link(&mut linker, module)?;
296    let instance = linker.instantiate_and_start(&mut store, module)?;
297    let memory = instance
298        .get_memory(&store, "memory")
299        .ok_or_else(|| wasm_error("grammar plugin has no memory export"))?;
300    Ok((store, instance, memory))
301}
302
303/// A compiled grammar plus a snapshot of guest memory taken right after its
304/// parser session was created, since compiling the highlight queries in
305/// `create_session` is expensive and every runtime would otherwise repeat it.
306pub(crate) struct Grammar {
307    module: Module,
308    memory: Vec<u8>,
309    session: i32,
310}
311
312impl Grammar {
313    pub(crate) fn new(bytes: &[u8]) -> Result<Self> {
314        let module = Module::new(engine(), bytes)?;
315        let (mut store, instance, memory) = instantiate(&module)?;
316        let session = instance
317            .get_typed_func::<(), i32>(&store, "create_session")?
318            .call(&mut store, ())?;
319        Ok(Self {
320            memory: memory.data(&store).to_vec(),
321            module,
322            session,
323        })
324    }
325
326    pub(crate) fn runtime(&self) -> Result<Runtime> {
327        let (mut store, instance, memory) = instantiate(&self.module)?;
328        let growth = self.memory.len().saturating_sub(memory.data_size(&store));
329        memory.grow(&mut store, growth.div_ceil(WASM_PAGE_SIZE) as u64)?;
330        memory.write(&mut store, 0, &self.memory)?;
331        Ok(Runtime {
332            store,
333            instance,
334            memory,
335            session: self.session,
336        })
337    }
338}
339
340pub(crate) struct Runtime {
341    store: Store<Host>,
342    instance: Instance,
343    memory: Memory,
344    session: i32,
345}
346
347impl Runtime {
348    fn call<P: WasmParams, R: WasmResults>(&mut self, name: &str, params: P) -> Result<R> {
349        self.instance
350            .get_typed_func::<P, R>(&self.store, name)?
351            .call(&mut self.store, params)
352    }
353
354    /// Parses `source`. Errors may be traps that leave guest state arbitrary,
355    /// so the runtime should be discarded after one.
356    pub(crate) fn parse(&mut self, source: &str) -> Result<ParseResult> {
357        let source_len = i32::try_from(source.len())
358            .map_err(|_| wasm_error("source is too large to highlight"))?;
359        let ptr: i32 = self.call("__wbindgen_export", (source_len, 1))?;
360        self.memory
361            .write(&mut self.store, ptr as usize, source.as_bytes())?;
362        let () = self.call("set_text", (self.session, ptr, source_len))?;
363
364        let retptr: i32 = self.call("__wbindgen_add_to_stack_pointer", -16)?;
365        let () = self.call("parse", (retptr, self.session))?;
366        let mut returned = [0; 12];
367        self.memory
368            .read(&self.store, retptr as usize, &mut returned)?;
369        let _: i32 = self.call("__wbindgen_add_to_stack_pointer", 16)?;
370        let [value, error, is_error] =
371            std::array::from_fn(|i| i32::from_le_bytes(returned[4 * i..][..4].try_into().unwrap()));
372
373        let host = self.store.data_mut();
374        if is_error != 0 {
375            return Err(wasm_error(match host.take(error)? {
376                Value::String(message) => message,
377                other => format!("{other:?}"),
378            }));
379        }
380        host.take(value)?.parse_result()
381    }
382}