From 2dea42e9039cdf47ca4927f62c69c6ae7ac2e399 Mon Sep 17 00:00:00 2001
From: Udi Finkelstein <github@udifink.com>
Date: Tue, 26 Sep 2017 09:11:25 +0300
Subject: [PATCH] Added $bits() for memories as well.

---
 frontends/ast/simplify.cc   | 28 ++++++++++++++++++++++++++--
 tests/simple/functions01.sv | 11 +++++------
 2 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/frontends/ast/simplify.cc b/frontends/ast/simplify.cc
index 608df0cb9..0cde34dc5 100644
--- a/frontends/ast/simplify.cc
+++ b/frontends/ast/simplify.cc
@@ -1870,19 +1870,43 @@ skip_dynamic_range_lvalue_expansion:;
 				goto apply_newNode;
 			}
 
-			if (str == "\\$size")
+			if (str == "\\$size" || str == "\\$bits")
 			{
 				if (children.size() != 1)
 					log_error("System function %s got %d arguments, expected 1 at %s:%d.\n",
 							RTLIL::unescape_id(str).c_str(), int(children.size()), filename.c_str(), linenum);
 
 				AstNode *buf = children[0]->clone();
+				int mem_depth = 1;
+				AstNode *id_ast = NULL;
+				
 				// Is this needed?
 				//while (buf->simplify(true, false, false, stage, width_hint, sign_hint, false)) { }
 				buf->detectSignWidth(width_hint, sign_hint);
+				if (str == "\\$bits") {
+					if (buf->type == AST_IDENTIFIER) {
+						id_ast = buf->id2ast;
+						if (id_ast == NULL && current_scope.count(buf->str))
+							id_ast = current_scope.at(buf->str);
+						if (!id_ast)
+							log_error("Failed to resolve identifier %s for width detection at %s:%d!\n", buf->str.c_str(), filename.c_str(), linenum);
+						if (id_ast->type == AST_MEMORY) {
+							AstNode *mem_range = id_ast->children[1];
+							if (mem_range->type == AST_RANGE) {
+								if (!mem_range->range_valid)
+									log_error("Failed to detect width of memory access `%s' at %s:%d!\n", mem_range->str.c_str(), filename.c_str(), linenum);
+								mem_depth = mem_range->range_left - mem_range->range_right + 1;
+							} else if (mem_range->type == AST_MULTIRANGE) {
+								for (auto n : mem_range->children) 
+									mem_depth *= (n->range_left - n->range_right + 1);
+							} else
+								log_error("Unknown memory depth AST type in `%s' at %s:%d!\n", mem_range->str.c_str(), filename.c_str(), linenum);
+						}
+					}
+				}
 				delete buf;
 
-				newNode = mkconst_int(width_hint, false);
+				newNode = mkconst_int(width_hint * mem_depth, false);
 				goto apply_newNode;
 			}
 
diff --git a/tests/simple/functions01.sv b/tests/simple/functions01.sv
index cd87a49b7..d6cd53e07 100644
--- a/tests/simple/functions01.sv
+++ b/tests/simple/functions01.sv
@@ -1,4 +1,5 @@
 module functions01;
+
 wire [3:0]x;
 wire [$size(x)-1:0]x_size;
 wire [$size({x, x})-1:0]xx_size;
@@ -7,11 +8,9 @@ wire [$size(y)-1:0]y_size;
 wire [3:0]z[0:5][0:7];
 wire [$size(z)-1:0]z_size;
 
-//
-// The following are not supported yet:
-//
-
-//wire [$bits(x)-1:0]x_bits;
-//wire [$bits({x, x})-1:0]xx_bits;
+wire [$bits(x)-1:0]x_bits;
+wire [$bits({x, x})-1:0]xx_bits;
+wire [$bits(y)-1:0]y_bits;
+wire [$bits(z)-1:0]z_bits;
 
 endmodule
-- 
2.30.2