2 files changed, 11 insertions, 2 deletions
diff --git a/kernel/rtlil.cc b/kernel/rtlil.cc
index 790ba52a3..f732b56b0 100644
--- a/kernel/rtlil.cc
+++ b/kernel/rtlil.cc
@@ -1565,13 +1565,21 @@ void RTLIL::Module::remove(const pool<RTLIL::Wire*> &wires)
 
 void RTLIL::Module::remove(RTLIL::Cell *cell)
 {
+	auto it = cells_.find(cell->name);
+	log_assert(it != cells_.end());
+	remove(it);
+}
+
+dict<RTLIL::IdString, RTLIL::Cell*>::iterator RTLIL::Module::remove(dict<RTLIL::IdString, RTLIL::Cell*>::iterator it)
+{
+	RTLIL::Cell *cell = it->second;
 	while (!cell->connections_.empty())
 		cell->unsetPort(cell->connections_.begin()->first);
 
-	log_assert(cells_.count(cell->name) != 0);
 	log_assert(refcount_cells_ == 0);
-	cells_.erase(cell->name);
+	it = cells_.erase(it);
 	delete cell;
+	return it;
 }
 
 void RTLIL::Module::rename(RTLIL::Wire *wire, RTLIL::IdString new_name)
diff --git a/kernel/rtlil.h b/kernel/rtlil.h
index f4fcf5dcf..4a0f8b4f8 100644
--- a/kernel/rtlil.h
+++ b/kernel/rtlil.h
@@ -1040,6 +1040,7 @@ public:
 	// Removing wires is expensive. If you have to remove wires, remove them all at once.
 	void remove(const pool<RTLIL::Wire*> &wires);
 	void remove(RTLIL::Cell *cell);
+	dict<RTLIL::IdString, RTLIL::Cell*>::iterator remove(dict<RTLIL::IdString, RTLIL::Cell*>::iterator it);
 
 	void rename(RTLIL::Wire *wire, RTLIL::IdString new_name);
 	void rename(RTLIL::Cell *cell, RTLIL::IdString new_name);