src/librustc_codegen_llvm/back/wasm.rs

   1 // Copyright 2018 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11 use std::collections::BTreeMap;
  12 use std::fs;
  13 use std::path::Path;
  14 use std::str;
  15
  16 use rustc_data_structures::fx::FxHashMap;
  17 use serialize::leb128;
  18
  19 // https://webassembly.github.io/spec/core/binary/modules.html#binary-importsec
  20 const WASM_IMPORT_SECTION_ID: u8 = 2;
  21
  22 const WASM_EXTERNAL_KIND_FUNCTION: u8 = 0;
  23 const WASM_EXTERNAL_KIND_TABLE: u8 = 1;
  24 const WASM_EXTERNAL_KIND_MEMORY: u8 = 2;
  25 const WASM_EXTERNAL_KIND_GLOBAL: u8 = 3;
  26
  27 /// Append all the custom sections listed in `sections` to the wasm binary
  28 /// specified at `path`.
  29 ///
  30 /// LLVM 6 which we're using right now doesn't have the ability to create custom
  31 /// sections in wasm files nor does LLD have the ability to merge these sections
  32 /// into one larger section when linking. It's expected that this will
  33 /// eventually get implemented, however!
  34 ///
  35 /// Until that time though this is a custom implementation in rustc to append
  36 /// all sections to a wasm file to the finished product that LLD produces.
  37 ///
  38 /// Support for this is landing in LLVM in https://reviews.llvm.org/D43097,
  39 /// although after that support will need to be in LLD as well.
  40 pub fn add_custom_sections(path: &Path, sections: &BTreeMap<String, Vec<u8>>) {
  41     if sections.len() == 0 {
  42         return
  43     }
  44
  45     let wasm = fs::read(path).expect("failed to read wasm output");
  46
  47     // see https://webassembly.github.io/spec/core/binary/modules.html#custom-section
  48     let mut wasm = WasmEncoder { data: wasm };
  49     for (section, bytes) in sections {
  50         // write the `id` identifier, 0 for a custom section
  51         wasm.byte(0);
  52
  53         // figure out how long our name descriptor will be
  54         let mut name = WasmEncoder::new();
  55         name.str(section);
  56
  57         // write the length of the payload followed by all its contents
  58         wasm.u32((bytes.len() + name.data.len()) as u32);
  59         wasm.data.extend_from_slice(&name.data);
  60         wasm.data.extend_from_slice(bytes);
  61     }
  62
  63     fs::write(path, &wasm.data).expect("failed to write wasm output");
  64 }
  65
  66 /// Rewrite the module imports are listed from in a wasm module given the field
  67 /// name to module name mapping in `import_map`.
  68 ///
  69 /// LLVM 6 which we're using right now doesn't have the ability to configure the
  70 /// module a wasm symbol is import from. Rather all imported symbols come from
  71 /// the bland `"env"` module unconditionally. Furthermore we'd *also* need
  72 /// support in LLD for preserving these import modules, which it unfortunately
  73 /// currently does not.
  74 ///
  75 /// This function is intended as a hack for now where we manually rewrite the
  76 /// wasm output by LLVM to have the correct import modules listed. The
  77 /// `#[wasm_import_module]` attribute in Rust translates to the module that each
  78 /// symbol is imported from, so here we manually go through the wasm file,
  79 /// decode it, rewrite imports, and then rewrite the wasm module.
  80 ///
  81 /// Support for this was added to LLVM in
  82 /// https://github.com/llvm-mirror/llvm/commit/0f32e1365, although support still
  83 /// needs to be added (AFAIK at the time of this writing) to LLD
  84 pub fn rewrite_imports(path: &Path, import_map: &FxHashMap<String, String>) {
  85     if import_map.len() == 0 {
  86         return
  87     }
  88
  89     let wasm = fs::read(path).expect("failed to read wasm output");
  90     let mut ret = WasmEncoder::new();
  91     ret.data.extend(&wasm[..8]);
  92
  93     // skip the 8 byte wasm/version header
  94     for (id, raw) in WasmSections(WasmDecoder::new(&wasm[8..])) {
  95         ret.byte(id);
  96         if id == WASM_IMPORT_SECTION_ID {
  97             info!("rewriting import section");
  98             let data = rewrite_import_section(
  99                 &mut WasmDecoder::new(raw),
 100                 import_map,
 101             );
 102             ret.bytes(&data);
 103         } else {
 104             info!("carry forward section {}, {} bytes long", id, raw.len());
 105             ret.bytes(raw);
 106         }
 107     }
 108
 109     fs::write(path, &ret.data).expect("failed to write wasm output");
 110
 111     fn rewrite_import_section(
 112         wasm: &mut WasmDecoder,
 113         import_map: &FxHashMap<String, String>,
 114     )
 115         -> Vec<u8>
 116     {
 117         let mut dst = WasmEncoder::new();
 118         let n = wasm.u32();
 119         dst.u32(n);
 120         info!("rewriting {} imports", n);
 121         for _ in 0..n {
 122             rewrite_import_entry(wasm, &mut dst, import_map);
 123         }
 124         return dst.data
 125     }
 126
 127     fn rewrite_import_entry(wasm: &mut WasmDecoder,
 128                             dst: &mut WasmEncoder,
 129                             import_map: &FxHashMap<String, String>) {
 130         // More info about the binary format here is available at:
 131         // https://webassembly.github.io/spec/core/binary/modules.html#import-section
 132         //
 133         // Note that you can also find the whole point of existence of this
 134         // function here, where we map the `module` name to a different one if
 135         // we've got one listed.
 136         let module = wasm.str();
 137         let field = wasm.str();
 138         let new_module = if module == "env" {
 139             import_map.get(field).map(|s| &**s).unwrap_or(module)
 140         } else {
 141             module
 142         };
 143         info!("import rewrite ({} => {}) / {}", module, new_module, field);
 144         dst.str(new_module);
 145         dst.str(field);
 146         let kind = wasm.byte();
 147         dst.byte(kind);
 148         match kind {
 149             WASM_EXTERNAL_KIND_FUNCTION => dst.u32(wasm.u32()),
 150             WASM_EXTERNAL_KIND_TABLE => {
 151                 dst.byte(wasm.byte()); // element_type
 152                 dst.limits(wasm.limits());
 153             }
 154             WASM_EXTERNAL_KIND_MEMORY => dst.limits(wasm.limits()),
 155             WASM_EXTERNAL_KIND_GLOBAL => {
 156                 dst.byte(wasm.byte()); // content_type
 157                 dst.bool(wasm.bool()); // mutable
 158             }
 159             b => panic!("unknown kind: {}", b),
 160         }
 161     }
 162 }
 163
 164 struct WasmSections<'a>(WasmDecoder<'a>);
 165
 166 impl<'a> Iterator for WasmSections<'a> {
 167     type Item = (u8, &'a [u8]);
 168
 169     fn next(&mut self) -> Option<(u8, &'a [u8])> {
 170         if self.0.data.len() == 0 {
 171             return None
 172         }
 173
 174         // see https://webassembly.github.io/spec/core/binary/modules.html#sections
 175         let id = self.0.byte();
 176         let section_len = self.0.u32();
 177         info!("new section {} / {} bytes", id, section_len);
 178         let section = self.0.skip(section_len as usize);
 179         Some((id, section))
 180     }
 181 }
 182
 183 struct WasmDecoder<'a> {
 184     data: &'a [u8],
 185 }
 186
 187 impl<'a> WasmDecoder<'a> {
 188     fn new(data: &'a [u8]) -> WasmDecoder<'a> {
 189         WasmDecoder { data }
 190     }
 191
 192     fn byte(&mut self) -> u8 {
 193         self.skip(1)[0]
 194     }
 195
 196     fn u32(&mut self) -> u32 {
 197         let (n, l1) = leb128::read_u32_leb128(self.data);
 198         self.data = &self.data[l1..];
 199         return n
 200     }
 201
 202     fn skip(&mut self, amt: usize) -> &'a [u8] {
 203         let (data, rest) = self.data.split_at(amt);
 204         self.data = rest;
 205         data
 206     }
 207
 208     fn str(&mut self) -> &'a str {
 209         let len = self.u32();
 210         str::from_utf8(self.skip(len as usize)).unwrap()
 211     }
 212
 213     fn bool(&mut self) -> bool {
 214         self.byte() == 1
 215     }
 216
 217     fn limits(&mut self) -> (u32, Option<u32>) {
 218         let has_max = self.bool();
 219         (self.u32(), if has_max { Some(self.u32()) } else { None })
 220     }
 221 }
 222
 223 struct WasmEncoder {
 224     data: Vec<u8>,
 225 }
 226
 227 impl WasmEncoder {
 228     fn new() -> WasmEncoder {
 229         WasmEncoder { data: Vec::new() }
 230     }
 231
 232     fn u32(&mut self, val: u32) {
 233         let at = self.data.len();
 234         leb128::write_u32_leb128(&mut self.data, at, val);
 235     }
 236
 237     fn byte(&mut self, val: u8) {
 238         self.data.push(val);
 239     }
 240
 241     fn bytes(&mut self, val: &[u8]) {
 242         self.u32(val.len() as u32);
 243         self.data.extend_from_slice(val);
 244     }
 245
 246     fn str(&mut self, val: &str) {
 247         self.bytes(val.as_bytes())
 248     }
 249
 250     fn bool(&mut self, b: bool) {
 251         self.byte(b as u8);
 252     }
 253
 254     fn limits(&mut self, limits: (u32, Option<u32>)) {
 255         self.bool(limits.1.is_some());
 256         self.u32(limits.0);
 257         if let Some(c) = limits.1 {
 258             self.u32(c);
 259         }
 260     }
 261 }