From 9e24eccc9dd755c65b416acba5834ce2c2cae6ed Mon Sep 17 00:00:00 2001
From: Andi Drebes <andi.drebes@zama.ai>
Date: Tue, 20 Sep 2022 15:56:26 +0200
Subject: [PATCH] enhance(compiler): Add operators == and != for Integer /
 Tensor Lambda Arguments

---
 .../concretelang/Support/LambdaArgument.h     | 30 +++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/compiler/include/concretelang/Support/LambdaArgument.h b/compiler/include/concretelang/Support/LambdaArgument.h
index 67358b339..62615e333 100644
--- a/compiler/include/concretelang/Support/LambdaArgument.h
+++ b/compiler/include/concretelang/Support/LambdaArgument.h
@@ -65,6 +65,16 @@ public:
   unsigned int getPrecision() const { return this->precision; }
   BackingIntType getValue() const { return this->value; }
 
+  template <typename OtherBackingIntType>
+  bool operator==(const IntLambdaArgument<OtherBackingIntType> &other) const {
+    return getValue() == other.getValue();
+  }
+
+  template <typename OtherBackingIntType>
+  bool operator!=(const IntLambdaArgument<OtherBackingIntType> &other) const {
+    return !(*this == other);
+  }
+
   static char ID;
 
 protected:
@@ -177,6 +187,26 @@ public:
     return this->value.data();
   }
 
+  template <typename OtherScalarArgumentT>
+  bool
+  operator==(const TensorLambdaArgument<OtherScalarArgumentT> &other) const {
+    if (getDimensions() != other.getDimensions())
+      return false;
+
+    for (auto pair : llvm::zip(value, other.value)) {
+      if (std::get<0>(pair) != std::get<1>(pair))
+        return false;
+    }
+
+    return true;
+  }
+
+  template <typename OtherScalarArgumentT>
+  bool
+  operator!=(const TensorLambdaArgument<OtherScalarArgumentT> &other) const {
+    return !(*this == other);
+  }
+
   static char ID;
 
 protected: