fix wrong API get_root_of_unity()

ingonyama-zk · Aug 20, 2024 · 98ad969 · 98ad969
1 parent 9759cef
commit 98ad969
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/icicle_v3/include/icicle/ntt.h b/icicle_v3/include/icicle/ntt.h
@@ -152,7 +152,7 @@ namespace icicle {
    * @return S Root of unity.
    */
   template <typename S>
-  S get_root_of_unity(uint64_t max_size);
+  eIcicleError get_root_of_unity(uint64_t max_size, S* rou);
 
   /**
    * @brief Gets the root of unity from the NTT domain for a given logarithmic size.

diff --git a/icicle_v3/src/ntt.cpp b/icicle_v3/src/ntt.cpp
@@ -68,16 +68,21 @@ namespace icicle {
   }
 
   /*************************** GET ROOT OF UNITY ***************************/
-  extern "C" scalar_t CONCAT_EXPAND(FIELD, get_root_of_unity)(uint64_t max_size)
+  extern "C" eIcicleError CONCAT_EXPAND(FIELD, get_root_of_unity)(uint64_t max_size, scalar_t* rou)
   {
     const auto log_max_size = static_cast<uint32_t>(std::ceil(std::log2(max_size)));
-    return scalar_t::omega(log_max_size);
+    if (scalar_t::get_omegas_count() < log_max_size) {
+      ICICLE_LOG_ERROR << "no root-of-unity of order " << log_max_size << " in field " << typeid(scalar_t).name();
+      return eIcicleError::INVALID_ARGUMENT;
+    }
+    *rou = scalar_t::omega(log_max_size);
+    return eIcicleError::SUCCESS;
   }
 
   template <>
-  scalar_t get_root_of_unity(uint64_t max_size)
+  eIcicleError get_root_of_unity(uint64_t max_size, scalar_t* rou)
   {
-    return CONCAT_EXPAND(FIELD, get_root_of_unity)(max_size);
+    return CONCAT_EXPAND(FIELD, get_root_of_unity)(max_size, rou);
   }
 
   /*************************** GET ROOT OF UNITY FROM DOMAIN ***************************/