Added support for bit/part select to mem2reg rewriter

2 files changed, 30 insertions, 0 deletions

diff --git a/frontends/ast/simplify.cc b/frontends/ast/simplify.cc
index 320c80d72..eee5a7b39 100644
--- a/frontends/ast/simplify.cc
+++ b/frontends/ast/simplify.cc
@@ -1974,6 +1974,8 @@ void AstNode::mem2reg_as_needed_pass2(std::set<AstNode*> &mem2reg_set, AstNode *
 				continue;
 			AstNode *cond_node = new AstNode(AST_COND, AstNode::mkconst_int(i, false, addr_bits), new AstNode(AST_BLOCK));
 			AstNode *assign_reg = new AstNode(type, new AstNode(AST_IDENTIFIER), new AstNode(AST_IDENTIFIER));
+			if (children[0]->children.size() == 2)
+				assign_reg->children[0]->children.push_back(children[0]->children[1]->clone());
 			assign_reg->children[0]->str = stringf("%s[%d]", children[0]->str.c_str(), i);
 			assign_reg->children[1]->str = id_data;
 			cond_node->children[1]->children.push_back(assign_reg);
@@ -1990,6 +1992,10 @@ void AstNode::mem2reg_as_needed_pass2(std::set<AstNode*> &mem2reg_set, AstNode *
 
 	if (type == AST_IDENTIFIER && id2ast && mem2reg_set.count(id2ast) > 0)
 	{
+		AstNode *bit_part_sel = NULL;
+		if (children.size() == 2)
+			bit_part_sel = children[1]->clone();
+
 		if (children[0]->children[0]->type == AST_CONSTANT)
 		{
 			int id = children[0]->children[0]->integer;
@@ -2073,6 +2079,9 @@ void AstNode::mem2reg_as_needed_pass2(std::set<AstNode*> &mem2reg_set, AstNode *
 			id2ast = NULL;
 			str = id_data;
 		}
+
+		if (bit_part_sel)
+			children.push_back(bit_part_sel);
 	}
 
 	assert(id2ast == NULL || mem2reg_set.count(id2ast) == 0);
diff --git a/tests/simple/memory.v b/tests/simple/memory.v
index aae3feace..21271b5e2 100644
--- a/tests/simple/memory.v
+++ b/tests/simple/memory.v
@@ -134,3 +134,24 @@ always @(posedge clk) begin
 end
 
 endmodule
+
+// ----------------------------------------------------------
+
+module test06(input clk, input rst, input [2:0] idx, input [7:0] din, output [7:0] dout);
+    (* gentb_constant=0 *) wire rst;
+    reg [7:0] test [0:7];
+    integer i;
+    always @(posedge clk or posedge rst) begin
+        if (rst) begin
+            for (i=0; i<8; i=i+1)
+                test[i] <= 0;
+        end else begin
+            test[0][2] <= din[1];
+            test[0][5] <= test[0][2];
+            test[idx][3] <= din[idx];
+            test[idx][6] <= test[idx][2];
+            test[idx][idx] <= !test[idx][idx];
+        end
+    end
+    assign dout = test[idx];
+endmodule