From 8c5f5c7db09a565e8b1c897cbd5d5f2160e6c8bd Mon Sep 17 00:00:00 2001
From: Daniel Lupu <dflupu@gmail.com>
Date: Mon, 2 May 2022 00:10:37 +0300
Subject: [PATCH] add rules for mod(mul(X, Y), A) & mod(add(X, Y), A)

---
 libevmasm/RuleList.h                          | 20 ++++++++++++++++++-
 .../combine_add_and_mod_constant.yul          | 13 ++++++++++++
 .../combine_mul_and_mod_constant.yul          | 13 ++++++++++++
 3 files changed, 45 insertions(+), 1 deletion(-)
 create mode 100644 test/libyul/yulOptimizerTests/expressionSimplifier/combine_add_and_mod_constant.yul
 create mode 100644 test/libyul/yulOptimizerTests/expressionSimplifier/combine_mul_and_mod_constant.yul

diff --git a/libevmasm/RuleList.h b/libevmasm/RuleList.h
index ec0970e65..1ed101455 100644
--- a/libevmasm/RuleList.h
+++ b/libevmasm/RuleList.h
@@ -279,7 +279,7 @@ std::vector<SimplificationRule<Pattern>> simplificationRuleListPart5(
 	Pattern B,
 	Pattern,
 	Pattern X,
-	Pattern
+	Pattern Y
 )
 {
 	using Word = typename Pattern::Word;
@@ -287,6 +287,24 @@ std::vector<SimplificationRule<Pattern>> simplificationRuleListPart5(
 
 	std::vector<SimplificationRule<Pattern>> rules;
 
+	// Replace MOD(MUL(X, Y), A) with MULMOD(X, Y, A) iff A=2**N
+	rules.push_back({
+		Builtins::MOD(Builtins::MUL(X, Y), A),
+		[=]() -> Pattern { return Builtins::MULMOD(X, Y, A); },
+		[=] {
+			return A.d() > 0 && ((A.d() & (A.d() - 1)) == 0);
+		}
+	});
+
+	// Replace MOD(ADD(X, Y), A) with ADDMOD(X, Y, A) iff A=2**N
+	rules.push_back({
+		Builtins::MOD(Builtins::ADD(X, Y), A),
+		[=]() -> Pattern { return Builtins::ADDMOD(X, Y, A); },
+		[=] {
+			return A.d() > 0 && ((A.d() & (A.d() - 1)) == 0);
+		}
+	});
+
 	// Replace MOD X, <power-of-two> with AND X, <power-of-two> - 1
 	for (size_t i = 0; i < Pattern::WordSize; ++i)
 	{
diff --git a/test/libyul/yulOptimizerTests/expressionSimplifier/combine_add_and_mod_constant.yul b/test/libyul/yulOptimizerTests/expressionSimplifier/combine_add_and_mod_constant.yul
new file mode 100644
index 000000000..f462486ea
--- /dev/null
+++ b/test/libyul/yulOptimizerTests/expressionSimplifier/combine_add_and_mod_constant.yul
@@ -0,0 +1,13 @@
+{
+    mstore(0, mod(add(mload(0), mload(1)), 32))
+}
+// ----
+// step: expressionSimplifier
+//
+// {
+//     {
+//         let _3 := mload(1)
+//         let _4 := 0
+//         mstore(_4, addmod(mload(_4), _3, 32))
+//     }
+// }
diff --git a/test/libyul/yulOptimizerTests/expressionSimplifier/combine_mul_and_mod_constant.yul b/test/libyul/yulOptimizerTests/expressionSimplifier/combine_mul_and_mod_constant.yul
new file mode 100644
index 000000000..93edb5319
--- /dev/null
+++ b/test/libyul/yulOptimizerTests/expressionSimplifier/combine_mul_and_mod_constant.yul
@@ -0,0 +1,13 @@
+{
+    mstore(0, mod(mul(mload(0), mload(1)), 32))
+}
+// ----
+// step: expressionSimplifier
+//
+// {
+//     {
+//         let _3 := mload(1)
+//         let _4 := 0
+//         mstore(_4, mulmod(mload(_4), _3, 32))
+//     }
+// }