1 files changed, 72 insertions, 5 deletions

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))
+}