feat: add coverage

convince coverage.py to cover vyper contracts by inspecting the frame.
thread filename through the Contract data structure so we can report on
it.

.gitignore

@@ -10,3 +10,5 @@ dist/
 .pytest_cache
 venv/
 .env
+
+.coverage

boa/coverage.py

@@ -0,0 +1,84 @@
+"""The titanoboa coverage plugin."""
+
+import coverage.plugin
+from vyper.ir import compile_ir
+
+import boa.interpret
+from boa.environment import Env, TracingCodeStream
+
+
+def coverage_init(registry, options):
+    registry.add_file_tracer(TitanoboaPlugin(options))
+
+
+class TitanoboaPlugin(coverage.plugin.CoveragePlugin):
+    def __init__(self, options):
+        pass
+
+    def file_tracer(self, filename):
+        if filename.endswith("boa/environment.py"):
+            return TitanoboaTracer()
+
+    def file_reporter(self, filename):
+        if filename.endswith(".vy"):
+            return TitanoboaReporter(filename)
+
+
+class TitanoboaTracer(coverage.plugin.FileTracer):
+    def __init__(self, env=None):
+        self.env = env or Env.get_singleton()
+
+    def _contract_for_frame(self, frame):
+        if frame.f_code.co_qualname != TracingCodeStream.__iter__.__qualname__:
+            return None
+        return frame.f_locals["self"]._contract
+
+    def dynamic_source_filename(self, filename, frame):
+        contract = self._contract_for_frame(frame)
+        if contract is None:
+            return None
+
+        return contract.filename
+
+    def has_dynamic_source_filename(self):
+        return True
+
+    # coverage.py requires us to inspect the python call frame to
+    # see what line number to produce. since every loop of the py-evm
+    # event loop uses CodeStream.__iter__, we intercept all calls to
+    # __iter__, and then back out the contract and lineno information
+    # from there.
+    def line_number_range(self, frame):
+        contract = self._contract_for_frame(frame)
+        if contract is None:
+            return super().line_number_range(frame)
+
+        pc = frame.f_locals["self"].program_counter
+        pc_map = contract.source_map["pc_pos_map"]
+
+        if (src_loc := pc_map.get(pc)) is None:
+            return (-1, -1)
+
+        (start_lineno, _, end_lineno, _) = src_loc
+        return start_lineno, end_lineno
+
+
+class TitanoboaReporter(coverage.plugin.FileReporter):
+    def __init__(self, filename, env=None):
+        super().__init__(filename)
+
+    def lines(self):
+        ret = set()
+        c = boa.interpret.compiler_data(self.source(), self.filename)
+
+        # source_map should really be in CompilerData
+        _, source_map = compile_ir.assembly_to_evm(c.assembly_runtime)
+
+        for _, v in source_map["pc_pos_map"].items():
+            if v is None:
+                continue
+            (start_lineno, _, end_lineno, _) = v
+            for i in range(start_lineno, end_lineno + 1):
+                ret.add(i)
+
+        return ret

boa/environment.py

@@ -156,12 +156,14 @@ class TracingCodeStream(CodeStream):
         "invalid_positions",
         "valid_positions",
         "program_counter",
+        "_contract",
     ]
 
-    def __init__(self, *args, start_pc=0, fake_codesize=None, **kwargs):
+    def __init__(self, *args, start_pc=0, fake_codesize=None, contract=None, **kwargs):
         super().__init__(*args, **kwargs)
         self._trace = []  # trace of opcodes that were run
         self.program_counter = start_pc  # configurable start PC
+        self._contract = contract  # for coverage
         self._fake_codesize = fake_codesize  # what CODESIZE returns
 
     def __iter__(self) -> Iterator[int]:
@@ -278,6 +280,7 @@ def __init__(self):
         self._profiled_contracts = {}
         self._cached_call_profiles = {}
         self._cached_line_profiles = {}
+        self._coverage_data = {}
 
         self.sha3_trace = {}
         self.sstore_trace = {}
@@ -298,8 +301,12 @@ def __init__(self, *args, **kwargs):
                 # super() hardcodes CodeStream into the ctor
                 # so we have to override it here
                 super().__init__(*args, **kwargs)
+                contract = env._contracts.get(
+                    to_checksum_address(self.msg.storage_address)
+                )
                 self.code = TracingCodeStream(
                     self.code._raw_code_bytes,
+                    contract=contract,
                     fake_codesize=getattr(self.msg, "_fake_codesize", None),
                     start_pc=getattr(self.msg, "_start_pc", 0),
                 )

boa/interpret.py

@@ -59,11 +59,11 @@ def load(filename: str, *args, **kwargs) -> _Contract:  # type: ignore
     if "name" in kwargs:
         name = kwargs.pop("name")
     with open(filename) as f:
-        return loads(f.read(), *args, name=name, **kwargs)
+        return loads(f.read(), *args, name=name, **kwargs, filename=filename)
 
 
-def loads(source_code, *args, as_blueprint=False, name=None, **kwargs):
-    d = loads_partial(source_code, name)
+def loads(source_code, *args, as_blueprint=False, name=None, filename=None, **kwargs):
+    d = loads_partial(source_code, name, filename=filename)
     if as_blueprint:
         return d.deploy_as_blueprint(**kwargs)
     else:
@@ -82,18 +82,18 @@ def loads_abi(json_str: str, *args, name: str = None, **kwargs) -> ABIContractFa
 
 
 def loads_partial(
-    source_code: str, name: str = None, dedent: bool = True
+    source_code: str, name: str = None, filename: str = None, dedent: bool = True
 ) -> VyperDeployer:
     name = name or "VyperContract"  # TODO handle this upstream in CompilerData
     if dedent:
         source_code = textwrap.dedent(source_code)
     data = compiler_data(source_code, name)
-    return VyperDeployer(data)
+    return VyperDeployer(data, filename=filename)
 
 
 def load_partial(filename: str) -> VyperDeployer:  # type: ignore
     with open(filename) as f:
-        return loads_partial(f.read(), name=filename)
+        return loads_partial(f.read(), name=filename, filename=filename)
 
 
 __all__ = ["BoaError"]

boa/vyper/contract.py

@@ -55,26 +55,35 @@
 
 
 class VyperDeployer:
-    def __init__(self, compiler_data):
+    def __init__(self, compiler_data, filename=None):
         self.compiler_data = compiler_data
 
         # force compilation so that if there are any errors in the contract,
         # we fail at load rather than at deploy time.
         _ = compiler_data.bytecode
 
+        self.filename = filename
+
     def __call__(self, *args, **kwargs):
         return self.deploy(*args, **kwargs)
 
     def deploy(self, *args, **kwargs):
-        return VyperContract(self.compiler_data, *args, **kwargs)
+        return VyperContract(
+            self.compiler_data, *args, filename=self.filename, **kwargs
+        )
 
     def deploy_as_blueprint(self, *args, **kwargs):
-        return VyperBlueprint(self.compiler_data, *args, **kwargs)
+        return VyperBlueprint(
+            self.compiler_data, *args, filename=self.filename, **kwargs
+        )
 
     def at(self, address: AddressType) -> "VyperContract":
         address = to_checksum_address(address)
         ret = VyperContract(
-            self.compiler_data, override_address=address, skip_initcode=True
+            self.compiler_data,
+            override_address=address,
+            skip_initcode=True,
+            filename=self.filename,
         )
         vm = ret.env.vm
         bytecode = vm.state.get_code(to_canonical_address(address))
@@ -88,14 +97,16 @@ def at(self, address: AddressType) -> "VyperContract":
 
 # a few lines of shared code between VyperBlueprint and VyperContract
 class _BaseContract:
-    def __init__(self, compiler_data, env=None):
+    def __init__(self, compiler_data, env=None, filename=None):
         self.compiler_data = compiler_data
 
         if env is None:
             env = Env.get_singleton()
 
         self.env = env
 
+        self.filename = filename
+
 
 # create a blueprint for use with `create_from_blueprint`.
 # uses a ERC5202 preamble, when calling `create_from_blueprint` will
@@ -107,10 +118,11 @@ def __init__(
         env=None,
         override_address=None,
         blueprint_preamble=b"\xFE\x71\x00",
+        filename=None,
     ):
         # note slight code duplication with VyperContract ctor,
         # maybe use common base class?
-        super().__init__(compiler_data, env)
+        super().__init__(compiler_data, env, filename)
 
         if blueprint_preamble is None:
             blueprint_preamble = b""
@@ -133,7 +145,7 @@ def __init__(
 
     @cached_property
     def deployer(self):
-        return VyperDeployer(self.compiler_data)
+        return VyperDeployer(self.compiler_data, filename=self.filename)
 
 
 class FrameDetail(dict):
@@ -435,8 +447,9 @@ def __init__(
         # whether or not to skip constructor
         skip_initcode=False,
         created_from=None,
+        filename=None,
     ):
-        super().__init__(compiler_data, env)
+        super().__init__(compiler_data, env, filename)
 
         self.created_from = created_from
 
@@ -516,7 +529,7 @@ def __repr__(self):
 
     @cached_property
     def deployer(self):
-        return VyperDeployer(self.compiler_data, env=self.env)
+        return VyperDeployer(self.compiler_data, env=self.env, filename=self.filename)
 
     # is this actually useful?
     def at(self, address):
@@ -689,6 +702,7 @@ def line_profile(self, computation=None):
         ret = LineProfile.from_single(self, computation)
         for child in computation.children:
             child_obj = self.env.lookup_contract(child.msg.code_address)
+            # TODO: child obj is opaque contract that calls back into known contract
             if child_obj is not None:
                 ret.merge(child_obj.line_profile(child))
         return ret