From 970bbf65bb022edf063ce9bb2b1242fbe9db194d Mon Sep 17 00:00:00 2001
From: Ales Katona <ales@katona.me>
Date: Mon, 11 Jan 2021 16:35:06 -0800
Subject: [PATCH] cleanup SegmentView

---
 src/formats/elf/output.rs  |  6 ++--
 src/formats/elf/segment.rs | 69 ++++++++++++++++++--------------------
 2 files changed, 35 insertions(+), 40 deletions(-)

diff --git a/src/formats/elf/output.rs b/src/formats/elf/output.rs
index af8c6e8..c478728 100644
--- a/src/formats/elf/output.rs
+++ b/src/formats/elf/output.rs
@@ -25,7 +25,7 @@ use super::ElfObject;
 pub struct ElfOutput {
     destination: PathBuf,
     file: ELF64,
-    input_data: OutputData,
+    input_data: SegmentView,
     writer: BufWriter<File>,
 }
 
@@ -48,7 +48,7 @@ impl ElfOutput {
         let result = Self {
             destination,
             file: elf,
-            input_data: OutputData::new(),
+            input_data: SegmentView::default(),
             writer,
         };
 
@@ -64,7 +64,7 @@ impl ElfOutput {
         let result = Self {
             destination,
             file: elf_bin_from_object(other),
-            input_data: OutputData::new(),
+            input_data: SegmentView::default(),
             writer,
         };
 
diff --git a/src/formats/elf/segment.rs b/src/formats/elf/segment.rs
index d122127..32b021a 100644
--- a/src/formats/elf/segment.rs
+++ b/src/formats/elf/segment.rs
@@ -1,35 +1,34 @@
+use std::iter::once;
+
 use crate::common::SectionInfo;
 use crate::common::{Relocatable, Section};
 use crate::error::Error;
 
 use super::ElfObject;
 
-const SI_TEXT: usize = 0;
-const SI_RODATA: usize = 1;
-const SI_DATA: usize = 2;
-const SI_BSS: usize = 3;
-
 pub enum SegmentType {
-    Text = SI_TEXT as isize,
-    Data = SI_DATA as isize,
-    Bss = SI_BSS as isize,
+    Text,
+    Data,
+    Bss,
 }
 
-const SEGMENT_NAMES: [&str; 4] = [".text", ".rodata", ".data", ".bss"];
-
-pub struct OutputData([Vec<SectionInfo<(usize, usize)>>; 4]);
+type SegmentSections = Vec<SectionInfo<(usize, usize)>>;
 
-impl<'data> OutputData {
-    pub fn new() -> Self {
-        Self([Vec::new(), Vec::new(), Vec::new(), Vec::new()])
-    }
+#[derive(Default)]
+pub struct SegmentView {
+    text: SegmentSections,
+    rodata: SegmentSections,
+    data: SegmentSections,
+    bss: SegmentSections,
+}
 
+impl<'data> SegmentView {
     pub fn append_section(&mut self, section: Section<(usize, usize)>) -> Result<(), Error> {
         match section {
-            Section::Text(si) => self.0[SI_TEXT].push(si),
-            Section::Data(si, false) => self.0[SI_DATA].push(si),
-            Section::Data(si, true) => self.0[SI_RODATA].push(si),
-            Section::Bss(si) => self.0[SI_BSS].push(si),
+            Section::Text(si) => self.text.push(si),
+            Section::Data(si, true) => self.rodata.push(si),
+            Section::Data(si, false) => self.data.push(si),
+            Section::Bss(si) => self.bss.push(si),
         }
 
         Ok(())
@@ -37,19 +36,21 @@ impl<'data> OutputData {
 
     pub fn sections_mut(
         &mut self,
-    ) -> impl Iterator<Item = (&'static str, &mut Vec<SectionInfo<(usize, usize)>>)> {
-        self.0
-            .iter_mut()
-            .enumerate()
-            .map(|(i, s)| (SEGMENT_NAMES[i], s))
+    ) -> impl Iterator<Item = (&'static str, &mut SegmentSections)> {
+        let text = once(&mut self.text).map(|s| (".text", s));
+        let rodata = once(&mut self.rodata).map(|s| (".rodata", s));
+        let data = once(&mut self.data).map(|s| (".data", s));
+        let bss = once(&mut self.bss).map(|s| (".bss", s));
+
+        text.chain(rodata).chain(data).chain(bss)
     }
 
     pub fn program_bytes<'l>(
         &'l self,
         objects: &'l Vec<ElfObject>,
     ) -> impl Iterator<Item = Result<&'l [u8], Error>> {
-        let text_iter = self.0[SI_TEXT].iter();
-        let rodata_iter = self.0[SI_RODATA].iter();
+        let text_iter = self.text.iter();
+        let rodata_iter = self.rodata.iter();
 
         let data1 = text_iter.filter_map(move |si| match si.data_index {
             None => None,
@@ -69,7 +70,7 @@ impl<'data> OutputData {
         &'l self,
         objects: &'l Vec<ElfObject>,
     ) -> impl Iterator<Item = Result<&'l [u8], Error>> {
-        let iter = self.0[SI_DATA]
+        let iter = self.data
             .iter()
             .filter_map(move |si| match si.data_index {
                 None => None,
@@ -81,8 +82,8 @@ impl<'data> OutputData {
 
     // .text + .rodata
     pub fn program_size(&self) -> u64 {
-        let text_iter = self.0[SI_TEXT].iter();
-        let rodata_iter = self.0[SI_RODATA].iter();
+        let text_iter = self.text.iter();
+        let rodata_iter = self.rodata.iter();
 
         let mut result = 0u64;
         for section in text_iter.chain(rodata_iter) {
@@ -95,7 +96,7 @@ impl<'data> OutputData {
     // data
     pub fn data_size(&self) -> u64 {
         let mut result = 0u64;
-        for section in self.0[SI_DATA].iter() {
+        for section in self.data.iter() {
             result += section.data_size
         }
 
@@ -105,16 +106,10 @@ impl<'data> OutputData {
     // bss
     pub fn bss_size(&self) -> u64 {
         let mut result = 0u64;
-        for section in self.0[SI_BSS].iter() {
+        for section in self.bss.iter() {
             result += section.data_size
         }
 
         result
     }
 }
-
-impl Default for OutputData {
-    fn default() -> Self {
-        Self::new()
-    }
-}