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

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