implement the TxsGraph class

Author: jonahar

py/parse_simulation_data.py

@@ -7,27 +7,7 @@
 
 from datatypes import BTC, TXID, btc_to_sat
 from paths import LN
-from txs_graph import build_txs_graph, get_downstream
-
-
-def get_htlcs_claimed_by_timeout(graph: DiGraph, commitment_txid: TXID) -> List[TXID]:
-    """
-    return a list of txids that claimed an HTLC output from the given
-    commitment transaction, using timeout-claim (as opposed to success-claim)
-    """
-    
-    # from the bolt:
-    # " HTLC-Timeout and HTLC-Success Transactions... are almost identical,
-    #   except the HTLC-timeout transaction is timelocked "
-    #
-    # i.e. if the child_tx has a non-zero locktime, it is an HTLC-timeout
-    # TODO: what about claiming of local/remote outputs? are they locked? check it
-    
-    return [
-        child_tx
-        for _, child_tx in graph.out_edges(commitment_txid)
-        if graph.nodes[child_tx]["tx"]["locktime"] > 0
-    ]
+from txs_graph.txs_graph import TxsGraph
 
 
 # ---------- txid-to-label functions ----------
@@ -110,19 +90,15 @@ def extract_bitcoin_funding_txids(simulation_outfile: str) -> Set[TXID]:
     return txids
 
 
-def get_all_direct_children(txid, graph: DiGraph) -> Set[TXID]:
-    return {txid for _, txid in graph.out_edges(txid)}
-
-
 def flatten(s: Iterable[Iterable[Any]]) -> List[Any]:
     return list(itertools.chain.from_iterable(s))
 
 
-def find_commitments(simulation_outfile: str, graph: DiGraph) -> List[TXID]:
+def find_commitments(simulation_outfile: str, graph: TxsGraph) -> List[TXID]:
     bitcoin_fundings = extract_bitcoin_funding_txids(simulation_outfile=simulation_outfile)
 
     ln_channel_fundings = flatten(
-        get_all_direct_children(txid, graph=graph)
+        graph.get_all_direct_children(txid)
         for txid in bitcoin_fundings
     )
 
@@ -132,7 +108,7 @@ def find_commitments(simulation_outfile: str, graph: DiGraph) -> List[TXID]:
         list(filter(
             # only keep those with the expected balance
             lambda child_txid: graph.edges[(channel_funding_txid, child_txid)]["value"] == LN_CHANNEL_BALANCE,
-            get_all_direct_children(txid=channel_funding_txid, graph=graph)
+            graph.get_all_direct_children(txid=channel_funding_txid)
         ))
         for channel_funding_txid in ln_channel_fundings
     )
@@ -151,24 +127,28 @@ def find_commitments(simulation_outfile: str, graph: DiGraph) -> List[TXID]:
                 f"Failed to find commitments. "
                 f"txid {commitment_txid[-4:]} expected to have exactly 1 input, but has {num_inputs}"
             )
-
+    
     return commitments
 
 
-def is_replaceable_by_fee(txid, graph: DiGraph) -> bool:
-    for input_dict in graph.nodes[txid]["tx"]["vin"]:
-        if input_dict['sequence'] < (0xffffffff - 1):
-            return True
-    return False
-
-
-def print_nsequence(txids: Iterable[TXID], graph: DiGraph):
-    for txid in txids:
-        print(f"txid {txid_to_short_txid(txid)}:")
-        for i, input_dict in enumerate(graph.nodes[txid]["tx"]["vin"]):
-            sequence: int = input_dict['sequence']
-            sequence_str = format(sequence, 'x')
-            print(f"input {i}: sequence={sequence_str}")
+def get_htlcs_claimed_by_timeout(graph: TxsGraph, commitment_txid: TXID) -> List[TXID]:
+    """
+    return a list of txids that claimed an HTLC output from the given
+    commitment transaction, using timeout-claim (as opposed to success-claim)
+    """
+    
+    # from the bolt:
+    # " HTLC-Timeout and HTLC-Success Transactions... are almost identical,
+    #   except the HTLC-timeout transaction is timelocked "
+    #
+    # i.e. if the child_tx has a non-zero locktime, it is an HTLC-timeout
+    # TODO: what about claiming of local/remote outputs? are they locked? check it
+    
+    return [
+        child_tx
+        for _, child_tx in graph.out_edges(commitment_txid)
+        if graph.nodes[child_tx]["tx"]["locktime"] > 0
+    ]
 
 
 def main(simulation_name):
@@ -177,26 +157,27 @@ def main(simulation_name):
     outfile = os.path.join(LN, "simulations", f"{simulation_name}.out")
     dotfile = os.path.join(LN, f"{simulation_name}.dot")
     jpgfile = os.path.join(LN, f"{simulation_name}.jpg")
-    
-    txs_graph = build_txs_graph(datadir)
+
+    txs_graph = TxsGraph.from_datadir(datadir)
 
     commitments = find_commitments(simulation_outfile=outfile, graph=txs_graph)
 
     for commitment_txid in commitments:
         short_txid = txid_to_short_txid(commitment_txid)
         num_outputs = len(txs_graph.nodes[commitment_txid]["tx"]["vout"])
-        replaceable = is_replaceable_by_fee(txid=commitment_txid, graph=txs_graph)
-        num_htlcs_stolen = len(get_htlcs_claimed_by_timeout(commitment_txid=commitment_txid, graph=txs_graph))
+        replaceable = txs_graph.is_replaceable_by_fee(txid=commitment_txid)
+        nsequence = txs_graph.get_minimal_nsequence(commitment_txid)
+        num_htlcs_stolen = len(get_htlcs_claimed_by_timeout(graph=txs_graph, commitment_txid=commitment_txid))
         print(
             f"commitment: {short_txid:<5} "
             f"num-outputs: {num_outputs:<4} "
             f"replaceable: {str(replaceable):<5} "
+            f"nsequence: {nsequence:x}  "
             f"htlcs-stolen: {num_htlcs_stolen}"
         )
 
     # this graph includes only interesting txs (no coinbase and other junk)
-    restricted_txs_graph = get_downstream(
-        graph=txs_graph,
+    restricted_txs_graph = txs_graph.get_downstream(
         sources=extract_bitcoin_funding_txids(simulation_outfile=outfile),
     )
 

py/tests/txs_graph_test.py

@@ -1,54 +1,69 @@
 import unittest
 
-from networkx import DiGraph
-
-from txs_graph import get_downstream
+from txs_graph.txs_graph import TxsGraph
 
 
 class TxsGraphTest(unittest.TestCase):
-    def test_downstream_1(self):
-        graph = DiGraph()
+    
+    def get_test_graph_1(self) -> TxsGraph:
+        graph = TxsGraph()
         graph.add_edge("a", "b")
         graph.add_edge("b", "c")
         graph.add_edge("c", "e")
         graph.add_edge("e", "d")
         graph.add_edge("d", "b")
         graph.add_edge("e", "f")
-        
-        downstream = get_downstream(graph, sources={"a"})
+        return graph
+    
+    def test_downstream_1_1(self):
+        graph = self.get_test_graph_1()
+        downstream = graph.get_downstream(sources={"a"})
         downstream_nodes = set(downstream.nodes)
         self.assertSetEqual(downstream_nodes, {"a", "b", "c", "d", "e", "f"})
-        
-        downstream = get_downstream(graph, sources={"b", "e"})
+    
+    def test_downstream_1_2(self):
+        graph = self.get_test_graph_1()
+        downstream = graph.get_downstream(sources={"b", "e"})
         downstream_nodes = set(downstream.nodes)
         self.assertSetEqual(downstream_nodes, {"b", "c", "d", "e", "f"})
-        
-        downstream = get_downstream(graph, sources={"e"})
+    
+    def test_downstream_1_3(self):
+        graph = self.get_test_graph_1()
+        downstream = graph.get_downstream(sources={"e"})
         downstream_nodes = set(downstream.nodes)
         self.assertSetEqual(downstream_nodes, {"b", "c", "d", "e", "f"})
-        
-        downstream = get_downstream(graph, sources={"f"})
+    
+    def test_downstream_1_4(self):
+        graph = self.get_test_graph_1()
+        downstream = graph.get_downstream(sources={"f"})
         downstream_nodes = set(downstream.nodes)
         self.assertSetEqual(downstream_nodes, {"f"})
 
-    def test_downstream_2(self):
-        graph = DiGraph()
+    def get_test_graph_2(self) -> TxsGraph:
+        graph = TxsGraph()
         graph.add_edge(1, 3)
         graph.add_edge(1, 4)
         graph.add_edge(2, 4)
         graph.add_edge(2, 5)
         graph.add_edge(4, 6)
         graph.add_edge(5, 6)
-        
-        downstream = get_downstream(graph, sources={1, 2})
+        return graph
+    
+    def test_downstream_2_1(self):
+        graph = self.get_test_graph_2()
+        downstream = graph.get_downstream(sources={1, 2})
         downstream_nodes = set(downstream.nodes)
         self.assertSetEqual(downstream_nodes, {1, 2, 3, 4, 5, 6})
-        
-        downstream = get_downstream(graph, sources={2})
+    
+    def test_downstream_2_2(self):
+        graph = self.get_test_graph_2()
+        downstream = graph.get_downstream(sources={2})
         downstream_nodes = set(downstream.nodes)
         self.assertSetEqual(downstream_nodes, {2, 4, 5, 6})
-        
-        downstream = get_downstream(graph, sources={3, 4})
+    
+    def test_downstream_2_3(self):
+        graph = self.get_test_graph_2()
+        downstream = graph.get_downstream(sources={3, 4})
         downstream_nodes = set(downstream.nodes)
         self.assertSetEqual(downstream_nodes, {3, 4, 6})
 

py/txs_graph/txs_graph.py

@@ -0,0 +1,88 @@
+from typing import Any, Iterable, List
+
+from networkx.algorithms.traversal.breadth_first_search import bfs_edges
+from networkx.classes.digraph import DiGraph
+
+from datatypes import TXID
+from txs_graph.txs_graph_utils import find_tx_fee, load_blocks, load_txs
+
+
+class TxsGraph(DiGraph):
+    
+    @staticmethod
+    def from_datadir(datadir: str) -> "TxsGraph":
+        """
+        read all block and transaction files in the given datadir and construct
+        a full transaction graph.
+        
+        Each node represents a transaction. the node's id is the txid and it has
+        the following attributes:
+            - "tx" - the full tx json, as returned by bitcoind
+            - "fee" - the tx fee
+            - "height" - the block height in which the tx was included
+        
+        Each edge has the following attributes:
+            - "value" the value in BTC of the output represented by this edge
+            
+        """
+        blocks = load_blocks(datadir)
+        txs = load_txs(datadir)
+        
+        txid_to_fee = {txid: find_tx_fee(txid, txs) for txid in txs.keys()}
+        txid_to_height = {
+            txid: block["height"]
+            for block in blocks.values()
+            for txid in block["tx"]
+        }
+        
+        graph = TxsGraph()
+        
+        # add all transactions
+        for txid in txs.keys():
+            graph.add_node(txid, tx=txs[txid], fee=txid_to_fee[txid], height=txid_to_height[txid])
+        
+        # add edges between transactions
+        for dest_txid, dest_tx in txs.items():
+            for entry in dest_tx["vin"]:
+                if "coinbase" in entry:
+                    continue  # coinbase transaction. no src
+                src_txid = entry["txid"]
+                index = entry["vout"]
+                value = txs[src_txid]["vout"][index]["value"]
+                graph.add_edge(src_txid, dest_txid, value=value)
+        
+        return graph
+    
+    def get_all_direct_children(self, txid: TXID) -> List[TXID]:
+        return [txid for _, txid in self.out_edges(txid)]
+    
+    def get_minimal_nsequence(self, txid: TXID) -> int:
+        """
+        return the minimal nsequence of an input in the given txid
+        """
+        return min(map(
+            lambda input_dict: input_dict['sequence'],
+            self.nodes[txid]["tx"]["vin"]
+        ))
+    
+    def is_replaceable_by_fee(self, txid: TXID) -> bool:
+        return self.get_minimal_nsequence(txid) < (0xffffffff - 1)
+    
+    def get_downstream(self, sources: Iterable[Any]) -> "TxsGraph":
+        """
+        return the downstream of sources in the given graph.
+        sources must be an iterable of existing node ids in the graph
+        """
+        sources = list(sources)  # sources may be iterable only once (e.g. map), so we copy
+        
+        downstream = TxsGraph()
+        downstream.add_nodes_from(
+            map(lambda src: (src, self.nodes[src]), sources)
+        )
+        for src in sources:
+            for v, u in bfs_edges(self, source=src):
+                downstream.add_node(v, **self.nodes[v])
+                downstream.add_node(u, **self.nodes[u])
+                downstream.add_edge(v, u, **self.edges[(v, u)])
+        
+        return downstream

py/txs_graph.py py/txs_graph/txs_graph_utils.py

@@ -1,12 +1,15 @@
 import json
 import os
-from typing import Any, Dict, Iterable
+from typing import Dict
 
-from networkx.algorithms.traversal.breadth_first_search import bfs_edges
 from networkx.classes.digraph import DiGraph
 
 from datatypes import BTC, Block, BlockHash, TX, TXID
 
+"""
+A collection of helper functions to build a complete TxsGraph
+"""
+
 
 def load_blocks(datadir: str) -> Dict[BlockHash, Block]:
     blocks = {}
@@ -97,18 +100,3 @@ def build_txs_graph(datadir: str) -> DiGraph:
             graph.add_edge(src_txid, dest_txid, value=value)
 
     return graph
-
-
-def get_downstream(graph: DiGraph, sources: Iterable[Any]) -> DiGraph:
-    """
-    return the downstream of sources in the given graph.
-    sources must be an iterable of existing node ids in the graph
-    """
-    sources = list(sources)  # sources may be iterable only once (e.g. map), so we copy
-    nodes_to_include = set(sources)
-    
-    # for each source, compute the set of its ancestors using bfs
-    nodes_to_include.update(
-        *({u for v, u in bfs_edges(graph, source=src)} for src in sources)
-    )
-    return graph.subgraph(nodes=nodes_to_include)