Improved mem2reg handling in ast simplifier

2 files changed, 35 insertions, 5 deletions

diff --git a/frontends/ast/ast.h b/frontends/ast/ast.h
index cdc56fba4..76314f88d 100644
--- a/frontends/ast/ast.h
+++ b/frontends/ast/ast.h
@@ -165,7 +165,7 @@ namespace AST
 		void expand_genblock(std::string index_var, std::string prefix, std::map<std::string, std::string> &name_map);
 		void replace_ids(std::map<std::string, std::string> &rules);
 		void mem2reg_as_needed_pass1(std::set<AstNode*> &mem2reg_set, std::set<AstNode*> &mem2reg_candidates, bool sync_proc, bool async_proc);
-		void mem2reg_as_needed_pass2(std::set<AstNode*> &mem2reg_set, AstNode *mod, AstNode *block);
+		void mem2reg_as_needed_pass2(std::set<AstNode*> &mem2reg_set, AstNode *mod, AstNode *block, AstNode *top_block);
 		void meminfo(int &mem_width, int &mem_size, int &addr_bits);
 
 		// create a human-readable text representation of the AST (for debugging)
diff --git a/frontends/ast/simplify.cc b/frontends/ast/simplify.cc
index 94423366e..2790a4a33 100644
--- a/frontends/ast/simplify.cc
+++ b/frontends/ast/simplify.cc
@@ -74,7 +74,7 @@ bool AstNode::simplify(bool const_fold, bool at_zero, bool in_lvalue, int stage)
 				}
 			}
 
-			mem2reg_as_needed_pass2(mem2reg_set, this, NULL);
+			mem2reg_as_needed_pass2(mem2reg_set, this, NULL, NULL);
 
 			for (size_t i = 0; i < children.size(); i++) {
 				if (mem2reg_set.count(children[i]) > 0) {
@@ -937,6 +937,9 @@ void AstNode::mem2reg_as_needed_pass1(std::set<AstNode*> &mem2reg_set, std::set<
 			}
 			mem2reg_candidates.insert(children[0]->id2ast);
 		}
+	
+	if (type == AST_MEMORY && attributes.count("\\mem2reg") > 0)
+		mem2reg_set.insert(this);
 
 	if (type == AST_ALWAYS) {
 		for (auto child : children) {
@@ -951,7 +954,7 @@ void AstNode::mem2reg_as_needed_pass1(std::set<AstNode*> &mem2reg_set, std::set<
 }
 
 // actually replace memories with registers
-void AstNode::mem2reg_as_needed_pass2(std::set<AstNode*> &mem2reg_set, AstNode *mod, AstNode *block)
+void AstNode::mem2reg_as_needed_pass2(std::set<AstNode*> &mem2reg_set, AstNode *mod, AstNode *block, AstNode *top_block)
 {
 	if (type == AST_BLOCK)
 		block = this;
@@ -976,6 +979,18 @@ void AstNode::mem2reg_as_needed_pass2(std::set<AstNode*> &mem2reg_set, AstNode *
 		wire_data->is_reg = true;
 		mod->children.push_back(wire_data);
 
+		assert(top_block != NULL);
+		std::vector<RTLIL::State> x_bits;
+		x_bits.push_back(RTLIL::State::Sx);
+
+		AstNode *assign_addr_x = new AstNode(AST_ASSIGN_EQ, new AstNode(AST_IDENTIFIER), AstNode::mkconst_bits(x_bits, false));
+		assign_addr_x->children[0]->str = id_addr;
+		top_block->children.insert(top_block->children.begin(), assign_addr_x);
+
+		AstNode *assign_data_x = new AstNode(AST_ASSIGN_EQ, new AstNode(AST_IDENTIFIER), AstNode::mkconst_bits(x_bits, false));
+		assign_data_x->children[0]->str = id_data;
+		top_block->children.insert(top_block->children.begin(), assign_data_x);
+
 		assert(block != NULL);
 		size_t assign_idx = 0;
 		while (assign_idx < block->children.size() && block->children[assign_idx] != this)
@@ -1048,6 +1063,17 @@ void AstNode::mem2reg_as_needed_pass2(std::set<AstNode*> &mem2reg_set, AstNode *
 		cond_node->children[1]->children.push_back(assign_reg);
 		case_node->children.push_back(cond_node);
 
+		if (top_block)
+		{
+			AstNode *assign_addr_x = new AstNode(AST_ASSIGN_EQ, new AstNode(AST_IDENTIFIER), AstNode::mkconst_bits(x_bits, false));
+			assign_addr_x->children[0]->str = id_addr;
+			top_block->children.insert(top_block->children.begin(), assign_addr_x);
+
+			AstNode *assign_data_x = new AstNode(AST_ASSIGN_EQ, new AstNode(AST_IDENTIFIER), AstNode::mkconst_bits(x_bits, false));
+			assign_data_x->children[0]->str = id_data;
+			top_block->children.insert(top_block->children.begin(), assign_data_x);
+		}
+
 		if (block)
 		{
 			size_t assign_idx = 0;
@@ -1075,8 +1101,12 @@ void AstNode::mem2reg_as_needed_pass2(std::set<AstNode*> &mem2reg_set, AstNode *
 
 	assert(id2ast == NULL || mem2reg_set.count(id2ast) == 0);
 
-	for (size_t i = 0; i < children.size(); i++)
-		children[i]->mem2reg_as_needed_pass2(mem2reg_set, mod, block);
+	auto children_list = children;
+	for (size_t i = 0; i < children_list.size(); i++) {
+		if (type == AST_ALWAYS && children_list[i]->type == AST_BLOCK)
+			top_block = children_list[i];
+		children_list[i]->mem2reg_as_needed_pass2(mem2reg_set, mod, block, top_block);
+	}
 }
 
 // calulate memory dimensions