get better instruction stream parsing workingHEADmaster

3 files changed, 109 insertions, 26 deletions

diff --git a/src/AddressingModes.scala b/src/AddressingModes.scala
index f511768..30ad696 100644
--- a/src/AddressingModes.scala
+++ b/src/AddressingModes.scala
@@ -6,8 +6,9 @@ import ixee.cryptopals.utils.FunctionUtils._
 import scala.collection.mutable.Map
 
 sealed trait AddressingMode {
-  def modeName: String
-  override def toString: String = modeName
+  def name: String
+  def wide: Boolean
+  override def toString: String = name
   def nameifyRegister(reg: Int): String = reg match {
     case 0 => "pc"
     case 1 => "sp"
@@ -15,9 +16,18 @@ sealed trait AddressingMode {
     case 3 => "cg"
     case num @ _ => s"r$num"
   }
+}
+
+sealed trait NormMode extends AddressingMode {
+  def wide = false
   def stringifyWith(reg: Int): String
 }
 
+sealed trait WideMode extends AddressingMode {
+  def wide = true
+  def stringifyWith(reg: Int, word: Int): String
+}
+
 sealed trait OneOpAddressing extends AddressingMode
 sealed trait TwoOpAddressing extends AddressingMode
 sealed trait TwoOpSourceAddressing extends TwoOpAddressing
@@ -30,10 +40,7 @@ object TwoOpSourceAddressingModes {
 }
 
 object OneOpAddressingModes {
-  case class AddressingMode(as: Int, name: String, stringifier: String => String) extends OneOpAddressing with TwoOpSourceAddressing {
-    def modeName = name
-    def stringifyWith(reg: Int): String = stringifier(nameifyRegister(reg))
-  }
+  trait AddressingMode extends OneOpAddressing with TwoOpSourceAddressing
 
   def apply(register: Int, mode: Int): Option[AddressingMode] = register match {
     case 0 => PCModes(mode) orElse General(mode)
@@ -44,15 +51,21 @@ object OneOpAddressingModes {
 
   trait Modes {
     private val modeMap: Map[Int, AddressingMode] = Map()
-    case class Mode(as: Int, name: String, stringifier: String => String) extends AddressingMode(as, name, stringifier) {
-      def Mode(as: Int, name: String, stringifier: String => String) = this(as, name, stringifier).tap(modeMap += as -> _)
+
+    case class Mode(as: Int, name: String, stringifier: String => String) extends AddressingMode with NormMode {
+      modeMap += as -> this
+      def stringifyWith(reg: Int): String = stringifier(nameifyRegister(reg))
+    }
+    case class ExtraByteMode(as: Int, name: String, stringifier: (String, Int) => String) extends AddressingMode with WideMode {
+      modeMap += as -> this
+      def stringifyWith(reg: Int, const: Int): String = stringifier(nameifyRegister(reg), const)
     }
 
     def apply(x: Int): Option[AddressingMode] = modeMap.get(x)
   }
 
   object PCModes extends Modes {
-    Mode(3, "NextWord", reg => s"@PC+ (next word)")
+    ExtraByteMode(3, "NextWord", (reg, byte) => s"${if (byte < 0) "-" else ""}0x${"%x".format(Math.abs(byte))}")
   }
   object SRModes extends Modes {
     Mode(2, "Const4", reg => "#4")
@@ -66,7 +79,7 @@ object OneOpAddressingModes {
   }
   object General extends Modes {
     Mode(0, "RegisterDirect", ident)
-    Mode(1, "Indexed", reg => s"@PC+($reg)")
+    ExtraByteMode(1, "Indexed", (reg, byte) => s"${if (byte < 0) "-" else ""}0x${"%x".format(Math.abs(byte))}($reg)")
     Mode(2, "RegisterIndirect", reg => s"@$reg")
     Mode(3, "IndirectAutoInc", reg => s"@$reg+")
   }
@@ -79,10 +92,7 @@ object OneOpAddressingModes {
 
 
 object TwoOpDestAddressingModes {
-  case class AddressingMode(as: Int, name: String, stringifier: String => String) extends TwoOpDestAddressing {
-    def modeName = name
-    def stringifyWith(reg: Int): String = stringifier(nameifyRegister(reg))
-  }
+  trait AddressingMode extends TwoOpDestAddressing
 
   def apply(register: Int, mode: Int): Option[AddressingMode] = register match {
     case 0 => PCModes(mode) orElse General(mode)
@@ -92,7 +102,15 @@ object TwoOpDestAddressingModes {
 
   trait Modes {
     private val modeMap: Map[Int, AddressingMode] = Map()
-    def Mode(as: Int, name: String, stringifier: String => String) = AddressingMode(as, name, stringifier).tap(modeMap += as -> _)
+
+    case class Mode(as: Int, name: String, stringifier: String => String) extends AddressingMode with NormMode {
+      modeMap += as -> this
+      def stringifyWith(reg: Int): String = stringifier(nameifyRegister(reg))
+    }
+    case class ExtraByteMode(as: Int, name: String, stringifier: (String, Int) => String) extends AddressingMode with WideMode {
+      modeMap += as -> this
+      def stringifyWith(reg: Int, const: Int): String = stringifier(nameifyRegister(reg), const)
+    }
 
     def apply(x: Int): Option[AddressingMode] = modeMap.get(x)
   }
@@ -101,11 +119,11 @@ object TwoOpDestAddressingModes {
     Mode(1, "Symbolic", reg => s"<??> (symbolic? x(PC)?")
   }
   object SRModes extends Modes {
-    Mode(1, "Absolute", reg => s"&@PC+ (value at nextWord)")
+    ExtraByteMode(1, "Absolute", (reg, byte) => s"&${"%4x".format(byte)}")
   }
   object General extends Modes {
     Mode(0, "RegisterDirect", reg => s"$reg")
-    Mode(1, "Indexed", reg => s"@PC+($reg)")
+    ExtraByteMode(1, "Indexed", (reg, byte) => s"${if (byte < 0) "-" else ""}0x${"%x".format(Math.abs(byte))}($reg)")
   }
 
   PCModes
diff --git a/src/Instruction.scala b/src/Instruction.scala
index 8a4e018..7723ef7 100644
--- a/src/Instruction.scala
+++ b/src/Instruction.scala
@@ -3,25 +3,54 @@ package ixee.re.disasm
 import ixee.cryptopals.utils.ByteUtils._
 import ixee.cryptopals.utils.FunctionUtils._
 
-import scala.collection.mutable.Map
+import scala.collection.mutable.{ArrayBuffer, Map}
 
-sealed trait Instruction
+sealed trait Instruction {
+  def isDone: Boolean
+  def giveWord(word: Int): Unit
+}
 
 case class NoOpInst(code: NoOpCode, offset: Int) extends Instruction {
   override def toString: String = {
     s"$code -> $offset"
   }
+  def giveWord(word: Int) = throw new RuntimeException("Too many bytes!")
+  def isDone = true
 }
 case class OneOpInst(code: OneOpCode, byte: Boolean, dest: Int, addressMode: OneOpAddressing) extends Instruction {
+  private var extraWord: Option[Int] = None
+  def neededWords = Seq(addressMode).count(_.wide)
+  def giveWord(word: Int) =
+    extraWord.fold(extraWord = Some(word))(_ => throw new RuntimeException("Too many bytes!"))
+  def isDone = neededWords == Seq(extraWord).flatten.length
+
   override def toString: String = {
     val byteStr = if (byte) ".B" else "  "
-    s"$code$byteStr ${addressMode.stringifyWith(dest)}"
+    s"$code$byteStr ${addressMode match {
+      case mode: NormMode => mode.stringifyWith(dest)
+      case mode: WideMode => mode.stringifyWith(dest, extraWord.get)
+    }}"
   }
 }
 case class TwoOpInst(code: TwoOpCode, byte: Boolean, src: Int, srcMode: TwoOpSourceAddressing, dest: Int, dstMode: TwoOpDestAddressing) extends Instruction {
+  private var extraWord1: Option[Int] = None
+  private var extraWord2: Option[Int] = None
+  def neededWords = Seq(srcMode, dstMode).count(_.wide)
+  def giveWord(word: Int) =
+    extraWord1.fold(extraWord1 = Some(word))(_ =>
+      extraWord2.fold(extraWord2 = Some(word))(_ => throw new RuntimeException("Too many bytes!"))
+    )
+  def isDone = neededWords == Seq(extraWord1, extraWord2).flatten.length
+
   override def toString: String = {
     val byteStr = if (byte) ".B" else "  "
-    s"$code$byteStr ${srcMode.stringifyWith(src)}, ${dstMode.stringifyWith(dest)}"
+    s"$code$byteStr ${srcMode match {
+      case mode: NormMode => mode.stringifyWith(src)
+      case mode: WideMode => mode.stringifyWith(src, extraWord1.get)
+    }}, ${dstMode match {
+      case mode: NormMode => mode.stringifyWith(dest)
+      case mode: WideMode => mode.stringifyWith(dest, extraWord2.orElse(if (srcMode.wide) None else extraWord1).get)
+    }}"
   }
 }
 
@@ -103,10 +132,48 @@ object Instruction {
 
   def twoOp(instruction: (Byte, Byte)): Option[TwoOpInst] = TwoOp.apply(instruction)
 
-  def isNoOp(instruction: (Byte ,Byte)) = noOp(instruction).isDefined
+  def isNoOp(instruction: (Byte, Byte)) = noOp(instruction).isDefined
 
   def isOneOp(instruction: (Byte, Byte)) = oneOp(instruction).isDefined
 
   def isTwoOp(instruction: (Byte, Byte)) = twoOp(instruction).isDefined
 }
 
+class InstructionParser(bytes: Seq[Byte]) {
+  val instructions = scala.collection.mutable.ArrayBuffer[Instruction]()
+  val bytePair = scala.collection.mutable.ArrayBuffer[Byte]()
+
+  for (byte <- bytes) {
+    bytePair match {
+      case ArrayBuffer() => bytePair += byte
+      case ArrayBuffer(first) => {
+        val second = byte
+        def parseInstr = {
+          instructions += Instruction((second, first)).getOrElse { throw new RuntimeException(s"Parse failure! Bytes = $first, $second") }
+          bytePair.clear
+        }
+
+        val lastIsDone = instructions.lastOption.map(_.isDone).getOrElse(true)
+
+        instructions.lastOption.map { instr =>
+          if (instr.isDone) {
+            parseInstr
+          } else {
+            instr.giveWord(second * 256 + first)
+            bytePair.clear
+          }
+        } getOrElse { parseInstr }
+      }
+      case buffer @ _ => throw new RuntimeException("Byte array overflow: " + buffer)
+    }
+  }
+
+  instructions.lastOption match {
+    case Some(instr) if !instr.isDone => throw new RuntimeException("Unable to parse instruction stream, last instruction is missing data.")
+    case _ => Unit
+  }
+}
+
+object InstructionParser {
+  def apply(bytes: Seq[String]) = new InstructionParser(bytes.map(java.lang.Integer.parseInt(_, 16).toByte))
+}
diff --git a/src/main.scala b/src/main.scala
index f186597..4d90601 100644
--- a/src/main.scala
+++ b/src/main.scala
@@ -1,8 +1,6 @@
 package ixee.re.disasm
 
 object main extends App {
-  def disassemble(bytes: Seq[String]): Seq[Option[Instruction]] = bytes.map(Instruction(_))
-
   val bytes = """
 3140 0044 1542 5c01
 75f3 35d0 085a 3f40
@@ -19,7 +17,7 @@ f000 fd3f 3040 0a46
 3e40 fcff 0e54 0e12
 0f12 3012 7d00 b012
 7a45 5f44 fcff 8f11
-""".replace("\n", " ").split(" ").filter(_ != "")
+""".replace("\n", " ").replace(" ", "").trim.grouped(2).toSeq
 
-  println(disassemble(bytes).flatten.mkString("\n"))
+  println(InstructionParser(bytes).instructions.mkString("\n"))
 }