1 files changed, 88 insertions, 1 deletions

diff --git a/frontends/ast/simplify.cc b/frontends/ast/simplify.cc
index 6b4b9e045..5c4dd290f 100644
--- a/frontends/ast/simplify.cc
+++ b/frontends/ast/simplify.cc
@@ -571,7 +571,7 @@ bool AstNode::simplify(bool const_fold, bool at_zero, bool in_lvalue, int stage,
 	static bool deep_recursion_warning = false;
 
 	if (recursion_counter++ == 1000 && deep_recursion_warning) {
-		log_warning("Deep recursion in AST simplifier.\nDoes this design contain insanely long expressions?\n");
+		log_warning("Deep recursion in AST simplifier.\nDoes this design contain overly long or deeply nested expressions, or excessive recursion?\n");
 		deep_recursion_warning = false;
 	}
 
@@ -3089,6 +3089,93 @@ skip_dynamic_range_lvalue_expansion:;
 				goto apply_newNode;
 			}
 
+			if (str == "\\$countbits") {
+				if (children.size() < 2)
+					log_file_error(filename, location.first_line, "System function %s got %d arguments, expected at least 2.\n",
+							RTLIL::unescape_id(str).c_str(), int(children.size()));
+
+				std::vector<RTLIL::State> control_bits;
+
+				// Determine which bits to count
+				for (size_t i = 1; i < children.size(); i++) {
+					AstNode *node = children[i];
+					while (node->simplify(true, false, false, stage, -1, false, false)) { }
+					if (node->type != AST_CONSTANT)
+						log_file_error(filename, location.first_line, "Failed to evaluate system function `%s' with non-constant control bit argument.\n", str.c_str());
+					if (node->bits.size() != 1)
+						log_file_error(filename, location.first_line, "Failed to evaluate system function `%s' with control bit width != 1.\n", str.c_str());
+					control_bits.push_back(node->bits[0]);
+				}
+
+				// Detect width of exp (first argument of $countbits)
+				int  exp_width = -1;
+				bool exp_sign  = false;
+				AstNode *exp = children[0];
+				exp->detectSignWidth(exp_width, exp_sign, NULL);
+
+				newNode = mkconst_int(0, false);
+
+				for (int i = 0; i < exp_width; i++) {
+					// Generate nodes for:  exp << i >> ($size(exp) - 1)
+					//                          ^^   ^^
+					AstNode *lsh_node = new AstNode(AST_SHIFT_LEFT, exp->clone(), mkconst_int(i, false));
+					AstNode *rsh_node = new AstNode(AST_SHIFT_RIGHT, lsh_node, mkconst_int(exp_width - 1, false));
+
+					AstNode *or_node = nullptr;
+
+					for (RTLIL::State control_bit : control_bits) {
+						// Generate node for:  (exp << i >> ($size(exp) - 1)) === control_bit
+						//                                                    ^^^
+						AstNode *eq_node = new AstNode(AST_EQX, rsh_node->clone(), mkconst_bits({control_bit}, false));
+
+						// Or the result for each checked bit value
+						if (or_node)
+							or_node = new AstNode(AST_LOGIC_OR, or_node, eq_node);
+						else
+							or_node = eq_node;
+					}
+
+					// We should have at least one element in control_bits,
+					// because we checked for the number of arguments above
+					log_assert(or_node != nullptr);
+
+					delete rsh_node;
+
+					// Generate node for adding with result of previous bit
+					newNode = new AstNode(AST_ADD, newNode, or_node);
+				}
+
+				goto apply_newNode;
+			}
+
+			if (str == "\\$countones" || str == "\\$isunknown" || str == "\\$onehot" || str == "\\$onehot0") {
+				if (children.size() != 1)
+					log_file_error(filename, location.first_line, "System function %s got %d arguments, expected 1.\n",
+							RTLIL::unescape_id(str).c_str(), int(children.size()));
+
+				AstNode *countbits = clone();
+				countbits->str = "\\$countbits";
+
+				if (str == "\\$countones") {
+					countbits->children.push_back(mkconst_bits({RTLIL::State::S1}, false));
+					newNode = countbits;
+				} else if (str == "\\$isunknown") {
+					countbits->children.push_back(mkconst_bits({RTLIL::Sx}, false));
+					countbits->children.push_back(mkconst_bits({RTLIL::Sz}, false));
+					newNode = new AstNode(AST_GT, countbits, mkconst_int(0, false));
+				} else if (str == "\\$onehot") {
+					countbits->children.push_back(mkconst_bits({RTLIL::State::S1}, false));
+					newNode = new AstNode(AST_EQ, countbits, mkconst_int(1, false));
+				} else if (str == "\\$onehot0") {
+					countbits->children.push_back(mkconst_bits({RTLIL::State::S1}, false));
+					newNode = new AstNode(AST_LE, countbits, mkconst_int(1, false));
+				} else {
+					log_abort();
+				}
+
+				goto apply_newNode;
+			}
+
 			if (current_scope.count(str) != 0 && current_scope[str]->type == AST_DPI_FUNCTION)
 			{
 				AstNode *dpi_decl = current_scope[str];