From c5c4679532bd971ddb0818222823b5bdb654469b Mon Sep 17 00:00:00 2001
From: "William D. Jones" <thor0505@comcast.net>
Date: Sun, 30 Dec 2018 05:17:39 -0500
Subject: [PATCH] hdl.xfrm: Add Assert and Assume abstract methods for
 StatementVisitor, implement for children.

---
 nmigen/back/pysim.py |  6 ++++++
 nmigen/hdl/xfrm.py   | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/nmigen/back/pysim.py b/nmigen/back/pysim.py
index 0bcb0b3..a1e43a5 100644
--- a/nmigen/back/pysim.py
+++ b/nmigen/back/pysim.py
@@ -297,6 +297,12 @@ class _StatementCompiler(StatementVisitor):
             lhs(state, normalize(rhs(state), shape))
         return run
 
+    def on_Assert(self, stmt):
+        raise NotImplementedError("Asserts not yet implemented for Simulator backend.") # :nocov:
+
+    def on_Assume(self, stmt):
+        pass # :nocov:
+
     def on_Switch(self, stmt):
         test  = self.rrhs_compiler(stmt.test)
         cases = []
diff --git a/nmigen/hdl/xfrm.py b/nmigen/hdl/xfrm.py
index e70f9d1..d7c142a 100644
--- a/nmigen/hdl/xfrm.py
+++ b/nmigen/hdl/xfrm.py
@@ -141,6 +141,14 @@ class StatementVisitor(metaclass=ABCMeta):
     def on_Assign(self, stmt):
         pass # :nocov:
 
+    @abstractmethod
+    def on_Assert(self, stmt):
+        pass # :nocov:
+
+    @abstractmethod
+    def on_Assume(self, stmt):
+        pass # :nocov:
+
     @abstractmethod
     def on_Switch(self, stmt):
         pass # :nocov:
@@ -155,6 +163,10 @@ class StatementVisitor(metaclass=ABCMeta):
     def on_statement(self, stmt):
         if type(stmt) is Assign:
             return self.on_Assign(stmt)
+        elif type(stmt) is Assert:
+            return self.on_Assert(stmt)
+        elif type(stmt) is Assume:
+            return self.on_Assume(stmt)
         elif isinstance(stmt, Switch):
             # Uses `isinstance()` and not `type() is` because nmigen.compat requires it.
             return self.on_Switch(stmt)
@@ -174,6 +186,12 @@ class StatementTransformer(StatementVisitor):
     def on_Assign(self, stmt):
         return Assign(self.on_value(stmt.lhs), self.on_value(stmt.rhs))
 
+    def on_Assert(self, stmt):
+        return Assert(self.on_value(stmt.test), _check=stmt._check, _en=stmt._en)
+
+    def on_Assume(self, stmt):
+        return Assume(self.on_value(stmt.test), _check=stmt._check, _en=stmt._en)
+
     def on_Switch(self, stmt):
         cases = OrderedDict((k, self.on_statement(s)) for k, s in stmt.cases.items())
         return Switch(self.on_value(stmt.test), cases)
@@ -294,6 +312,10 @@ class SwitchCleaner(StatementVisitor):
     def on_Assign(self, stmt):
         return stmt
 
+    on_Assert = on_Assign
+
+    on_Assume = on_Assign
+
     def on_Switch(self, stmt):
         cases = OrderedDict((k, self.on_statement(s)) for k, s in stmt.cases.items())
         if any(len(s) for s in cases.values()):
@@ -338,6 +360,10 @@ class LHSGroupAnalyzer(StatementVisitor):
     def on_Assign(self, stmt):
         self.unify(*stmt._lhs_signals())
 
+    on_Assert = on_Assign
+
+    on_Assume = on_Assign
+
     def on_Switch(self, stmt):
         for case_stmts in stmt.cases.values():
             self.on_statements(case_stmts)
@@ -362,6 +388,13 @@ class LHSGroupFilter(SwitchCleaner):
         if any_lhs_signal in self.signals:
             return stmt
 
+    def on_Assert(self, stmt):
+        any_lhs_signal = next(iter(stmt._lhs_signals()))
+        if any_lhs_signal in self.signals:
+            return stmt
+
+    on_Assume = on_Assert
+
 
 class _ControlInserter(FragmentTransformer):
     def __init__(self, controls):
-- 
2.30.2