pvm optimization (#160)

* parse code in init

* fix gas

* update and fix

PolkaVM/Sources/PolkaVM/Engine.swift

@@ -59,15 +59,17 @@ public class Engine {
     func step(program: ProgramCode, context: ExecutionContext) -> ExecOutcome {
         let pc = context.state.pc
         let skip = program.skip(pc)
-        let startIndex = program.code.startIndex + Int(pc)
-        let endIndex = startIndex + 1 + Int(skip)
-        let data = if endIndex <= program.code.endIndex {
-            program.code[startIndex ..< endIndex]
-        } else {
-            program.code[startIndex ..< min(program.code.endIndex, endIndex)] + Data(repeating: 0, count: endIndex - program.code.endIndex)
+        let inst = program.getInstructionAt(pc: pc)
+
+        guard let inst else {
+            return .exit(.panic(.invalidInstructionIndex))
         }
-        guard let inst = InstructionTable.parse(data) else {
-            return .exit(.panic(.invalidInstruction))
+
+        // TODO: check after GP specified the behavior
+        if context.state.program.basicBlockIndices.contains(pc) {
+            let blockGas = context.state.program.getBlockGasCosts(pc: pc)
+            context.state.consumeGas(blockGas)
+            logger.debug("consumed \(blockGas) gas for block at pc: \(pc)")
         }
 
         logger.debug("executing \(inst)", metadata: ["skip": "\(skip)", "pc": "\(context.state.pc)"])

PolkaVM/Sources/PolkaVM/ExecOutcome.swift

@@ -1,7 +1,7 @@
 public enum ExitReason {
     public enum PanicReason {
         case trap
-        case invalidInstruction
+        case invalidInstructionIndex
         case invalidDynamicJump
         case invalidBranch
     }

PolkaVM/Sources/PolkaVM/Instruction.swift

@@ -28,8 +28,6 @@ public class ExecutionContext {
 
 extension Instruction {
     public func execute(context: ExecutionContext, skip: UInt32) -> ExecOutcome {
-        context.state.consumeGas(gasCost())
-        logger.debug("consumed \(gasCost()) gas")
         do {
             let execRes = try _executeImpl(context: context)
             if case .exit = execRes {
@@ -38,8 +36,6 @@ extension Instruction {
             logger.debug("execution success! updating pc...")
             return updatePC(context: context, skip: skip)
         } catch let e as Memory.Error {
-            // this passes test vector
-            context.state.consumeGas(gasCost())
             logger.debug("memory error: \(e)")
             return .exit(.pageFault(e.address))
         } catch let e {

PolkaVM/Sources/PolkaVM/Instructions/Instructions.swift

@@ -32,7 +32,7 @@ public enum Instructions {
     public struct Trap: Instruction {
         public static var opcode: UInt8 { 0 }
 
-        public init(data _: Data) {}
+        public init(data _: Data = .init()) {}
 
         public func _executeImpl(context _: ExecutionContext) -> ExecOutcome {
             .exit(.panic(.trap))
@@ -1178,7 +1178,7 @@ public enum Instructions {
     // MARK: Instruction with Arguments of Two Registers and Two Immediates (5.11)
 
     public struct LoadImmJumpInd: Instruction {
-        public static var opcode: UInt8 { 10 }
+        public static var opcode: UInt8 { 42 }
 
         public let ra: Registers.Index
         public let rb: Registers.Index

PolkaVM/Sources/PolkaVM/ProgramCode.swift

@@ -7,6 +7,7 @@ public class ProgramCode {
         case invalidJumpTableEncodeSize
         case invalidCodeLength
         case invalidDataLength
+        case invalidInstruction
     }
 
     public enum Constants {
@@ -22,7 +23,11 @@ public class ProgramCode {
     public let code: Data
     private let bitmask: Data
 
-    public let basicBlockIndices: Set<UInt32>
+    // parsed stuff
+    public private(set) var basicBlockIndices: Set<UInt32> = []
+    private var skipCache: [UInt32: UInt32] = [:]
+    private var instCache: [UInt32: Instruction] = [:]
+    private var blockGasCosts: [UInt32: Gas] = [:]
 
     public init(_ blob: Data) throws(Error) {
         self.blob = blob
@@ -62,9 +67,74 @@ public class ProgramCode {
             throw Error.invalidDataLength
         }
 
-        bitmask = blob[codeEndIndex ..< slice.endIndex]
+        // mark bitmask bits longer than codeLength as 1
+        var bitmaskData = blob[codeEndIndex ..< slice.endIndex]
+        let fullBytes = Int(codeLength) / 8
+        let remainingBits = Int(codeLength) % 8
+        if remainingBits > 0 {
+            let mask: UInt8 = ~0 << remainingBits
+            bitmaskData[codeEndIndex + fullBytes] |= mask
+        }
+        bitmask = bitmaskData
+
+        try parseCode(code: code, bitmask: bitmask)
+    }
+
+    /// traverse the program code, collect basic block indices, cache skips and gas costs
+    private func parseCode(code: Data, bitmask: Data) throws(Error) {
+        var i = UInt32(0)
+        basicBlockIndices.insert(0)
+        var currentBlockStart = i
+        var currentBlockGasCost = Gas(0)
+        while i < code.count {
+            let skip = ProgramCode.skip(start: i, bitmask: bitmask)
+            skipCache[i] = skip
+
+            let inst = try parseInstruction(startIndex: code.startIndex + Int(i), skip: skip)
+            instCache[i] = inst
+            currentBlockGasCost += inst.gasCost()
+
+            let opcode = code[relative: Int(i)]
+            if BASIC_BLOCK_INSTRUCTIONS.contains(opcode) {
+                // block end
+                blockGasCosts[currentBlockStart] = currentBlockGasCost
+                // next block
+                basicBlockIndices.insert(i + skip + 1)
+                currentBlockStart = i + skip + 1
+                currentBlockGasCost = Gas(0)
+            }
+            i += skip + 1
+        }
+        blockGasCosts[currentBlockStart] = currentBlockGasCost
+        // trap at the end
+        instCache[i] = Instructions.Trap()
+        basicBlockIndices.insert(i)
+        blockGasCosts[i] = Instructions.Trap().gasCost()
+    }
+
+    private func parseInstruction(startIndex: Int, skip: UInt32) throws(Error) -> Instruction {
+        let endIndex = startIndex + Int(skip) + 1
+        let data = if endIndex <= code.endIndex {
+            code[startIndex ..< endIndex]
+        } else {
+            code[startIndex ..< min(code.endIndex, endIndex)] + Data(repeating: 0, count: endIndex - code.endIndex)
+        }
+        guard let inst = InstructionTable.parse(data) else {
+            throw Error.invalidInstruction
+        }
+        return inst
+    }
 
-        basicBlockIndices = ProgramCode.getBasicBlockIndices(code: code, bitmask: bitmask)
+    public func getInstructionAt(pc: UInt32) -> Instruction? {
+        instCache[pc]
+    }
+
+    public func getBlockGasCosts(pc: UInt32) -> Gas {
+        blockGasCosts[pc] ?? Gas(0)
+    }
+
+    public func skip(_ pc: UInt32) -> UInt32 {
+        skipCache[pc] ?? 0
     }
 
     public static func skip(start: UInt32, bitmask: Data) -> UInt32 {
@@ -91,26 +161,6 @@ public class ProgramCode {
 
         return idx
     }
-
-    public func skip(_ start: UInt32) -> UInt32 {
-        ProgramCode.skip(start: start, bitmask: bitmask)
-    }
-
-    /// traverse the program code and collect basic block indices
-    private static func getBasicBlockIndices(code: Data, bitmask: Data) -> Set<UInt32> {
-        // TODO: parse the instructions here and so we don't need to do skip calculation multiple times
-        var res: Set<UInt32> = [0]
-        var i = UInt32(0)
-        while i < code.count {
-            let opcode = code[relative: Int(i)]
-            let skip = ProgramCode.skip(start: i, bitmask: bitmask)
-            if BASIC_BLOCK_INSTRUCTIONS.contains(opcode) {
-                res.insert(i + skip + 1)
-            }
-            i += skip + 1
-        }
-        return res
-    }
 }
 
 extension ProgramCode: Equatable {