From 4f96bd186eca93370dbdcc3e0588476c4ec1bae2 Mon Sep 17 00:00:00 2001
From: Juliya Smith <jules@apeworx.io>
Date: Mon, 24 Jun 2024 16:36:44 -0500
Subject: [PATCH] feat: wip events in calltrees

---
 evm_trace/base.py    | 10 ++++++++--
 evm_trace/display.py | 42 +++++++++++++++++++++++++++++++++++++++---
 evm_trace/geth.py    |  8 +++++++-
 3 files changed, 54 insertions(+), 6 deletions(-)

diff --git a/evm_trace/base.py b/evm_trace/base.py
index 77a5f0b..8e39d3d 100644
--- a/evm_trace/base.py
+++ b/evm_trace/base.py
@@ -21,12 +21,18 @@ class Event(BaseModel):
     An event emitted during a CALL.
     """
 
-    calldata: HexBytes = HexBytes("")
-    """Event calldata (inputs)."""
+    data: list[HexBytes] = []
+    """The remaining event data besides the topics."""
 
     selector: HexBytes
     """The selector hash of the event."""
 
+    topics: list[HexBytes] = []
+    """Event topics."""
+    #
+    # def __repr__(self) -> str:
+    #     return f"{self.calldata.hex()}"
+
 
 class CallTreeNode(BaseModel):
     """
diff --git a/evm_trace/display.py b/evm_trace/display.py
index fb53d9d..e6ca93c 100644
--- a/evm_trace/display.py
+++ b/evm_trace/display.py
@@ -15,8 +15,16 @@ def get_tree_display(call: "CallTreeNode") -> str:
 
 
 class TreeRepresentation:
-    FILE_MIDDLE_PREFIX = "├──"
-    FILE_LAST_PREFIX = "└──"
+    """
+    A class for creating a simple tree-representation of a call-tree node.
+
+    **NOTE**: We purposely are not using the rich library here to keep
+    evm-trace small and simple while sill offering a nice stringified
+    version of a :class:`~evm_trace.base.CallTreeNode`.
+    """
+
+    MIDDLE_PREFIX = "├──"
+    LAST_PREFIX = "└──"
     PARENT_PREFIX_MIDDLE = "    "
     PARENT_PREFIX_LAST = "│   "
 
@@ -32,10 +40,17 @@ def __init__(
 
     @property
     def depth(self) -> int:
+        """
+        The depth in the call tree, such as the
+        number of calls deep.
+        """
         return self.call.depth
 
     @property
     def title(self) -> str:
+        """
+        The title of the node representation, including address, calldata, and return-data.
+        """
         call_type = self.call.call_type.value
         address_hex_str = self.call.address.hex() if self.call.address else None
 
@@ -81,6 +96,15 @@ def make_tree(
         parent: Optional["TreeRepresentation"] = None,
         is_last: bool = False,
     ) -> Iterator["TreeRepresentation"]:
+        """
+        Create a node representation object from a :class:`~evm_trace.base.CallTreeNode`.
+
+        Args:
+            root (:class:`~evm_trace.base.CallTreeNode`): The call-tree node to display.
+            parent (Optional[:class:`~evm_trace.display.TreeRepresentation`]): The parent
+              node of this node.
+            is_last (bool): True if a leaf-node.
+        """
         displayable_root = cls(root, parent=parent, is_last=is_last)
         yield displayable_root
 
@@ -95,15 +119,27 @@ def make_tree(
             count += 1
 
     def __str__(self) -> str:
+        """
+        The representation str via ``calling str()``.
+        :return:
+        """
         if self.parent is None:
             return self.title
 
-        filename_prefix = self.FILE_LAST_PREFIX if self.is_last else self.FILE_MIDDLE_PREFIX
+        filename_prefix = self.LAST_PREFIX if self.is_last else self.MIDDLE_PREFIX
 
         parts = [f"{filename_prefix} {self.title}"]
+
+        # if events := self.call.events:
+        # TODO!
+        # events_str = "\n".join([str(e) for e in events])
+
         parent = self.parent
         while parent and parent.parent is not None:
             parts.append(self.PARENT_PREFIX_MIDDLE if parent.is_last else self.PARENT_PREFIX_LAST)
             parent = parent.parent
 
         return "".join(reversed(parts))
+
+    def __repr__(self) -> str:
+        return str(self)
diff --git a/evm_trace/geth.py b/evm_trace/geth.py
index 984d080..0d086f2 100644
--- a/evm_trace/geth.py
+++ b/evm_trace/geth.py
@@ -300,9 +300,15 @@ def _create_node(
         elif frame.op.startswith("LOG") and len(frame.op) > 3 and frame.op[3].isnumeric():
             num_topics = int(frame.op[3])
             start_calldata_idx = -num_topics - 3
+
+            # TODO: Test this more.
+            data_key = frame.stack[start_calldata_idx]
+            data_idx = frame.memory.root.index(data_key) + 1
+            event_data = frame.memory.root[data_idx:]
             event = Event(
+                data=event_data,
+                topics=[HexBytes(t) for t in frame.stack[start_calldata_idx:-4]],
                 selector=frame.stack[-num_topics],
-                calldata=HexBytes(b"".join(frame.stack[start_calldata_idx:-4])),
             )
             if "events" in node_kwargs:
                 node_kwargs["events"].append(event)