From 52b8b884f7125f8dcb5db957fbcddefdfb28c1f2 Mon Sep 17 00:00:00 2001
From: Dag Lem <dag@nimrod.no>
Date: Tue, 8 Aug 2023 13:53:05 +0200
Subject: [PATCH] Corrections and optimizations for nordshift

* Corrects corner cases
* Handles variable slices on the form dst[i*stride +: width] = src
---
 frontends/ast/simplify.cc | 43 +++++++++++++++++++++++++++++----------
 1 file changed, 32 insertions(+), 11 deletions(-)

diff --git a/frontends/ast/simplify.cc b/frontends/ast/simplify.cc
index 2bd1bba18fe..8f53f9ffa59 100644
--- a/frontends/ast/simplify.cc
+++ b/frontends/ast/simplify.cc
@@ -2389,18 +2389,29 @@ bool AstNode::simplify(bool const_fold, bool in_lvalue, int stage, int width_hin
 			int stride = 1;
 			int div_stride = 1;
 
-			// Optimization: Remove multiplication of index by stride for two-dimensional packed arrays and
-			// variable slices on the form dst = src[i*w +: w].
-			if (!source_offset && shift_expr->type == AST_MUL &&
-			    ((shift_expr->children[0]->type == AST_CONSTANT && (int)shift_expr->children[0]->integer == result_width) ||
-			     (shift_expr->children[1]->type == AST_CONSTANT && (int)shift_expr->children[1]->integer == result_width)))
+			// Extract (index)*(width) from non_opt_range pattern (@selfsz@((index)*(width)))+(0)).
+			AstNode *lsb_expr =
+				shift_expr->type == AST_ADD && shift_expr->children[0]->type == AST_SELFSZ &&
+				shift_expr->children[1]->type == AST_CONSTANT && shift_expr->children[1]->integer == 0 ?
+				shift_expr->children[0]->children[0] :
+				shift_expr;
+
+			// Optimization: Extract stride from indexing of two-dimensional packed arrays and
+			// variable slices on the form dst = src[i*stride +: width].
+			if (!source_offset && lsb_expr->type == AST_MUL &&
+			    (lsb_expr->children[0]->type == AST_CONSTANT || lsb_expr->children[1]->type == AST_CONSTANT))
 			{
-				int var_i = shift_expr->children[0]->type == AST_CONSTANT;
-				AstNode *tmp = shift_expr->children[var_i]->clone();
-				stride = result_width;
+				int const_i = lsb_expr->children[1]->type == AST_CONSTANT;
+				stride = (int)lsb_expr->children[const_i]->integer;
+
+				// Remove multiplication of index by stride.
+				// FIXME: Counterintuitively, this can yield higher resource usage. Disable for now.
+#if 0
 				div_stride = stride;
+				AstNode *tmp = lsb_expr->children[1 - const_i]->clone();
 				delete shift_expr;
 				shift_expr = tmp;
+#endif
 			}
 			else if (member_node)
 			{
@@ -2419,11 +2430,21 @@ bool AstNode::simplify(bool const_fold, bool in_lvalue, int stage, int width_hin
 				}
 			}
 
+			// Limit case conditions to possible index range.
+			int case_width_hint = -1;
+			bool case_sign_hint = true;
+			shift_expr->detectSignWidth(case_width_hint, case_sign_hint);
+			int max_bits = min(case_width_hint, 31 + case_sign_hint);
+			int max_offset = (1u << (max_bits - case_sign_hint)) - 1;
+			int min_offset = case_sign_hint ? -(1u << (max_bits - 1)) : 0;
+
 			assign = new AstNode(AST_CASE, shift_expr);
-			for (int i = 0; i < source_width; i += stride) {
+			for (int i = (1 - result_width)/stride*stride; i < source_width; i += stride) {
 				int start_bit = source_offset + i;
-				int end_bit = std::min(start_bit+result_width,source_width) - 1;
-				AstNode *cond = new AstNode(AST_COND, node_int(start_bit/div_stride));
+				int end_bit = std::min(start_bit + result_width, source_offset + source_width) - 1;
+				if (start_bit/div_stride < min_offset || start_bit/div_stride > max_offset)
+					continue;
+				AstNode *cond = new AstNode(AST_COND, mkconst_int(start_bit/div_stride, case_sign_hint, case_width_hint));
 				AstNode *rvalue = clone();
 				rvalue->delete_children();
 				if (member_node)