ChainGrapher/src/interfaces/iota_interface.py at main · dbiton/ChainGrapher · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
from typing import Dict, Set, List
from interfaces.interface import Interface
import os
from collections import Counter

# This is for IOTA Rebased (after switching from tangle to MoveVM)
class IotaInterface(Interface):

    def __init__(self):
        rpc_url = os.getenv("IOTA_RPC_URL")
        super().__init__(True, rpc_url)

    def _fetch_checkpoint(self, checkpoint_id: int) -> dict:
        payload = {
            "jsonrpc": "2.0",
            "id": 1,
            "method": "iota_getCheckpoint",
            "params": [str(checkpoint_id)]
        }
        response = self._post_with_retry(payload)
        return response

    def _fetch_txs(self, txs_ids: List[str]) -> List[dict]:
        all_results = []
        for i in range(0, len(txs_ids), 50):
            batch = txs_ids[i:i + 50]
            payload = {
                "jsonrpc": "2.0",
                "id": 1,
                "method": "iota_multiGetTransactionBlocks",
                "params": [
                    batch,
                    {
                        "showInput": True,
                        "showRawInput": True,
                        "showEffects": True,
                        "showEvents": True,
                        "showObjectChanges": True,
                        "showBalanceChanges": True,
                        "showRawEffects": True
                    }
                ]
            }
            batch_results = self._post_with_retry(payload)
            all_results.extend(batch_results)
        return all_results

    def fetch(self, checkpoint_id):
        checkpoint = self._fetch_checkpoint(checkpoint_id)
        if not checkpoint:
            print(f"Checkpoint {checkpoint_id} | Checkpoint not found. Exiting.")
            exit()
        txs_ids = checkpoint.get("transactions", [])
        print(f"Checkpoint {checkpoint_id} | {len(txs_ids)} transactions")
        txs = self._fetch_txs(txs_ids)
        if not txs:
            print(f"Checkpoint {checkpoint_id} | Transactions not found. Exiting.")
            exit()
        return checkpoint_id, checkpoint, txs

    def get_conflict_graph(self, checkpoint_trace):
        checkpoint_trace, txs_traces = checkpoint_trace
        writes: Dict[str, Set[str]] = {}
        reads: Dict[str, Set[str]] = {}
        for tx_trace in txs_traces:
            tx_id = tx_trace['digest']
            tx_reads, tx_writes = self._parse_tx(tx_trace)
            reads[tx_id] = tx_reads
            writes[tx_id] = tx_writes
        txs = [tx_trace["digest"] for tx_trace in txs_traces]
        return self._create_conflict_graph_from_readset_writeset(txs, reads, writes)

    def _get_tx_type(self, tx):
        return tx['transaction']['data']['transaction']['kind']

    def get_additional_metrics(self, block_number, trace) -> Dict[str, float]:
        checkpoint_trace, txs_traces = trace

        txs_types = [self._get_tx_type(tx) for tx in txs_traces]
        txs_type_counter = Counter(txs_types)

        user_kinds = {
            "ProgrammableTransaction", "TransferObject", "TransferSui",
            "Pay", "PaySui", "PayAllSui", "SplitCoin", "MergeCoin", "Publish"
        }

        system_kinds = {
            "ConsensusCommitPrologue", "ConsensusCommitPrologueV1",
            "ChangeEpoch", "Genesis", "RandomnessStateUpdate"
        }

        return {
            "user_tx_count": sum(txs_type_counter.get(k, 0) for k in user_kinds),
            "system_tx_count": sum(txs_type_counter.get(k, 0) for k in system_kinds),
            "block_number": block_number,
            "txs": len(txs_traces)
        }

    def _parse_tx(self, tx):
        write_addrs = {
            change['objectId']
            for change in tx.get('objectChanges', [])
            if 'objectId' in change
        }
        shared_addrs = {
            obj['objectId']
            for obj in tx.get('effects', {}).get('sharedObjects', [])
        }
        inputs = (
            tx.get('transaction', {})
            .get('data', {})
            .get('transaction', {})
            .get('inputs', [])
        )
        inputs_addrs = {
            inp['objectId']
            for inp in inputs
            if inp.get('type') == 'object'
        }
        read_addrs = (inputs_addrs | shared_addrs) - write_addrs
        return read_addrs, write_addrs