From c5bb1d3fca39670db7a7a9b1699f2d7f3a3c5749 Mon Sep 17 00:00:00 2001
From: whitequark <cz@m-labs.hk>
Date: Sat, 19 Jan 2019 00:08:51 +0000
Subject: [PATCH] hdl.ast: give Assert and Assume their own src_loc.

This helps with patterns like `Assert(fsm.ongoing("IDLE"))`, which
would otherwise point into nMigen internals.
---
 nmigen/back/rtlil.py |  4 ++--
 nmigen/hdl/ast.py    | 38 ++++++++++++--------------------------
 nmigen/hdl/xfrm.py   | 15 +++++++++------
 3 files changed, 23 insertions(+), 34 deletions(-)

diff --git a/nmigen/back/rtlil.py b/nmigen/back/rtlil.py
index 7c0dda7..4d78b82 100644
--- a/nmigen/back/rtlil.py
+++ b/nmigen/back/rtlil.py
@@ -617,7 +617,7 @@ class _StatementCompiler(xfrm.StatementVisitor):
         self.state.rtlil.cell("$assert", ports={
             "\\A": check_wire,
             "\\EN": en_wire,
-        }, src=src(stmt.test.src_loc))
+        }, src=src(stmt.src_loc))
 
     def on_Assume(self, stmt):
         self(stmt._check.eq(stmt.test))
@@ -628,7 +628,7 @@ class _StatementCompiler(xfrm.StatementVisitor):
         self.state.rtlil.cell("$assume", ports={
             "\\A": check_wire,
             "\\EN": en_wire,
-        }, src=src(stmt.test.src_loc))
+        }, src=src(stmt.src_loc))
 
     def on_Switch(self, stmt):
         self._check_rhs(stmt.test)
diff --git a/nmigen/hdl/ast.py b/nmigen/hdl/ast.py
index 3ae5370..bd66d45 100644
--- a/nmigen/hdl/ast.py
+++ b/nmigen/hdl/ast.py
@@ -910,19 +910,21 @@ class Assign(Statement):
         return "(eq {!r} {!r})".format(self.lhs, self.rhs)
 
 
-class Assert(Statement):
+class Property(Statement):
     def __init__(self, test, _check=None, _en=None):
+        self.src_loc = tracer.get_src_loc()
+
         self.test = Value.wrap(test)
 
         self._check = _check
         if self._check is None:
-            self._check = Signal(reset_less=True, name="$assert$check")
-            self._check.src_loc = self.test.src_loc
+            self._check = Signal(reset_less=True, name="${}$check".format(self._kind))
+            self._check.src_loc = self.src_loc
 
         self._en = _en
         if _en is None:
-            self._en = Signal(reset_less=True, name="$assert$en")
-            self._en.src_loc = self.test.src_loc
+            self._en = Signal(reset_less=True, name="${}$en".format(self._kind))
+            self._en.src_loc = self.src_loc
 
     def _lhs_signals(self):
         return ValueSet((self._en, self._check))
@@ -931,31 +933,15 @@ class Assert(Statement):
         return self.test._rhs_signals()
 
     def __repr__(self):
-        return "(assert {!r})".format(self.test)
-
-
-class Assume(Statement):
-    def __init__(self, test, _check=None, _en=None):
-        self.test = Value.wrap(test)
-
-        self._check = _check
-        if self._check is None:
-            self._check = Signal(reset_less=True, name="$assume$check")
-            self._check.src_loc = self.test.src_loc
+        return "({} {!r})".format(self._kind, self.test)
 
-        self._en = _en
-        if self._en is None:
-            self._en = Signal(reset_less=True, name="$assume$en")
-            self._en.src_loc = self.test.src_loc
 
-    def _lhs_signals(self):
-        return ValueSet((self._en, self._check))
+class Assert(Property):
+    _kind = "assert"
 
-    def _rhs_signals(self):
-        return self.test._rhs_signals()
 
-    def __repr__(self):
-        return "(assume {!r})".format(self.test)
+class Assume(Property):
+    _kind = "assume"
 
 
 class Switch(Statement):
diff --git a/nmigen/hdl/xfrm.py b/nmigen/hdl/xfrm.py
index 473139b..ca43cf5 100644
--- a/nmigen/hdl/xfrm.py
+++ b/nmigen/hdl/xfrm.py
@@ -190,18 +190,21 @@ class StatementVisitor(metaclass=ABCMeta):
 
     def on_statement(self, stmt):
         if type(stmt) is Assign:
-            return self.on_Assign(stmt)
+            new_stmt = self.on_Assign(stmt)
         elif type(stmt) is Assert:
-            return self.on_Assert(stmt)
+            new_stmt = self.on_Assert(stmt)
         elif type(stmt) is Assume:
-            return self.on_Assume(stmt)
+            new_stmt = self.on_Assume(stmt)
         elif isinstance(stmt, Switch):
             # Uses `isinstance()` and not `type() is` because nmigen.compat requires it.
-            return self.on_Switch(stmt)
+            new_stmt = self.on_Switch(stmt)
         elif isinstance(stmt, Iterable):
-            return self.on_statements(stmt)
+            new_stmt = self.on_statements(stmt)
         else:
-            return self.on_unknown_statement(stmt)
+            new_stmt = self.on_unknown_statement(stmt)
+        if hasattr(stmt, "src_loc") and hasattr(new_stmt, "src_loc"):
+            new_stmt.src_loc = stmt.src_loc
+        return new_stmt
 
     def __call__(self, value):
         return self.on_statement(value)
-- 
2.30.2