From 20bc37de4ad3d643d636747defbd99914f3ac880 Mon Sep 17 00:00:00 2001
From: Alexey Pashinov <pashinov93@gmail.com>
Date: Fri, 9 Feb 2024 13:35:02 +0100
Subject: [PATCH] Add workchain id to accounts producer data

---
 server/src/lib.rs |  4 +++-
 server/src/ws.rs  | 27 ++++++++++++++++++++-------
 2 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/server/src/lib.rs b/server/src/lib.rs
index fbd1cae..3f99740 100644
--- a/server/src/lib.rs
+++ b/server/src/lib.rs
@@ -253,7 +253,9 @@ impl RpcState {
             storage.update(block_id, block, shard_state)?;
         }
 
-        self.ws_producer.handle_block(block).await?;
+        self.ws_producer
+            .handle_block(block_id.shard_id.workchain_id(), block)
+            .await?;
 
         Ok(())
     }
diff --git a/server/src/ws.rs b/server/src/ws.rs
index 915889f..2993888 100644
--- a/server/src/ws.rs
+++ b/server/src/ws.rs
@@ -8,8 +8,8 @@ use futures_util::{SinkExt, StreamExt};
 use rustc_hash::FxHashMap;
 use serde::{Deserialize, Serialize};
 use tokio::sync::Mutex;
-use ton_block::{Deserializable, HashmapAugType};
-use ton_types::HashmapType;
+use ton_block::{Deserializable, HashmapAugType, MsgAddressInt};
+use ton_types::{AccountId, HashmapType};
 
 use crate::server::Server;
 
@@ -47,15 +47,15 @@ async fn handle_socket(client_id: uuid::Uuid, state: Arc<Server>, socket: WebSoc
 
 #[derive(Default)]
 pub struct WsProducer {
-    clients: Mutex<FxHashMap<uuid::Uuid, SplitSink<WebSocket, axum::extract::ws::Message>>>,
+    clients: Mutex<FxHashMap<uuid::Uuid, SplitSink<WebSocket, Message>>>,
 }
 
 impl WsProducer {
-    pub async fn handle_block(&self, block: &ton_block::Block) -> Result<()> {
+    pub async fn handle_block(&self, workchain_id: i32, block: &ton_block::Block) -> Result<()> {
         let extra = block.read_extra()?;
         let account_blocks = extra.read_account_blocks()?;
 
-        let mut accounts = FxHashMap::default();
+        let mut accounts = Vec::with_capacity(account_blocks.len()?);
         account_blocks.iterate_with_keys(|account, value| {
             let mut lt = 0;
             value.transactions().iterate_slices(|_, mut value| {
@@ -67,7 +67,14 @@ impl WsProducer {
 
                 Ok(true)
             })?;
-            accounts.insert(account.into_vec(), lt);
+
+            let address =
+                MsgAddressInt::with_standart(None, workchain_id as i8, AccountId::from(account))?;
+
+            accounts.push(AccountInfo {
+                account: nekoton_proto::utils::addr_to_bytes(&address).to_vec(),
+                account_lt: lt,
+            });
 
             Ok(true)
         })?;
@@ -75,10 +82,16 @@ impl WsProducer {
         let message = bincode::serialize(&accounts)?;
         let mut clients = self.clients.lock().await;
         for (_, client) in clients.iter_mut() {
-            let message = axum::extract::ws::Message::Binary(message.clone());
+            let message = Message::Binary(message.clone());
             client.send(message).await?;
         }
 
         Ok(())
     }
 }
+
+#[derive(Clone, Serialize, Deserialize)]
+struct AccountInfo {
+    pub account: Vec<u8>,
+    pub account_lt: u64,
+}