Added support for FOR loops in function calls in parameters

3 files changed, 43 insertions, 0 deletions

diff --git a/frontends/ast/ast.cc b/frontends/ast/ast.cc
index 7e199cd5..56e9393b 100644
--- a/frontends/ast/ast.cc
+++ b/frontends/ast/ast.cc
@@ -136,6 +136,7 @@ std::string AST::type2str(AstNodeType type)
 	X(AST_COND)
 	X(AST_DEFAULT)
 	X(AST_FOR)
+	X(AST_WHILE)
 	X(AST_GENVAR)
 	X(AST_GENFOR)
 	X(AST_GENIF)
diff --git a/frontends/ast/ast.h b/frontends/ast/ast.h
index ec4f4f62..f42bc35f 100644
--- a/frontends/ast/ast.h
+++ b/frontends/ast/ast.h
@@ -116,6 +116,7 @@ namespace AST
 		AST_COND,
 		AST_DEFAULT,
 		AST_FOR,
+		AST_WHILE,
 
 		AST_GENVAR,
 		AST_GENFOR,
diff --git a/frontends/ast/simplify.cc b/frontends/ast/simplify.cc
index 85dc1d3c..236843d5 100644
--- a/frontends/ast/simplify.cc
+++ b/frontends/ast/simplify.cc
@@ -1891,10 +1891,51 @@ AstNode *AstNode::eval_const_function(AstNode *fcall)
 
 			variables[stmt->children.at(0)->str].val = stmt->children.at(1)->bitsAsConst(variables[stmt->children.at(0)->str].val.bits.size());
 
+			delete block->children.front();
 			block->children.erase(block->children.begin());
 			continue;
 		}
 
+		if (stmt->type == AST_FOR)
+		{
+			block->children.insert(block->children.begin(), stmt->children.at(0));
+			stmt->children.at(3)->children.push_back(stmt->children.at(2));
+			stmt->children.erase(stmt->children.begin() + 2);
+			stmt->children.erase(stmt->children.begin());
+			stmt->type = AST_WHILE;
+			continue;
+		}
+
+		if (stmt->type == AST_WHILE)
+		{
+			AstNode *cond = stmt->children.at(0)->clone();
+			cond->replace_variables(variables, fcall);
+			while (cond->simplify(true, false, false, 1, -1, false, true)) { }
+
+			if (cond->type != AST_CONSTANT)
+				log_error("Non-constant expression in constant function at %s:%d (called from %s:%d).\n",
+						stmt->filename.c_str(), stmt->linenum, fcall->filename.c_str(), fcall->linenum);
+
+			if (cond->asBool()) {
+				block->children.insert(block->children.begin(), stmt->children.at(1)->clone());
+			} else {
+				delete block->children.front();
+				block->children.erase(block->children.begin());
+			}
+
+			delete cond;
+			continue;
+		}
+
+		if (stmt->type == AST_BLOCK)
+		{
+			block->children.erase(block->children.begin());
+			block->children.insert(block->children.begin(), stmt->children.begin(), stmt->children.end());
+			stmt->children.clear();
+			delete stmt;
+			continue;
+		}
+
 		log_error("Unsupported language construct in constant function at %s:%d (called from %s:%d).\n",
 				stmt->filename.c_str(), stmt->linenum, fcall->filename.c_str(), fcall->linenum);
 		log_abort();