From ed0ef9a42aead44b04c517a53928ab14a77d9160 Mon Sep 17 00:00:00 2001
From: Mikhail Aksenov <maksenov@nil.foundation>
Date: Mon, 19 Feb 2024 13:19:34 +0200
Subject: [PATCH] Enable constant folding for fields

---
 llvm/lib/IR/ConstantFold.cpp | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/llvm/lib/IR/ConstantFold.cpp b/llvm/lib/IR/ConstantFold.cpp
index cd93474324fa..e21204b494da 100644
--- a/llvm/lib/IR/ConstantFold.cpp
+++ b/llvm/lib/IR/ConstantFold.cpp
@@ -30,6 +30,7 @@
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/ZK/FieldArithmetics.h"
 using namespace llvm;
 using namespace llvm::PatternMatch;
 
@@ -1221,6 +1222,32 @@ Constant *llvm::ConstantFoldBinaryInstruction(unsigned Opcode, Constant *C1,
     }
   }
 
+  if (auto *CF1 = dyn_cast<ConstantField>(C1)) {
+    if (auto *CF2 = dyn_cast<ConstantField>(C2)) {
+      FieldElem Elem1 = CF1->getValue();
+      FieldElem Elem2 = CF2->getValue();
+      assert(Elem1.getKind() == Elem2.getKind());
+      FieldOperation Op;
+      switch (Opcode) {
+      case Instruction::Add:
+        Op = F_Add;
+        break;
+      case Instruction::Sub:
+        Op = F_Sub;
+        break;
+      case Instruction::Mul:
+        Op = F_Mul;
+        break;
+      case Instruction::UDiv:
+      case Instruction::SDiv:
+        Op = F_Div;
+        break;
+      }
+      FieldElem Res = FieldBinOp(Op, Elem1, Elem2);
+      return ConstantField::get(llvm::cast<GaloisFieldType>(C1->getType()), Res);
+    }
+  }
+
   if (ConstantExpr *CE1 = dyn_cast<ConstantExpr>(C1)) {
     // There are many possible foldings we could do here.  We should probably
     // at least fold add of a pointer with an integer into the appropriate