Fix alignment calculation and padding

Aligned values are a little funky, because the size of the line can't be
determined without knowing what address we're at. So the
DataLines iterator now takes a start address, so it is able to correctly
calculate alignment padding.

Also, the length of the data section uses the DataLines iterator now so
there isn't a complex re-implementation of the same logic.

Signed-off-by: Alek Ratzloff <alekratz@gmail.com>

src/libvm/examples/deadbeef.asm

@@ -25,9 +25,7 @@
         halt
     .export main
 
-}
+    .align u64
-
-.section iv $0x2800 {
     dz:
         mov %status, $255
         halt

src/libvm/src/obj/assemble.rs

@@ -59,6 +59,7 @@ impl Asm for DataSection {
     fn assemble(&self, session: &mut AsmSession) -> Result<Self::Out> {
         let names = names::get_section_names(self)?;
         session.name_stack.push(names);
+
         let content_len = self.len() as u64;
         let (start, end) = match self.org {
             SectionOrg::Start(start) => (start, start + content_len),
@@ -78,9 +79,13 @@ impl Asm for DataSection {
         }
 
         let mut contents = Vec::with_capacity(content_len as usize);
-        for (pos, line) in self.lines() {
+        for (pos, line) in self.lines(start) {
+            let expected_len = pos - (start as usize);
+            assert!(expected_len >= contents.len(), "next line size would cause section to shrink");
+            // resize contents if necessary, this pads out aligned values
+            contents.resize(expected_len, 0);
             contents.extend(line.assemble(session)?);
-            session.pos = start + (pos as u64);
+            session.pos = pos as u64;
         }
         assert_eq!(
             contents.len() as u64,
@@ -89,8 +94,11 @@ impl Asm for DataSection {
             self.name
         );
         assert_eq!(
-            session.pos - start, content_len
+            session.pos - start, content_len,
+            "in section {}",
+            self.name
         );
+
         session.name_stack.pop();
         Ok(obj::DataSection {
             name: self.name.clone(),

src/libvm/src/obj/assemble/names.rs

@@ -37,25 +37,26 @@ pub fn get_section_names(section: &DataSection) -> Result<Names> {
     // At the end of name gathering, if any exports have not been removed, then we know those
     // are exports whose names are undefined and we can return an UnknownExport error.
 
+    let start = section.org.start();
+
     // get exported names
-    for (_, line) in section.lines() {
+    for (_, line) in section.lines(start) {
         if let DataLine::Export(name) = line {
             exports.insert(name);
         }
     }
 
     // get names
-    for (pos, line) in section.lines() {
+    for (pos, line) in section.lines(start) {
         if let DataLine::Label(name) = line {
             if names.contains_key(name) {
                 return Err(AsmError::DuplicateLabel { name: name.clone() });
             }
             let export = exports.remove(name);
-            let start = section.org.start();
             
             names.insert(name.clone(), Name {
                 name: name.clone(),
-                addr: Addr(start + (pos as u64)),
+                addr: Addr(pos as u64),
                 export,
             });
         }

src/libvm/src/obj/syn/ast.rs

@@ -33,13 +33,16 @@ pub struct DataSection {
 
 impl DataSection {
     pub fn len(&self) -> usize {
-        self.blocks.iter()
+        let start = self.org.start();
-            .map(AlignedBlock::len)
+        let mut end = start;
-            .sum()
+        for (pos, _) in self.lines(start) {
+            end = pos as u64;
+        }
+        (end - start) as usize
     }
 
-    pub fn lines<'a>(&'a self) -> DataLines<'a> {
+    pub fn lines<'a>(&'a self, start: u64) -> DataLines<'a> {
-        DataLines::new(&self.blocks)
+        DataLines::new(&self.blocks, start)
     }
 }
 
@@ -51,12 +54,12 @@ pub struct DataLines<'a> {
 }
 
 impl<'a> DataLines<'a> {
-    fn new(blocks: &'a Vec<AlignedBlock>) -> Self {
+    fn new(blocks: &'a Vec<AlignedBlock>, start: u64) -> Self {
         DataLines {
             blocks,
             block_idx: 0,
             line_idx: 0,
-            pos: 0,
+            pos: start as usize,
         }
     }
 }
@@ -73,8 +76,11 @@ impl<'a> Iterator for DataLines<'a> {
             // next block - advance the position by the padding amount
             self.block_idx += 1;
             self.line_idx = 0;
-            let pos = self.pos;
+            // if there's a next block, update the padding
-            self.pos += block.padding_for(pos);
+            if self.block_idx < self.blocks.len() {
+                let block = &self.blocks[self.block_idx];
+                self.pos += block.padding_for(self.pos);
+            }
             self.next()
         } else {
             let pos = self.pos;
@@ -107,16 +113,14 @@ pub struct AlignedBlock {
 }
 
 impl AlignedBlock {
-    pub fn len(&self) -> usize {
-        let block_len = self.block.iter()
-            .map(DataLine::len)
-            .sum();
-        block_len + self.padding_for(block_len)
-    }
-
     pub fn padding_for(&self, len: usize) -> usize {
         let align = self.alignment.len();
-        len % align
+        let padding = len % align;
+        if padding == 0 {
+            padding
+        } else {
+            align - padding
+        }
     }
 }