From 7638ef80da2ce2dad7c70aae05836083129d370e Mon Sep 17 00:00:00 2001
From: colepoirier <colepoirier@gmail.com>
Date: Sun, 7 Jun 2020 14:09:17 -0700
Subject: [PATCH] Add TrapMainStage.trap() convenience function to set trap
 address and PC to begin from on return

---
 src/soc/fu/trap/main_stage.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/src/soc/fu/trap/main_stage.py b/src/soc/fu/trap/main_stage.py
index 5f4b95c9..650dc8a8 100644
--- a/src/soc/fu/trap/main_stage.py
+++ b/src/soc/fu/trap/main_stage.py
@@ -73,6 +73,17 @@ class TrapMainStage(PipeModBase):
         self.fields = DecodeFields(SignalBitRange, [self.i.ctx.op.insn])
         self.fields.create_specs()
 
+    def trap(self, m, addr, trap_addr):
+        comb  = m.d.comb
+        nia_o, srr0_o = self.o.nia, self.o.srr0
+
+        comb += nia_o.data.eq(trap_addr)
+        comb += nia_o.ok.eq(1)
+
+        comb += srr0_o.data.eq(addr) # addr to begin from on return
+        comb += srro_o.ok.eq(1)
+
+
     def ispec(self):
         return TrapInputData(self.pspec)
 
-- 
2.30.2