Save hashes to optimise merkle proof calcs.

This commit is contained in:
FiveMovesAhead 2024-12-26 21:51:03 +08:00
parent c8e437854a
commit 53a9d252bc
2 changed files with 12 additions and 13 deletions

View File

@ -10,7 +10,7 @@ import time
import zlib
from threading import Thread
from common.structs import OutputData, MerkleProof
from common.merkle_tree import MerkleTree
from common.merkle_tree import MerkleTree, MerkleHash
logger = logging.getLogger(os.path.splitext(os.path.basename(__file__))[0])
PENDING_BATCH_IDS = set()
@ -126,22 +126,21 @@ def send_results(session, master_ip, master_port, tig_worker_path, download_wasm
time.sleep(2)
else:
with open(f"{output_folder}/hashes.zlib", "rb") as f:
hashes = json.loads(zlib.decompress(f.read()).decode())
with open(f"{output_folder}/data.zlib", "rb") as f:
leafs = [
OutputData.from_dict(x)
for x in json.loads(zlib.decompress(f.read()).decode())
]
leafs = json.loads(zlib.decompress(f.read()).decode())
merkle_tree = MerkleTree(
[x.to_merkle_hash() for x in leafs],
[MerkleHash.from_str(x) for x in hashes],
batch["batch_size"]
)
proofs_to_submit = [
MerkleProof(
dict(
leaf=leafs[n - batch["start_nonce"]],
branch=merkle_tree.calc_merkle_branch(branch_idx=n - batch["start_nonce"])
).to_dict()
branch=merkle_tree.calc_merkle_branch(branch_idx=n - batch["start_nonce"]).to_str()
)
for n in batch["sampled_nonces"]
]

View File

@ -4,7 +4,7 @@ use clap::{arg, Command};
use futures::stream::{self, StreamExt};
use serde_json::json;
use std::{fs, path::PathBuf, sync::Arc};
use tig_structs::core::BenchmarkSettings;
use tig_structs::core::*;
use tig_utils::{compress_obj, dejsonify, jsonify, MerkleHash, MerkleTree};
use tig_worker::OutputData;
use tokio::runtime::Runtime;
@ -272,8 +272,8 @@ fn compute_batch(
}
if let Some(path) = output_folder {
dump.sort_by_key(|data| data.nonce);
let file_path = path.join("data.zlib");
fs::write(&file_path, compress_obj(&dump))?;
fs::write(&path.join("data.zlib"), compress_obj(&dump))?;
fs::write(&path.join("hashes.zlib"), compress_obj(&hashes))?;
}
let tree = MerkleTree::new(hashes, batch_size as usize)?;