From 5b46a74b7fde70f7ad5d54b5e313b2ac9567b410 Mon Sep 17 00:00:00 2001
From: youben11 <ayoub.benaissa@zama.ai>
Date: Tue, 15 Nov 2022 19:02:13 +0100
Subject: [PATCH] feat(CAPI): API to get the width of an eint/esint

---
 compiler/include/concretelang-c/Dialect/FHE.h |  6 ++++++
 compiler/lib/CAPI/Dialect/FHE/FHE.cpp         | 13 +++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/compiler/include/concretelang-c/Dialect/FHE.h b/compiler/include/concretelang-c/Dialect/FHE.h
index d7683b95a..e56e21a8c 100644
--- a/compiler/include/concretelang-c/Dialect/FHE.h
+++ b/compiler/include/concretelang-c/Dialect/FHE.h
@@ -36,6 +36,12 @@ fheEncryptedSignedIntegerTypeGetChecked(MlirContext context, unsigned width);
 /// If the type is an EncryptedSignedInteger
 MLIR_CAPI_EXPORTED bool fheTypeIsAnEncryptedSignedIntegerType(MlirType);
 
+/// \brief Get bitwidth of the encrypted integer type.
+///
+/// \return bitwidth of the encrypted integer or 0 if it's not an encrypted
+/// integer
+MLIR_CAPI_EXPORTED unsigned fheTypeIntegerWidthGet(MlirType);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/compiler/lib/CAPI/Dialect/FHE/FHE.cpp b/compiler/lib/CAPI/Dialect/FHE/FHE.cpp
index 0c0f97571..d06a45b76 100644
--- a/compiler/lib/CAPI/Dialect/FHE/FHE.cpp
+++ b/compiler/lib/CAPI/Dialect/FHE/FHE.cpp
@@ -61,3 +61,16 @@ MlirTypeOrError fheEncryptedSignedIntegerTypeGetChecked(MlirContext ctx,
                                                         unsigned width) {
   return IntegerTypeGetChecked<EncryptedSignedIntegerType>(ctx, width);
 }
+
+unsigned fheTypeIntegerWidthGet(MlirType integerType) {
+  mlir::Type type = unwrap(integerType);
+  auto eint = type.dyn_cast_or_null<EncryptedIntegerType>();
+  if (eint) {
+    return eint.getWidth();
+  }
+  auto esint = type.dyn_cast_or_null<EncryptedSignedIntegerType>();
+  if (esint) {
+    return esint.getWidth();
+  }
+  return 0;
+}