deltatrace-deterministic-replayable-late/deltatrace/dsl.py

from __future__ import annotations
from dataclasses import dataclass, field
from typing import Dict, Any, Optional, List
import time

@dataclass
class LocalEvent:
  event_id: str
  event_type: str  # e.g., "MDTick", "Signal", etc.
  asset: str
  venue: str
  timestamp: float
  payload: Dict[str, Any]
  metadata: Dict[str, Any] = field(default_factory=dict)

@dataclass
class PlanDelta:
  delta_id: str
  timestamp: float
  author: str
  contract_id: str
  delta_payload: Dict[str, Any]
  signature: str

@dataclass
class OrderEvent:
  order_id: str
  timestamp: float
  side: str
  quantity: int
  price: float
  provenance: str = ""
  delta_id: Optional[str] = None

@dataclass
class FillEvent:
  fill_id: str
  timestamp: float
  order_id: str
  quantity: int
  price: float
  venue: str = ""
  delta_id: Optional[str] = None

@dataclass
class RiskCheck:
  check_id: str
  timestamp: float
  result: bool
  budget_used: float
  details: Dict[str, Any] = field(default_factory=dict)

@dataclass
class AuditLog:
  log_id: str
  timestamp: float
  action: str
  payload: Dict[str, Any]
  signature: str

@dataclass
class Metadata:
  version: str
  nonce: str
  source_adapter: str

@dataclass
class TraceGraph:
  nodes: List[object] = field(default_factory=list)
  edges: List[tuple] = field(default_factory=list)  # (src_id, dst_id, tag)

  def add_node(self, node: object) -> object:
    self.nodes.append(node)
    return node

  def add_edge(self, src_id: str, dst_id: str, tag: str) -> None:
    self.edges.append((src_id, dst_id, tag))

  def merkle_root(self) -> str:
    """Deterministic Merkle root for the trace graph.

    This implementation avoids Python's non-deterministic hash() across
    runs by serializing each node to a canonical JSON representation using
    its public fields. The resulting bytestrings are hashed in a Merkle-like
    reduction to produce a stable root string.
    """
    import hashlib
    import json

    def _canonicalize(obj: object) -> Dict[str, object]:
      # If the node is a dataclass or has __dict__, extract a stable dict
      if hasattr(obj, "__dict__"):
        # Ensure all nested attributes are also representable
        data = getattr(obj, "__dict__")
        # Deeply serialize to JSON-safe form by iterating keys
        return {k: v for k, v in data.items()}
      # Fallback: convert to string representation
      return {"repr": str(obj)}

    items = []
    for n in self.nodes:
      canon = _canonicalize(n)
      # Use a stable JSON representation of the node
      serialized = json.dumps(canon, sort_keys=True, separators=(",", ":"))
      items.append(serialized.encode())
    if not items:
      return ""
    level = [hashlib.sha256(i).hexdigest() for i in items]
    while len(level) > 1:
      next_level = []
      for i in range(0, len(level), 2):
        left = level[i]
        right = level[i + 1] if i + 1 < len(level) else left
        next_level.append(hashlib.sha256((left + right).encode()).hexdigest())
      level = next_level
    return level[0]

  def merkle_proof(self, index: int) -> List[Dict[str, str]]:
    """Compute a Merkle-like proof for the node at the given index.

    This returns a list of dicts, each describing the sibling hash at a
    particular tree level required to recompute the Merkle root. Each entry
    has:
      - direction: 'left' or 'right' indicating the position of the sibling
      - hash: the hexadecimal digest of the sibling node at that level

    Note: This is a lightweight, deterministic proof suitable for auditing
    provenance in MVP deployments. It relies on a fixed leaf ordering (the
    order of self.nodes) and a deterministic leaf hashing similar to merkle_root.
    """
    import hashlib
    import json

    def _canonicalize(obj: object) -> Dict[str, object]:
      if hasattr(obj, "__dict__"):
        data = getattr(obj, "__dict__")
        return {k: v for k, v in data.items()}
      return {"repr": str(obj)}

    if index < 0 or index >= len(self.nodes):
      raise IndexError("index out of range for merkle_proof")

    # Build leaf hashes from canonicalized node representations
    leaves: List[str] = []
    for n in self.nodes:
      canon = _canonicalize(n)
      serialized = json.dumps(canon, sort_keys=True, separators=(",", ":"))
      leaves.append(hashlib.sha256(serialized.encode()).hexdigest())

    proof: List[Dict[str, str]] = []
    idx = index
    level = leaves
    while len(level) > 1:
      # Sibling hash for the current index
      if idx % 2 == 0:
        sib_index = idx + 1
        direction = "right"
      else:
        sib_index = idx - 1
        direction = "left"
      if sib_index < len(level):
        proof.append({"direction": direction, "hash": level[sib_index]})
      else:
        # If no true sibling (odd last item), replicate the node itself
        proof.append({"direction": direction, "hash": level[idx]})
      # Compute next level
      next_level: List[str] = []
      for i in range(0, len(level), 2):
        left = level[i]
        right = level[i + 1] if i + 1 < len(level) else level[i]
        next_level.append(hashlib.sha256((left + right).encode()).hexdigest())
      level = next_level
      idx = idx // 2

    return proof