walrus_meter/
lib.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4/*!
5# `walrus-meter`
6
7This crate provides instrumentation-based metering for Wasm modules using the
8`walrus` library.
9
10Provide the name of a function to call, and the function will be called
11immediately before executing a sequence of instructions with that sequence's
12cost.
13*/
14
15pub mod costs;
16
17#[doc(inline)]
18pub use costs::Costs;
19
20use std::collections::HashMap;
21
22/// The type of errors raised by this crate.
23pub type Error = anyhow::Error;
24
25#[derive(Default)]
26struct InfoItem {
27    // The minimum cost of entering this sequence (after any prepayment)
28    min_cost: usize,
29    // The number of enclosing blocks this sequence can escape
30    escape_velocity: usize,
31}
32
33struct Instrument<C> {
34    costs: C,
35    spend: walrus::FunctionId,
36    info: Info,
37    // Walrus IR doesn't preserve the relative block indices, so we have to
38    // reconstruct them by keeping a stack.
39    stack: Vec<(walrus::ir::InstrSeqId, InfoItem)>,
40}
41
42#[derive(Default)]
43struct Info {
44    items: HashMap<walrus::ir::InstrSeqId, InfoItem>,
45}
46
47#[facile::facade]
48impl InstrExt for walrus::ir::Instr {
49    fn jump_targets(&self) -> Vec<walrus::ir::InstrSeqId> {
50        use walrus::ir::Instr::*;
51        match self {
52            Br(walrus::ir::Br { block })
53            | BrIf(walrus::ir::BrIf { block })
54            | BrOnCast(walrus::ir::BrOnCast { block, .. })
55            | BrOnCastFail(walrus::ir::BrOnCastFail { block, .. })
56            | BrOnNull(walrus::ir::BrOnNull { block })
57            | BrOnNonNull(walrus::ir::BrOnNonNull { block }) => vec![*block],
58            BrTable(walrus::ir::BrTable { blocks, default }) => blocks
59                .iter()
60                .cloned()
61                .chain(std::iter::once(*default))
62                .collect(),
63            _ => vec![],
64        }
65    }
66
67    fn branches(&self) -> Vec<walrus::ir::InstrSeqId> {
68        use walrus::ir::Instr::*;
69
70        match self {
71            IfElse(walrus::ir::IfElse {
72                consequent,
73                alternative,
74            }) => vec![*consequent, *alternative],
75
76            Block(walrus::ir::Block { seq }) => vec![*seq],
77
78            _ => vec![],
79        }
80    }
81
82    fn entry_points(&self) -> Vec<walrus::ir::InstrSeqId> {
83        if let walrus::ir::Instr::Loop(walrus::ir::Loop { seq }) = self {
84            vec![*seq]
85        } else {
86            self.branches()
87        }
88    }
89}
90
91impl Info {
92    fn can_escape(&self, instr: &walrus::ir::Instr) -> bool {
93        !instr.jump_targets().is_empty()
94            || instr
95                .entry_points()
96                .iter()
97                .map(|seq_id| self.items.get(seq_id).unwrap().escape_velocity)
98                .max()
99                .unwrap_or_default()
100                > 1
101    }
102}
103
104struct BasicBlocks<'info, 'instrs> {
105    info: &'info Info,
106    instrs: &'instrs [(walrus::ir::Instr, walrus::InstrLocId)],
107}
108
109impl<'info, 'instrs> Iterator for BasicBlocks<'info, 'instrs> {
110    type Item = &'instrs [(walrus::ir::Instr, walrus::InstrLocId)];
111
112    fn next(&mut self) -> Option<Self::Item> {
113        if self.instrs.is_empty() {
114            None
115        } else {
116            let pos = self
117                .instrs
118                .iter()
119                .position(|(instr, _loc)| self.info.can_escape(instr))
120                .map_or(self.instrs.len(), |pos| pos + 1);
121            let ret = &self.instrs[..pos];
122            self.instrs = &self.instrs[pos..];
123            Some(ret)
124        }
125    }
126}
127
128impl<C: Costs> Instrument<C> {
129    fn new(costs: C, spend: walrus::FunctionId) -> Self {
130        Self {
131            costs,
132            spend,
133            info: Info::default(),
134            stack: vec![],
135        }
136    }
137
138    fn enter_seq(&mut self, id: walrus::ir::InstrSeqId) {
139        self.stack.push((id, InfoItem::default()))
140    }
141
142    fn exit_seq(&mut self) -> (walrus::ir::InstrSeqId, InfoItem) {
143        self.stack.pop().expect("enter should be paired with exit")
144    }
145
146    fn info(&self, seq: walrus::ir::InstrSeqId) -> &InfoItem {
147        self.info.items.get(&seq).unwrap()
148    }
149
150    fn relative_label(&self, seq: walrus::ir::InstrSeqId) -> usize {
151        self.stack
152            .iter()
153            .rev()
154            .position(|(s, _info)| *s == seq)
155            .expect("Wasm should be well-formed")
156    }
157
158    fn escape_velocity(&self, instr: &walrus::ir::Instr) -> usize {
159        if instr.is_return()
160            || instr.is_return_call_indirect()
161            || instr.is_return_call()
162            || instr.is_return_call_ref()
163        {
164            self.stack.len()
165        } else {
166            std::cmp::max(
167                instr
168                    .entry_points()
169                    .into_iter()
170                    .map(|seq| self.info(seq).escape_velocity)
171                    .max()
172                    .unwrap_or_default()
173                    .saturating_sub(1),
174                instr
175                    .jump_targets()
176                    .into_iter()
177                    .map(|seq| self.relative_label(seq) + 1)
178                    .max()
179                    .unwrap_or_default(),
180            )
181        }
182    }
183
184    fn basic_blocks<'instr>(
185        &self,
186        seq: &'instr [(walrus::ir::Instr, walrus::InstrLocId)],
187    ) -> impl Iterator<Item = &'instr [(walrus::ir::Instr, walrus::InstrLocId)]> {
188        BasicBlocks {
189            info: &self.info,
190            instrs: seq,
191        }
192    }
193
194    fn extract_min_cost(&mut self, instr: &walrus::ir::Instr) -> usize {
195        let nested_cost = instr
196            .branches()
197            .iter()
198            .map(|seq| self.info(*seq).min_cost)
199            .min()
200            .unwrap_or_default();
201        for nest in instr.branches() {
202            self.info.items.get_mut(&nest).unwrap().min_cost -= nested_cost;
203        }
204        self.costs.instruction(instr) as usize + nested_cost
205    }
206
207    fn call_spend(&self, output: &mut Vec<(walrus::ir::Instr, walrus::InstrLocId)>, cost: usize) {
208        let cost = i64::try_from(cost).unwrap();
209        if cost > 0 {
210            output.push((
211                walrus::ir::Const {
212                    value: walrus::ir::Value::I64(cost),
213                }
214                .into(),
215                walrus::InstrLocId::default(),
216            ));
217
218            output.push((
219                walrus::ir::Call { func: self.spend }.into(),
220                walrus::InstrLocId::default(),
221            ));
222        }
223    }
224}
225
226// Build up information about the function.
227impl<'instr, C: Costs> walrus::ir::Visitor<'instr> for Instrument<C> {
228    fn start_instr_seq(&mut self, seq: &'instr walrus::ir::InstrSeq) {
229        self.enter_seq(seq.id());
230    }
231
232    fn end_instr_seq(&mut self, seq: &'instr walrus::ir::InstrSeq) {
233        let escape_velocity = seq
234            .iter()
235            .map(|(instr, _loc)| self.escape_velocity(instr))
236            .max()
237            .unwrap_or_default();
238
239        let mut min_cost = 0;
240        for (instr, _loc) in self
241            .basic_blocks(seq)
242            .take(1)
243            .flatten()
244            .collect::<Vec<_>>()
245        {
246            min_cost += self.extract_min_cost(instr);
247        }
248
249        let (id, mut info) = self.exit_seq();
250        debug_assert_eq!(id, seq.id());
251
252        info.escape_velocity = escape_velocity;
253        info.min_cost = min_cost;
254
255        self.info.items.insert(id, info);
256    }
257}
258
259impl<C: Costs> walrus::ir::VisitorMut for Instrument<C> {
260    fn end_instr_seq_mut(&mut self, seq: &mut walrus::ir::InstrSeq) {
261        let seq_min_cost = self.info(seq.id()).min_cost;
262        for (index, basic_block) in self
263            .basic_blocks(&std::mem::take(&mut seq.instrs))
264            .collect::<Vec<_>>()
265            .into_iter()
266            .enumerate()
267        {
268            let cost = if index == 0 {
269                seq_min_cost
270            } else {
271                basic_block
272                    .iter()
273                    .map(|(instr, _loc)| self.extract_min_cost(instr))
274                    .sum()
275            };
276
277            self.call_spend(&mut seq.instrs, cost);
278            seq.instrs.extend_from_slice(basic_block);
279        }
280    }
281}
282
283/// Injects the named spend function and instrument all functions to call it
284/// according to the given costs.
285pub fn instrument(
286    module: &[u8],
287    costs: impl Costs,
288    (spend_module, spend_name): (&str, &str),
289) -> Result<Vec<u8>, Error> {
290    let mut module = walrus::Module::from_buffer(module)?;
291    let spend_ty = module.types.add(&[walrus::ValType::I64], &[]);
292    let (spend, _) = module.add_import_func(spend_module, spend_name, spend_ty);
293
294    for (_id, func) in module.funcs.iter_local_mut() {
295        let mut visitor = Instrument::new(&costs, spend);
296        walrus::ir::dfs_in_order(&mut visitor, func, func.entry_block());
297        walrus::ir::dfs_pre_order_mut(&mut visitor, func, func.entry_block());
298    }
299
300    Ok(module.emit_wasm())
301}
walrus_meter/lib.rs

walrus_meter/
lib.rs
