From dba37cb71c124df5dd3168d4ea133715539b176b Mon Sep 17 00:00:00 2001
From: Arni Hod <arnon@starkware.co>
Date: Wed, 26 Jun 2024 14:39:07 +0300
Subject: [PATCH] feat: use starknet_api builtin for name trait

---
 crates/gateway/src/compilation.rs | 28 +++++++++++++++++++++++-----
 crates/gateway/src/utils.rs       | 30 ++++++++++++++++++++++++++++--
 crates/gateway/src/utils_test.rs  | 20 +++++++++++++++++++-
 3 files changed, 70 insertions(+), 8 deletions(-)

diff --git a/crates/gateway/src/compilation.rs b/crates/gateway/src/compilation.rs
index d52bbd91..8b9b44cd 100644
--- a/crates/gateway/src/compilation.rs
+++ b/crates/gateway/src/compilation.rs
@@ -8,12 +8,13 @@ use cairo_lang_starknet_classes::casm_contract_class::{
 };
 use starknet_api::core::CompiledClassHash;
 use starknet_api::rpc_transaction::RPCDeclareTransaction;
+use starknet_api::transaction::Builtin;
 use starknet_sierra_compile::compile::compile_sierra_to_casm;
 use starknet_sierra_compile::errors::CompilationUtilError;
 use starknet_sierra_compile::utils::into_contract_class_for_compilation;
 
 use crate::errors::{GatewayError, GatewayResult};
-use crate::utils::is_subsequence;
+use crate::utils::{is_subsequence, IntoEnumIteratorExt};
 
 #[cfg(test)]
 #[path = "compilation_test.rs"]
@@ -60,16 +61,33 @@ pub fn compile_contract_class(declare_tx: &RPCDeclareTransaction) -> GatewayResu
     Ok(class_info)
 }
 
+// TODO(Arni): Consider moving into Starknet-api.
+// List of supported builtins. These builtins can be used by contracts directly.
+// This is an explicit function so that it is explicitly decides which builtins are supported.
+// If new builtins are added, they should be added here.
+fn is_os_supported_builtin(builtin: &Builtin) -> bool {
+    match builtin {
+        Builtin::RangeCheck
+        | Builtin::Pedersen
+        | Builtin::Poseidon
+        | Builtin::EcOp
+        | Builtin::Ecdsa
+        | Builtin::Bitwise
+        | Builtin::SegmentArena => true,
+        Builtin::Keccak => false,
+    }
+}
+
 // TODO(Arni): Add to a config.
 // TODO(Arni): Use the Builtin enum from Starknet-api, and explicitly tag each builtin as supported
 // or unsupported so that the compiler would alert us on new builtins.
 fn supported_builtins() -> &'static Vec<String> {
     static SUPPORTED_BUILTINS: OnceLock<Vec<String>> = OnceLock::new();
     SUPPORTED_BUILTINS.get_or_init(|| {
-        // The OS expects this order for the builtins.
-        const SUPPORTED_BUILTIN_NAMES: [&str; 7] =
-            ["pedersen", "range_check", "ecdsa", "bitwise", "ec_op", "poseidon", "segment_arena"];
-        SUPPORTED_BUILTIN_NAMES.iter().map(|builtin| builtin.to_string()).collect::<Vec<String>>()
+        Builtin::iter()
+            .filter(is_os_supported_builtin)
+            .map(|builtin| builtin.name().to_string())
+            .collect::<Vec<String>>()
     })
 }
 
diff --git a/crates/gateway/src/utils.rs b/crates/gateway/src/utils.rs
index 6cebd08d..4c50e1db 100644
--- a/crates/gateway/src/utils.rs
+++ b/crates/gateway/src/utils.rs
@@ -10,8 +10,9 @@ use starknet_api::rpc_transaction::{
     RPCDeclareTransaction, RPCDeployAccountTransaction, RPCInvokeTransaction, RPCTransaction,
 };
 use starknet_api::transaction::{
-    DeclareTransaction, DeclareTransactionV3, DeployAccountTransaction, DeployAccountTransactionV3,
-    InvokeTransaction, InvokeTransactionV3, Tip, TransactionHash, TransactionHasher,
+    Builtin, DeclareTransaction, DeclareTransactionV3, DeployAccountTransaction,
+    DeployAccountTransactionV3, InvokeTransaction, InvokeTransactionV3, Tip, TransactionHash,
+    TransactionHasher,
 };
 use starknet_mempool_types::mempool_types::ThinTransaction;
 
@@ -176,3 +177,28 @@ pub fn is_subsequence<T: Eq>(subsequence: &[T], sequence: &[T]) -> bool {
 
     offset == subsequence.len()
 }
+
+// TODO(Arni): Remove the trait IntoEnumIteratorExt once EnumIter is implemented in starknet API.
+pub trait IntoEnumIteratorExt {
+    /// Returns an iterator over all the builtins in the order the Starknet OS expects.
+    fn iter() -> impl Iterator<Item = Builtin>;
+}
+
+// TODO(Arni): When the trait IntoEnumIteratorExt is removed, Make sure the order is maintained or
+// address the different order.
+impl IntoEnumIteratorExt for Builtin {
+    fn iter() -> impl Iterator<Item = Builtin> {
+        // The OS expects this order for the builtins.
+        vec![
+            Builtin::Pedersen,
+            Builtin::RangeCheck,
+            Builtin::Ecdsa,
+            Builtin::Bitwise,
+            Builtin::EcOp,
+            Builtin::Poseidon,
+            Builtin::SegmentArena,
+            Builtin::Keccak,
+        ]
+        .into_iter()
+    }
+}
diff --git a/crates/gateway/src/utils_test.rs b/crates/gateway/src/utils_test.rs
index 7b9d2fdb..3e2e096c 100644
--- a/crates/gateway/src/utils_test.rs
+++ b/crates/gateway/src/utils_test.rs
@@ -1,7 +1,8 @@
 use pretty_assertions::assert_eq;
 use rstest::rstest;
+use starknet_api::transaction::Builtin;
 
-use crate::utils::is_subsequence;
+use crate::utils::{is_subsequence, IntoEnumIteratorExt};
 
 #[rstest]
 #[case::empty(
@@ -76,3 +77,20 @@ fn test_is_subsequence(
 ) {
     assert_eq!(is_subsequence(subsequence, sequence), expected_result);
 }
+
+#[test]
+fn test_os_order_iter() {
+    let expected_order = vec![
+        Builtin::Pedersen,
+        Builtin::RangeCheck,
+        Builtin::Ecdsa,
+        Builtin::Bitwise,
+        Builtin::EcOp,
+        Builtin::Poseidon,
+        Builtin::SegmentArena,
+        Builtin::Keccak,
+    ];
+
+    let actual_order = Builtin::iter().collect::<Vec<_>>();
+    assert_eq!(actual_order, expected_order);
+}