client, node: Handle snapshot in explorer

client/src/main/kotlin/com/r3corda/client/model/ContractStateModel.kt

@@ -1,20 +1,23 @@
 package com.r3corda.client.model
 
+import com.r3corda.client.fxutils.foldToObservableList
 import com.r3corda.contracts.asset.Cash
 import com.r3corda.core.contracts.ContractState
 import com.r3corda.core.contracts.StateAndRef
 import com.r3corda.core.contracts.StateRef
-import com.r3corda.client.fxutils.foldToObservableList
 import com.r3corda.node.services.monitor.ServiceToClientEvent
 import com.r3corda.node.services.monitor.StateSnapshotMessage
 import javafx.collections.ObservableList
 import kotlinx.support.jdk8.collections.removeIf
 import rx.Observable
 
-class StatesDiff<out T : ContractState>(
+sealed class StatesModification<out T : ContractState>{
+    class Diff<out T : ContractState>(
             val added: Collection<StateAndRef<T>>,
             val removed: Collection<StateRef>
-)
+    ) : StatesModification<T>()
+    class Reset<out T : ContractState>(val states: Collection<StateAndRef<T>>) : StatesModification<T>()
+}
 
 /**
  * This model exposes the list of owned contract states.
@@ -24,16 +27,45 @@ class ContractStateModel {
     private val snapshot: Observable<StateSnapshotMessage> by observable(WalletMonitorModel::snapshot)
     private val outputStates = serviceToClient.ofType(ServiceToClientEvent.OutputState::class.java)
 
-    val contractStatesDiff = outputStates.map { StatesDiff(it.produced, it.consumed) }
+    val contractStatesDiff: Observable<StatesModification.Diff<ContractState>> =
+            outputStates.map { StatesModification.Diff(it.produced, it.consumed) }
     // We filter the diff first rather than the complete contract state list.
-    // TODO wire up snapshot once it holds StateAndRefs
+    val cashStatesModification: Observable<StatesModification<Cash.State>> = Observable.merge(
-    val cashStatesDiff = contractStatesDiff.map {
+            arrayOf(
-        StatesDiff(it.added.filterIsInstance<StateAndRef<Cash.State>>(), it.removed)
+                    contractStatesDiff.map {
+                        StatesModification.Diff(it.added.filterCashStateAndRefs(), it.removed)
+                    },
+                    snapshot.map {
+                        StatesModification.Reset(it.contractStates.filterCashStateAndRefs())
                     }
+            )
+    )
     val cashStates: ObservableList<StateAndRef<Cash.State>> =
-            cashStatesDiff.foldToObservableList(Unit) { statesDiff, _accumulator, observableList ->
+            cashStatesModification.foldToObservableList(Unit) { statesDiff, _accumulator, observableList ->
+                when (statesDiff) {
+                    is StatesModification.Diff -> {
                         observableList.removeIf { it.ref in statesDiff.removed }
                         observableList.addAll(statesDiff.added)
                     }
+                    is StatesModification.Reset -> {
+                        observableList.setAll(statesDiff.states)
+                    }
+                }
+            }
+
+
+    companion object {
+        private fun Collection<StateAndRef<ContractState>>.filterCashStateAndRefs(): List<StateAndRef<Cash.State>> {
+            return this.map { stateAndRef ->
+                @Suppress("UNCHECKED_CAST")
+                if (stateAndRef.state.data is Cash.State) {
+                    // Kotlin doesn't unify here for some reason
+                    stateAndRef as StateAndRef<Cash.State>
+                } else {
+                    null
+                }
+            }.filterNotNull()
+        }
+    }
 
 }

node/src/main/kotlin/com/r3corda/node/services/monitor/Messages.kt

@@ -2,6 +2,7 @@ package com.r3corda.node.services.monitor
 
 import com.r3corda.core.contracts.ClientToServiceCommand
 import com.r3corda.core.contracts.ContractState
+import com.r3corda.core.contracts.StateAndRef
 import com.r3corda.core.messaging.SingleMessageRecipient
 import com.r3corda.protocols.DirectRequestMessage
 
@@ -14,6 +15,6 @@ data class DeregisterRequest(override val replyToRecipient: SingleMessageRecipie
                              override val sessionID: Long) : DirectRequestMessage
 
 data class DeregisterResponse(val success: Boolean)
-data class StateSnapshotMessage(val contractStates: Collection<ContractState>, val protocolStates: Collection<String>)
+data class StateSnapshotMessage(val contractStates: Collection<StateAndRef<ContractState>>, val protocolStates: Collection<String>)
 
 data class ClientToServiceCommandMessage(override val sessionID: Long, override val replyToRecipient: SingleMessageRecipient, val command: ClientToServiceCommand) : DirectRequestMessage

node/src/main/kotlin/com/r3corda/node/services/monitor/WalletMonitorService.kt

@@ -135,7 +135,7 @@ class WalletMonitorService(services: ServiceHubInternal, val smm: StateMachineMa
     fun processRegisterRequest(req: RegisterRequest) {
         try {
             listeners.add(RegisteredListener(req.replyToRecipient, req.sessionID))
-            val stateMessage = StateSnapshotMessage(services.walletService.currentWallet.states.map { it.state.data }.toList(),
+            val stateMessage = StateSnapshotMessage(services.walletService.currentWallet.states.toList(),
                     smm.allStateMachines.map { it.javaClass.name })
             net.send(net.createMessage(STATE_TOPIC, DEFAULT_SESSION_ID, stateMessage.serialize().bits), req.replyToRecipient)