cleanup

1 files changed, 87 insertions, 81 deletions

diff --git a/elf.py b/elf.py
index 919b690..ae81253 100755
--- a/elf.py
+++ b/elf.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 import sys
-from subx import clean, classify, parse_segment, parse_instr
+import subv
 
 """
 Takes a packed & surveyed hex stream (with section headers) and packs it into a
@@ -18,86 +18,92 @@ See full documentation here:
 https://sifive.cdn.prismic.io/sifive/4d063bf8-3ae6-4db6-9843-ee9076ebadf7_fe310-g000.pdf
 """
 
-segments = []
-segment = None
-for line in sys.stdin:
-  line = clean(line)
-  type = classify(line)
-
-  if type == 'segment':
-    (name, addr) = parse_segment(line)
-    segment = { 'name': name, 'addr': addr, 'content': [] }
-    segments.append(segment)
-  elif type == 'instr':
-    if segment == None:
-      raise ValueError("label or code outside of segment!")
-
-    instr = parse_instr(line)
-    segment['content'] += instr
-  else:
-    raise ValueError("elf input should contain only segments and instructions!")
-entrypoint = 0x80000000
-cursor = 0
-
 def w(b):
-  global cursor
-  cursor += len(b)
-  sys.stdout.buffer.write(b)
+    sys.stdout.buffer.write(b)
+    return len(b)
 
 def wi(num, size=2):
-  global cursor
-  cursor += size
-  sys.stdout.buffer.write(num.to_bytes(size, byteorder='little'))
-
-def padto(offset):
-  missing = offset - cursor
-  if missing < 0:
-    raise ValueError("already past offset!")
-  w(b'\x00' * missing) 
-
-def write_elf_header():
-  w(b"\x7fELF")     # ELF magic
-  wi(0x010101, 4)   # 32bit, little endian
-  wi(0, 8)          # reserved
-  wi(0x02)          # e_type
-  wi(0xf3)          # e_machine
-  wi(0x01, 4)       # e_version
-  wi(entrypoint, 4) # e_entry
-  wi(0x34, 4)       # e_phoff (Program Header offset)
-  wi(0x00, 4)       # e_shoff (Section Header offset, unused)
-  wi(0x0004, 4)     # e_flags
-  wi(0x34)          # e_ehsize
-  wi(0x20)          # e_phentsize
-  wi(len(segments)) # e_phnum
-  wi(0x28)          # e_shentsize
-  wi(0x0)           # e_shnum
-  wi(0x0)           # e_shstrndx
-
-def write_program_header(segment):
-  align = 0x1000
-  offset = (1 + cursor // align) * align
-  start = segment['addr']
-  size = len(segment['content'])
-  segment['offset'] = offset
-
-  if offset % align != start % align:
-    print("{:02x} {:02x}".format(offset, offset % align))
-    print("{:02x} {:02x}".format(start, start % align))
-    raise ValueError("improper alignment")
-
-  wi(0x1, 4)    # p_type
-  wi(offset, 4) # p_offset
-  wi(start, 4)  # p_vaddr
-  wi(start, 4)  # p_paddr
-  wi(size, 4)   # p_filesz
-  wi(size, 4)   # p_memsz
-  wi(0x5, 4)    # p_flags, 0x5=rx, 0x6=rw
-  wi(align, 4)  # p_align
-
-write_elf_header()
-for seg in segments:
-  write_program_header(seg)
-for seg in segments:
-  padto(seg['offset'])
-  for part in seg['content']:
-    wi(part[0], 1)
+    sys.stdout.buffer.write(num.to_bytes(size, byteorder='little'))
+    return size
+
+def padto(offset, cursor):
+    missing = offset - cursor
+    if missing < 0:
+        raise ValueError("already past offset!")
+    w(b'\x00' * missing)
+    return offset
+
+def write_elf_header(segments):
+    c = 0
+
+    code = next(s for s in segments if s['name'] == 'code')
+    entrypoint = code['addr']
+
+    c += w(b"\x7fELF")     # ELF magic
+    c += wi(0x010101, 4)   # 32bit, little endian
+    c += wi(0, 8)          # reserved
+    c += wi(0x02)          # e_type
+    c += wi(0xf3)          # e_machine
+    c += wi(0x01, 4)       # e_version
+    c += wi(entrypoint, 4) # e_entry
+    c += wi(0x34, 4)       # e_phoff (Program Header offset)
+    c += wi(0x00, 4)       # e_shoff (Section Header offset, unused)
+    c += wi(0x0004, 4)     # e_flags
+    c += wi(0x34)          # e_ehsize
+    c += wi(0x20)          # e_phentsize
+    c += wi(len(segments)) # e_phnum
+    c += wi(0x28)          # e_shentsize
+    c += wi(0x0)           # e_shnum
+    c += wi(0x0)           # e_shstrndx
+    return c
+
+def write_program_header(segment, c):
+    align = 0x1000
+    offset = (1 + c // align) * align
+    start = segment['addr']
+    size = len(segment['content'])
+    segment['offset'] = offset
+
+    if offset % align != start % align:
+        print("{:02x} {:02x}".format(offset, offset % align))
+        print("{:02x} {:02x}".format(start, start % align))
+        raise ValueError("improper alignment")
+
+    c += wi(0x1, 4)    # p_type
+    c += wi(offset, 4) # p_offset
+    c += wi(start, 4)  # p_vaddr
+    c += wi(start, 4)  # p_paddr
+    c += wi(size, 4)   # p_filesz
+    c += wi(size, 4)   # p_memsz
+    c += wi(0x5, 4)    # p_flags, 0x5=rx, 0x6=rw
+    c += wi(align, 4)  # p_align
+    return c
+
+if __name__ == "__main__":
+    segments = []
+    segment = None
+    for line in sys.stdin:
+        line = subv.parse(line)
+
+        if line['type'] == 'segment':
+            (name, addr) = line['segment']
+            segment = { 'name': name, 'addr': addr, 'content': [] }
+            segments.append(segment)
+        elif line['type'] == 'instr':
+            if segment == None:
+              raise ValueError("label or code outside of segment!")
+
+            segment['content'] += line['instr']
+        else:
+            raise ValueError("elf input should contain only segments and instructions!")
+
+    cursor = 0
+
+    c = write_elf_header(segments)
+    for seg in segments:
+        c = write_program_header(seg, c)
+
+    for seg in segments:
+        c = padto(seg['offset'], c)
+        for part in seg['content']:
+            c += wi(part[0], 1)