write: add StringTable::new_in_order

Use this for write::elf::SinglePhaseWriter

src/write/elf/writer.rs

@@ -38,6 +38,7 @@ mod private {
 
     #[allow(private_interfaces)]
     pub trait ModeSealed {
+        fn string_table() -> StringTable<'static>;
         fn reserve(&self, buffer: &mut dyn WritableBuffer) -> Result<()>;
         fn need_offset(offset: u64) -> bool;
         fn set_offset(dest: &mut u64, offset: u64);
@@ -171,13 +172,13 @@ impl<'a, M: Mode> Writer<'a, M> {
             written_segment_num: 0,
             written_section_num: 0,
 
-            shstrtab: StringTable::default(),
+            shstrtab: M::string_table(),
             shstrtab_str_id: None,
             shstrtab_offset: 0,
             shstrtab_size: 0,
 
             need_strtab: false,
-            strtab: StringTable::default(),
+            strtab: M::string_table(),
             strtab_str_id: None,
             strtab_index: SectionIndex(0),
             strtab_offset: 0,
@@ -195,7 +196,7 @@ impl<'a, M: Mode> Writer<'a, M> {
             symtab_shndx_data: Vec::new(),
 
             need_dynstr: false,
-            dynstr: StringTable::default(),
+            dynstr: M::string_table(),
             dynstr_str_id: None,
             dynstr_index: SectionIndex(0),
             dynstr_offset: 0,
@@ -1620,9 +1621,12 @@ pub type SinglePhaseWriter<'a> = Writer<'a, SinglePhase>;
 ///
 /// Items must be written before they are referenced:
 /// - write section data before section headers
-/// - write string tables before they are referenced
 /// - write metadata section headers before they are referenced (e.g. `.strtab` before `.symtab`)
 ///
+/// Strings in strings tables are written in the order they are added, without suffix merging.
+/// This means that string offsets are known before the string table is written.
+/// For example, [`Writer::write_symbol`] may be called before [`Writer::write_strtab`].
+///
 /// The one exception is the file header, which must be written first, but needs to
 /// reference the program header table and section header table. To support this, you can
 /// call [`Writer::write_file_header`] to write placeholders, then after everything
@@ -1640,6 +1644,10 @@ pub struct SinglePhase(());
 impl Mode for SinglePhase {}
 #[allow(private_interfaces)]
 impl ModeSealed for SinglePhase {
+    fn string_table() -> StringTable<'static> {
+        StringTable::new_in_order(1)
+    }
+
     fn reserve(&self, _buffer: &mut dyn WritableBuffer) -> Result<()> {
         Ok(())
     }
@@ -1871,6 +1879,10 @@ pub struct TwoPhase {
 impl Mode for TwoPhase {}
 #[allow(private_interfaces)]
 impl ModeSealed for TwoPhase {
+    fn string_table() -> StringTable<'static> {
+        StringTable::new()
+    }
+
     fn reserve(&self, buffer: &mut dyn WritableBuffer) -> Result<()> {
         buffer
             .reserve(self.len)

src/write/string.rs

@@ -16,10 +16,36 @@ pub struct StringId(usize);
 pub(crate) struct StringTable<'a> {
     strings: IndexSet<&'a [u8]>,
     offsets: Vec<u32>,
+    size: u64,
+    in_order: bool,
+    // Only set for in_order.
+    base: u32,
 }
 
 impl<'a> StringTable<'a> {
+    /// Construct a new string table.
+    #[allow(dead_code)]
+    pub(crate) fn new() -> Self {
+        StringTable::default()
+    }
+
+    /// Construct a new string table that writes the strings in
+    /// the order they are added.
+    ///
+    /// This does not perform suffix merging.
+    #[allow(dead_code)]
+    pub(crate) fn new_in_order(base: u32) -> Self {
+        StringTable {
+            strings: IndexSet::default(),
+            offsets: Vec::new(),
+            size: base.into(),
+            in_order: true,
+            base,
+        }
+    }
+
     /// Return true if the string table contains no strings.
+    #[allow(dead_code)]
     pub(crate) fn is_empty(&self) -> bool {
         self.strings.is_empty()
     }
@@ -33,9 +59,13 @@ impl<'a> StringTable<'a> {
     /// The string must not contain a null byte; this is asserted here for
     /// debug builds, and checked by `write` for all builds.
     pub(crate) fn add(&mut self, string: &'a [u8]) -> StringId {
-        debug_assert!(self.offsets.is_empty());
+        debug_assert!(self.in_order || self.offsets.is_empty());
         debug_assert!(!string.contains(&0));
-        let id = self.strings.insert_full(string).0;
+        let (id, new) = self.strings.insert_full(string);
+        if new && self.in_order {
+            self.offsets.push(self.size as u32);
+            self.size += string.len() as u64 + 1;
+        }
         StringId(id)
     }
 
@@ -58,7 +88,11 @@ impl<'a> StringTable<'a> {
 
     /// Return the offset of the given string.
     ///
-    /// Must be called after [`Self::write`].
+    /// Must be called after [`Self::write`] unless the string table was
+    /// constructed with [`Self::new_in_order`].
+    ///
+    /// When using `new_in_order`, the offset returned will be truncated if it
+    /// overflows `u32`. Reporting the overflow is postponed to `write`.
     ///
     /// Panics if `id` is invalid or `write` was not called.
     pub(crate) fn get_offset(&self, id: StringId) -> u32 {
@@ -79,13 +113,23 @@ impl<'a> StringTable<'a> {
     /// - any string contains a null byte
     /// - the string table size is > `u32::MAX`
     pub(crate) fn write(&mut self, base: u32, w: &mut Vec<u8>) -> Result<u32> {
-        if !self.offsets.is_empty() {
+        if !self.in_order && !self.offsets.is_empty() {
             return Err(Error("string table already written".into()));
         }
         if self.strings.iter().any(|s| s.contains(&0)) {
             return Err(Error("string table entry contains null byte".into()));
         }
 
+        if self.in_order {
+            debug_assert_eq!(self.base, base);
+            for string in &self.strings {
+                w.extend_from_slice(string);
+                w.push(0);
+            }
+            return u32::try_from(self.size)
+                .map_err(|_| Error("string table size overflow".into()));
+        }
+
         let mut ids: Vec<_> = (0..self.strings.len()).collect();
         sort(&mut ids, 1, &self.strings);
 
@@ -115,6 +159,11 @@ impl<'a> StringTable<'a> {
     /// null byte.
     #[allow(dead_code)]
     pub(crate) fn size(&self, base: usize) -> usize {
+        if self.in_order {
+            debug_assert_eq!(self.base as usize, base);
+            return self.size as usize;
+        }
+
         // TODO: cache this result?
         let mut ids: Vec<_> = (0..self.strings.len()).collect();
         sort(&mut ids, 1, &self.strings);
@@ -198,8 +247,7 @@ mod tests {
         let id2 = table.add(b"bar");
         let id3 = table.add(b"foobar");
 
-        let mut data = Vec::new();
-        data.push(0);
+        let mut data = vec![0];
         assert_eq!(table.write(1, &mut data), Ok(12));
         assert_eq!(data, b"\0foobar\0foo\0");
 
@@ -211,4 +259,27 @@ mod tests {
         let mut data = Vec::new();
         assert!(table.write(1, &mut data).is_err());
     }
+
+    #[test]
+    fn string_table_in_order() {
+        let mut table = StringTable::new_in_order(1);
+        let id0 = table.add(b"");
+        let id1 = table.add(b"foo");
+        let id2 = table.add(b"bar");
+        let id3 = table.add(b"foobar");
+
+        let mut data = vec![0];
+        assert_eq!(table.write(1, &mut data), Ok(17));
+        assert_eq!(data, b"\0\0foo\0bar\0foobar\0");
+
+        assert_eq!(table.get_offset(id0), 1);
+        assert_eq!(table.get_offset(id1), 2);
+        assert_eq!(table.get_offset(id2), 6);
+        assert_eq!(table.get_offset(id3), 10);
+
+        // Not documented or expected to be needed, but multiple writes aren't prevented.
+        let mut data2 = vec![0];
+        assert_eq!(table.write(1, &mut data2), Ok(17));
+        assert_eq!(data, data2);
+    }
 }