vendor/object-0.20.0/src/write/string.rs

   1 use indexmap::IndexSet;
   2 use std::vec::Vec;
   3
   4 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
   5 pub(crate) struct StringId(usize);
   6
   7 #[derive(Debug, Default)]
   8 pub(crate) struct StringTable<'a> {
   9     strings: IndexSet<&'a [u8]>,
  10     offsets: Vec<usize>,
  11 }
  12
  13 impl<'a> StringTable<'a> {
  14     /// Add a string to the string table.
  15     ///
  16     /// Panics if the string table has already been written, or
  17     /// if the string contains a null byte.
  18     pub fn add(&mut self, string: &'a [u8]) -> StringId {
  19         assert!(self.offsets.is_empty());
  20         assert!(!string.contains(&0));
  21         let id = self.strings.insert_full(string).0;
  22         StringId(id)
  23     }
  24
  25     /// Return the offset of the given string.
  26     ///
  27     /// Panics if the string table has not been written, or
  28     /// if the string is not in the string table.
  29     pub fn get_offset(&self, id: StringId) -> usize {
  30         self.offsets[id.0]
  31     }
  32
  33     /// Append the string table to the given `Vec`, and
  34     /// calculate the list of string offsets.
  35     ///
  36     /// `base` is the initial string table offset. For example,
  37     /// this should be 1 for ELF, to account for the initial
  38     /// null byte (which must have been written by the caller).
  39     pub fn write(&mut self, base: usize, w: &mut Vec<u8>) {
  40         assert!(self.offsets.is_empty());
  41
  42         let mut ids: Vec<_> = (0..self.strings.len()).collect();
  43         sort(&mut ids, 1, &self.strings);
  44
  45         self.offsets = vec![0; ids.len()];
  46         let mut offset = base;
  47         let mut previous = &[][..];
  48         for id in ids {
  49             let string = self.strings.get_index(id).unwrap();
  50             if previous.ends_with(string) {
  51                 self.offsets[id] = offset - string.len() - 1;
  52             } else {
  53                 self.offsets[id] = offset;
  54                 w.extend_from_slice(string);
  55                 w.push(0);
  56                 offset += string.len() + 1;
  57                 previous = string;
  58             }
  59         }
  60     }
  61 }
  62
  63 // Multi-key quicksort.
  64 //
  65 // Ordering is such that if a string is a suffix of at least one other string,
  66 // then it is placed immediately after one of those strings. That is:
  67 // - comparison starts at the end of the string
  68 // - shorter strings come later
  69 //
  70 // Based on the implementation in LLVM.
  71 fn sort(mut ids: &mut [usize], mut pos: usize, strings: &IndexSet<&[u8]>) {
  72     loop {
  73         if ids.len() <= 1 {
  74             return;
  75         }
  76
  77         let pivot = byte(ids[0], pos, strings);
  78         let mut lower = 0;
  79         let mut upper = ids.len();
  80         let mut i = 1;
  81         while i < upper {
  82             let b = byte(ids[i], pos, strings);
  83             if b > pivot {
  84                 ids.swap(lower, i);
  85                 lower += 1;
  86                 i += 1;
  87             } else if b < pivot {
  88                 upper -= 1;
  89                 ids.swap(upper, i);
  90             } else {
  91                 i += 1;
  92             }
  93         }
  94
  95         sort(&mut ids[..lower], pos, strings);
  96         sort(&mut ids[upper..], pos, strings);
  97
  98         if pivot == 0 {
  99             return;
 100         }
 101         ids = &mut ids[lower..upper];
 102         pos += 1;
 103     }
 104 }
 105
 106 fn byte(id: usize, pos: usize, strings: &IndexSet<&[u8]>) -> u8 {
 107     let string = strings.get_index(id).unwrap();
 108     let len = string.len();
 109     if len >= pos {
 110         string[len - pos]
 111     } else {
 112         // We know the strings don't contain null bytes.
 113         0
 114     }
 115 }
 116
 117 #[cfg(test)]
 118 mod tests {
 119     use super::*;
 120
 121     #[test]
 122     fn string_table() {
 123         let mut table = StringTable::default();
 124         let id0 = table.add(b"");
 125         let id1 = table.add(b"foo");
 126         let id2 = table.add(b"bar");
 127         let id3 = table.add(b"foobar");
 128
 129         let mut data = Vec::new();
 130         data.push(0);
 131         table.write(1, &mut data);
 132         assert_eq!(data, b"\0foobar\0foo\0");
 133
 134         assert_eq!(table.get_offset(id0), 11);
 135         assert_eq!(table.get_offset(id1), 8);
 136         assert_eq!(table.get_offset(id2), 4);
 137         assert_eq!(table.get_offset(id3), 1);
 138     }
 139 }