Merge pull request #806 from alephium/update-contract-code-storage

Use key-value storage for contract code

app/src/main/scala/org/alephium/app/ServerUtils.scala

@@ -1081,8 +1081,8 @@ class ServerUtils(implicit
   ): Try[ContractState] = {
     val result = for {
       state          <- worldState.getContractState(contractId)
-      codeRecord     <- worldState.getContractCode(state.codeHash)
-      contract       <- codeRecord.code.toContract().left.map(IOError.Serde)
+      code           <- worldState.getContractCode(state)
+      contract       <- code.toContract().left.map(IOError.Serde)
       contractOutput <- worldState.getContractAsset(state.contractOutputRef)
     } yield ContractState(
       Address.contract(contractId),

flow/src/main/scala/org/alephium/flow/io/Storages.scala

@@ -55,7 +55,7 @@ object Storages {
     val logCounterStorage = RocksDBKeyValueStorage[ContractId, Int](db, LogCounter, writeOptions)
     val logStorage        = LogStorage(logStateStorage, logRefStorage, logCounterStorage)
     val trieImmutableStateStorage =
-      RocksDBKeyValueStorage[Hash, ContractImmutableState](db, Trie, writeOptions)
+      RocksDBKeyValueStorage[Hash, ContractStorageImmutableState](db, Trie, writeOptions)
     val worldStateStorage =
       WorldStateRockDBStorage(
         trieStorage,

flow/src/main/scala/org/alephium/flow/io/WorldStateStorage.scala

@@ -23,12 +23,12 @@ import org.alephium.io._
 import org.alephium.io.RocksDBSource.{ColumnFamily, Settings}
 import org.alephium.protocol.Hash
 import org.alephium.protocol.model.BlockHash
-import org.alephium.protocol.vm.{ContractImmutableState, WorldState}
+import org.alephium.protocol.vm.{ContractStorageImmutableState, WorldState}
 import org.alephium.protocol.vm.event.LogStorage
 
 trait WorldStateStorage extends KeyValueStorage[BlockHash, WorldState.Hashes] {
   val trieStorage: KeyValueStorage[Hash, SparseMerkleTrie.Node]
-  val trieImmutableStateStorage: KeyValueStorage[Hash, ContractImmutableState]
+  val trieImmutableStateStorage: KeyValueStorage[Hash, ContractStorageImmutableState]
   val logStorage: LogStorage
 
   override def storageKey(key: BlockHash): ByteString =
@@ -54,7 +54,7 @@ trait WorldStateStorage extends KeyValueStorage[BlockHash, WorldState.Hashes] {
 object WorldStateRockDBStorage {
   def apply(
       trieStorage: KeyValueStorage[Hash, SparseMerkleTrie.Node],
-      trieImmutableStateStorage: KeyValueStorage[Hash, ContractImmutableState],
+      trieImmutableStateStorage: KeyValueStorage[Hash, ContractStorageImmutableState],
       logStorage: LogStorage,
       storage: RocksDBSource,
       cf: ColumnFamily,
@@ -74,7 +74,7 @@ object WorldStateRockDBStorage {
 
 class WorldStateRockDBStorage(
     val trieStorage: KeyValueStorage[Hash, SparseMerkleTrie.Node],
-    val trieImmutableStateStorage: KeyValueStorage[Hash, ContractImmutableState],
+    val trieImmutableStateStorage: KeyValueStorage[Hash, ContractStorageImmutableState],
     val logStorage: LogStorage,
     storage: RocksDBSource,
     cf: ColumnFamily,

flow/src/test/scala/org/alephium/flow/core/VMSpec.scala

@@ -297,7 +297,8 @@ class VMSpec extends AlephiumSpec {
       val contractKey = ContractId.from(Hex.from(contractId).get).get
       worldState.contractState.exists(contractKey) isE existed
       worldState.outputState.exists(contractAssetRef) isE existed
-      worldState.codeState.exists(contract.hash) isE existed
+      worldState.codeState.exists(contract.hash) isE false
+      worldState.contractImmutableState.exists(contract.hash) isE true // keep history state always
     }
 
     def getContractAsset(contractId: ContractId, chainIndex: ChainIndex): ContractOutput = {

io/src/main/scala/org/alephium/io/CachedKV.scala

@@ -21,6 +21,8 @@ import scala.collection.mutable
 import org.alephium.util.discard
 
 abstract class CachedKV[K, V, C >: Modified[V] <: Cache[V]] extends MutableKV[K, V, Unit] {
+  def unit: Unit = ()
+
   def underlying: ReadableKV[K, V]
 
   def caches: mutable.Map[K, C]

io/src/main/scala/org/alephium/io/KeyValueStorage.scala

@@ -50,6 +50,7 @@ trait KeyValueStorage[K, V]
     extends AbstractKeyValueStorage[K, V]
     with RawKeyValueStorage
     with MutableKV[K, V, Unit] {
+  def unit: Unit = ()
 
   protected def storageKey(key: K): ByteString = serialize(key)
 

io/src/main/scala/org/alephium/io/MutableKV.scala

@@ -20,6 +20,8 @@ trait MutableKV[K, V, T] extends ReadableKV[K, V] {
   def remove(key: K): IOResult[T]
 
   def put(key: K, value: V): IOResult[T]
+
+  def unit: T
 }
 
 object MutableKV {

io/src/main/scala/org/alephium/io/SparseMerkleTrie.scala

@@ -554,6 +554,8 @@ final class SparseMerkleTrie[K: Serde, V: Serde](
 ) extends SparseMerkleTrieBase[K, V, SparseMerkleTrie[K, V]] {
   import SparseMerkleTrie._
 
+  def unit: SparseMerkleTrie[K, V] = this
+
   def getNode(hash: Hash): IOResult[Node] = storage.get(hash)
 
   def applyActions(result: TrieUpdateActions): IOResult[SparseMerkleTrie[K, V]] = {
@@ -601,6 +603,8 @@ final class InMemorySparseMerkleTrie[K: Serde, V: Serde](
 ) extends SparseMerkleTrieBase[K, V, Unit] {
   import SparseMerkleTrie._
 
+  def unit: Unit = ()
+
   def getNode(hash: Hash): IOResult[Node] = {
     cache.get(hash) match {
       case Some(node) => Right(node)

protocol/src/main/scala/org/alephium/protocol/vm/ContractState.scala

@@ -70,8 +70,6 @@ sealed trait ContractState {
   def updateOutputRef(ref: ContractOutputRef): ContractStorageState
 }
 
-object ContractState {}
-
 final case class ContractLegacyState private (
     codeHash: Hash,
     initialStateHash: Hash,

protocol/src/main/scala/org/alephium/protocol/vm/WorldState.scala

@@ -29,7 +29,7 @@ import org.alephium.util.AVector
 trait WorldState[T, R1, R2, R3] {
   def outputState: MutableKV[TxOutputRef, TxOutput, R1]
   def contractState: MutableKV[ContractId, ContractStorageState, R2]
-  def contractImmutableState: MutableKV[Hash, ContractImmutableState, Unit]
+  def contractImmutableState: MutableKV[Hash, ContractStorageImmutableState, Unit]
   def codeState: MutableKV[Hash, WorldState.CodeRecord, R3]
 
   @SuppressWarnings(Array("org.wartremover.warts.AsInstanceOf"))
@@ -74,7 +74,7 @@ trait WorldState[T, R1, R2, R3] {
     contractState.get(id).flatMap {
       case mutable: ContractMutableState =>
         contractImmutableState.get(mutable.immutableStateHash) map {
-          case immutable: ContractImmutableState =>
+          case Left(immutable: ContractImmutableState) =>
             ContractNewState(immutable, mutable)
           case _ => throw new RuntimeException("Invalid contract state")
         }
@@ -86,8 +86,15 @@ trait WorldState[T, R1, R2, R3] {
     contractState.exists(id)
   }
 
-  def getContractCode(id: Hash): IOResult[WorldState.CodeRecord] = {
-    codeState.get(id)
+  def getContractCode(state: ContractState): IOResult[StatefulContract.HalfDecoded] = {
+    state match {
+      case _: ContractLegacyState => codeState.get(state.codeHash).map(_.code)
+      case _: ContractNewState =>
+        contractImmutableState.get(state.codeHash).map {
+          case Right(code) => code
+          case _           => throw new RuntimeException("Invalid contract state")
+        }
+    }
   }
 
   def getContractAsset(id: ContractId): IOResult[ContractOutput] = {
@@ -120,8 +127,8 @@ trait WorldState[T, R1, R2, R3] {
   def getContractObj(key: ContractId): IOResult[StatefulContractObject] = {
     for {
       state <- getContractState(key)
-      code  <- getContractCode(state.codeHash)
-    } yield state.toObject(key, code.code)
+      code  <- getContractCode(state)
+    } yield state.toObject(key, code)
   }
 
   def addAsset(outputRef: TxOutputRef, output: TxOutput): IOResult[T]
@@ -175,16 +182,26 @@ trait WorldState[T, R1, R2, R3] {
   def removeContractFromVM(contractKey: ContractId): IOResult[T]
 
   protected def removeContractCode(
-      currentState: ContractState,
-      currentRecord: WorldState.CodeRecord
+      currentState: ContractState
   ): IOResult[R3] = {
-    if (currentRecord.refCount > 1) {
-      codeState.put(
-        currentState.codeHash,
-        currentRecord.copy(refCount = currentRecord.refCount - 1)
-      )
-    } else {
-      codeState.remove(currentState.codeHash)
+    currentState match {
+      case _: ContractLegacyState => removeContractCodeDeprecated(currentState)
+      case _: ContractNewState    => Right(codeState.unit) // We keep the code as history state
+    }
+  }
+
+  protected def removeContractCodeDeprecated(
+      currentState: ContractState
+  ): IOResult[R3] = {
+    codeState.get(currentState.codeHash).flatMap { currentRecord =>
+      if (currentRecord.refCount > 1) {
+        codeState.put(
+          currentState.codeHash,
+          currentRecord.copy(refCount = currentRecord.refCount - 1)
+        )
+      } else {
+        codeState.remove(currentState.codeHash)
+      }
     }
   }
 
@@ -280,7 +297,7 @@ object WorldState {
   final case class Persisted(
       outputState: SparseMerkleTrie[TxOutputRef, TxOutput],
       contractState: SparseMerkleTrie[ContractId, ContractStorageState],
-      contractImmutableState: KeyValueStorage[Hash, ContractImmutableState],
+      contractImmutableState: KeyValueStorage[Hash, ContractStorageImmutableState],
       codeState: SparseMerkleTrie[Hash, CodeRecord],
       logStorage: LogStorage
   ) extends ImmutableWorldState {
@@ -365,14 +382,13 @@ object WorldState {
       for {
         newOutputState   <- outputState.put(outputRef, output)
         newContractState <- contractState.put(contractId, state.mutable)
-        _         <- contractImmutableState.put(state.mutable.immutableStateHash, state.immutable)
-        recordOpt <- codeState.getOpt(code.hash)
-        newCodeState <- codeState.put(code.hash, CodeRecord.from(code, recordOpt))
+        _ <- contractImmutableState.put(state.mutable.immutableStateHash, Left(state.immutable))
+        _ <- contractImmutableState.put(state.codeHash, Right(code))
       } yield Persisted(
         newOutputState,
         newContractState,
         contractImmutableState,
-        newCodeState,
+        codeState,
         logStorage
       )
     }
@@ -415,8 +431,7 @@ object WorldState {
       for {
         state            <- getContractState(contractKey)
         newContractState <- contractState.remove(contractKey)
-        codeRecord       <- codeState.get(state.codeHash)
-        newCodeState     <- removeContractCode(state, codeRecord)
+        newCodeState     <- removeContractCode(state)
       } yield Persisted(
         outputState,
         newContractState,
@@ -450,7 +465,7 @@ object WorldState {
   sealed abstract class AbstractCached extends MutableWorldState {
     def outputState: MutableKV[TxOutputRef, TxOutput, Unit]
     def contractState: MutableKV[ContractId, ContractStorageState, Unit]
-    def contractImmutableState: MutableKV[Hash, ContractImmutableState, Unit]
+    def contractImmutableState: MutableKV[Hash, ContractStorageImmutableState, Unit]
     def codeState: MutableKV[Hash, CodeRecord, Unit]
     def logState: MutableLog
 
@@ -484,11 +499,10 @@ object WorldState {
     ): IOResult[Unit] = {
       val state = ContractNewState.unsafe(code, immFields, mutFields, outputRef)
       for {
-        _         <- outputState.put(outputRef, output)
-        _         <- contractState.put(contractId, state.mutable)
-        _         <- contractImmutableState.put(state.mutable.immutableStateHash, state.immutable)
-        recordOpt <- codeState.getOpt(code.hash)
-        _         <- codeState.put(code.hash, CodeRecord.from(code, recordOpt))
+        _ <- outputState.put(outputRef, output)
+        _ <- contractState.put(contractId, state.mutable)
+        _ <- contractImmutableState.put(state.mutable.immutableStateHash, Left(state.immutable))
+        _ <- contractImmutableState.put(state.codeHash, Right(code))
       } yield ()
     }
 
@@ -533,11 +547,11 @@ object WorldState {
       val migratedState = state.migrate(newCode, newImmFields, newMutFields)
       for {
         _ <- updateContract(contractId, migratedState.mutable)
-        _ <- contractImmutableState.put(migratedState.immutableStateHash, migratedState.immutable)
-        codeRecord       <- codeState.get(state.codeHash)
-        _                <- removeContractCode(state, codeRecord)
-        newCodeRecordOpt <- codeState.getOpt(newCode.hash)
-        _ <- codeState.put(newCode.hash, CodeRecord.from(newCode.toHalfDecoded(), newCodeRecordOpt))
+        _ <- contractImmutableState.put(
+          migratedState.immutableStateHash,
+          Left(migratedState.immutable)
+        )
+        _ <- contractImmutableState.put(newCode.hash, Right(newCode.toHalfDecoded()))
       } yield true
     }
 
@@ -548,10 +562,9 @@ object WorldState {
     // Contract output is already removed by the VM
     def removeContractFromVM(contractId: ContractId): IOResult[Unit] = {
       for {
-        state      <- getContractState(contractId)
-        codeRecord <- codeState.get(state.codeHash)
-        _          <- removeContractCode(state, codeRecord)
-        _          <- contractState.remove(contractId)
+        state <- getContractState(contractId)
+        _     <- removeContractCode(state)
+        _     <- contractState.remove(contractId)
       } yield ()
     }
 
@@ -566,7 +579,7 @@ object WorldState {
   final case class Cached(
       outputState: CachedSMT[TxOutputRef, TxOutput],
       contractState: CachedSMT[ContractId, ContractStorageState],
-      contractImmutableState: CachedKVStorage[Hash, ContractImmutableState],
+      contractImmutableState: CachedKVStorage[Hash, ContractStorageImmutableState],
       codeState: CachedSMT[Hash, CodeRecord],
       logState: CachedLog
   ) extends AbstractCached {
@@ -599,7 +612,7 @@ object WorldState {
   final case class Staging(
       outputState: StagingSMT[TxOutputRef, TxOutput],
       contractState: StagingSMT[ContractId, ContractStorageState],
-      contractImmutableState: StagingKVStorage[Hash, ContractImmutableState],
+      contractImmutableState: StagingKVStorage[Hash, ContractStorageImmutableState],
       codeState: StagingSMT[Hash, CodeRecord],
       logState: StagingLog
   ) extends AbstractCached {
@@ -624,7 +637,7 @@ object WorldState {
 
   def emptyPersisted(
       trieStorage: KeyValueStorage[Hash, SparseMerkleTrie.Node],
-      trieImmutableStateStorage: KeyValueStorage[Hash, ContractImmutableState],
+      trieImmutableStateStorage: KeyValueStorage[Hash, ContractStorageImmutableState],
       logStorage: LogStorage
   ): Persisted = {
     val genesisRef  = ContractOutputRef.forSMT
@@ -647,7 +660,7 @@ object WorldState {
 
   def emptyCached(
       trieStorage: KeyValueStorage[Hash, SparseMerkleTrie.Node],
-      trieImmutableStateStorage: KeyValueStorage[Hash, ContractImmutableState],
+      trieImmutableStateStorage: KeyValueStorage[Hash, ContractStorageImmutableState],
       logStorage: LogStorage
   ): Cached = {
     emptyPersisted(trieStorage, trieImmutableStateStorage, logStorage).cached()
@@ -656,7 +669,7 @@ object WorldState {
   final case class Hashes(outputStateHash: Hash, contractStateHash: Hash, codeStateHash: Hash) {
     def toPersistedWorldState(
         trieStorage: KeyValueStorage[Hash, SparseMerkleTrie.Node],
-        trieImmutableStateStorage: KeyValueStorage[Hash, ContractImmutableState],
+        trieImmutableStateStorage: KeyValueStorage[Hash, ContractStorageImmutableState],
         logStorage: LogStorage
     ): Persisted = {
       val outputState = SparseMerkleTrie[TxOutputRef, TxOutput](outputStateHash, trieStorage)
@@ -668,7 +681,7 @@ object WorldState {
 
     def toCachedWorldState(
         trieStorage: KeyValueStorage[Hash, SparseMerkleTrie.Node],
-        trieImmutableStateStorage: KeyValueStorage[Hash, ContractImmutableState],
+        trieImmutableStateStorage: KeyValueStorage[Hash, ContractStorageImmutableState],
         logStorage: LogStorage
     ): Cached = {
       toPersistedWorldState(trieStorage, trieImmutableStateStorage, logStorage).cached()

protocol/src/main/scala/org/alephium/protocol/vm/package.scala

@@ -50,4 +50,6 @@ package object vm {
   val destroyContractEventIndex: Val.I256 = Val.I256(I256.from(-2))
   val debugEventIndex: Val.I256           = Val.I256(I256.from(-3))
   // scalastyle:on magic.number
+
+  type ContractStorageImmutableState = Either[ContractImmutableState, StatefulContract.HalfDecoded]
 }

protocol/src/test/scala/org/alephium/protocol/vm/InstrSpec.scala

@@ -2629,8 +2629,8 @@ class InstrSpec extends AlephiumSpec with NumericHelpers {
       }
       contractOutput.tokens.toSet is allTokens.toSet
       contractOutput.amount is attoAlphAmount
-      val contractRecord = frame.ctx.worldState.getContractCode(contractState.codeHash).rightValue
-      contractRecord.code.toContract() isE contract
+      val code = frame.ctx.worldState.getContractCode(contractState).rightValue
+      code.toContract() isE contract
     }
   }
 

protocol/src/test/scala/org/alephium/protocol/vm/VMFactory.scala

@@ -28,7 +28,7 @@ trait VMFactory extends StorageFixture {
     val storage = newDBStorage()
     val trieDb  = newDB[Hash, SparseMerkleTrie.Node](storage, RocksDBSource.ColumnFamily.Trie)
     val trieImmutableStateStorage =
-      newDB[Hash, ContractImmutableState](storage, RocksDBSource.ColumnFamily.Trie)
+      newDB[Hash, ContractStorageImmutableState](storage, RocksDBSource.ColumnFamily.Trie)
     val logDb        = newDB[LogStatesId, LogStates](storage, RocksDBSource.ColumnFamily.Log)
     val logRefDb     = newDB[Byte32, AVector[LogStateRef]](storage, RocksDBSource.ColumnFamily.Log)
     val logCounterDb = newDB[ContractId, Int](storage, RocksDBSource.ColumnFamily.LogCounter)