From cac3156ecb2ca9a16a45abaa9e5c443073b0a475 Mon Sep 17 00:00:00 2001
From: Federica <fedemoletta@hotmail.com>
Date: Thu, 18 Apr 2024 11:14:58 -0300
Subject: [PATCH] Complete API

---
 vm/src/tests/mod.rs          | 13 +------------
 vm/src/types/builtin_name.rs | 24 +++++++++++++++++++++---
 2 files changed, 22 insertions(+), 15 deletions(-)

diff --git a/vm/src/tests/mod.rs b/vm/src/tests/mod.rs
index b33725a9cc..c4e3ea49fc 100644
--- a/vm/src/tests/mod.rs
+++ b/vm/src/tests/mod.rs
@@ -310,17 +310,6 @@ fn get_casm_contract_builtins(
         .unwrap()
         .builtins
         .iter()
-        .map(|s| match s.as_str() {
-            "output" => BuiltinName::output,
-            "range_check" => BuiltinName::range_check,
-            "pedersen" => BuiltinName::pedersen,
-            "ecdsa" => BuiltinName::ecdsa,
-            "keccak" => BuiltinName::keccak,
-            "bitwise" => BuiltinName::bitwise,
-            "ec_op" => BuiltinName::ec_op,
-            "poseidon" => BuiltinName::poseidon,
-            "segment_arena" => BuiltinName::segment_arena,
-            _ => panic!("Invalid builtin {}", s),
-        })
+        .map(|s| BuiltinName::from_str(s).expect("Invalid builtin name"))
         .collect()
 }
diff --git a/vm/src/types/builtin_name.rs b/vm/src/types/builtin_name.rs
index 2d3e0ad9e4..1de03d1f1b 100644
--- a/vm/src/types/builtin_name.rs
+++ b/vm/src/types/builtin_name.rs
@@ -85,8 +85,8 @@ impl BuiltinName {
         }
     }
 
-    pub(crate) fn from_suffixed_string(suffixed_str: &String) -> Option<Self> {
-        match suffixed_str.as_str() {
+    pub fn from_str_with_suffix(suffixed_str: &str) -> Option<Self> {
+        match suffixed_str {
             OUTPUT_BUILTIN_NAME_WITH_SUFFIX => Some(BuiltinName::output),
             RANGE_CHECK_BUILTIN_NAME_WITH_SUFFIX => Some(BuiltinName::range_check),
             HASH_BUILTIN_NAME_WITH_SUFFIX => Some(BuiltinName::pedersen),
@@ -102,6 +102,24 @@ impl BuiltinName {
             _ => None,
         }
     }
+
+    pub fn from_str(str: &str) -> Option<Self> {
+        match str {
+            OUTPUT_BUILTIN_NAME => Some(BuiltinName::output),
+            RANGE_CHECK_BUILTIN_NAME => Some(BuiltinName::range_check),
+            HASH_BUILTIN_NAME => Some(BuiltinName::pedersen),
+            SIGNATURE_BUILTIN_NAME => Some(BuiltinName::ecdsa),
+            KECCAK_BUILTIN_NAME => Some(BuiltinName::keccak),
+            BITWISE_BUILTIN_NAME => Some(BuiltinName::bitwise),
+            EC_OP_BUILTIN_NAME => Some(BuiltinName::ec_op),
+            POSEIDON_BUILTIN_NAME => Some(BuiltinName::poseidon),
+            SEGMENT_ARENA_BUILTIN_NAME => Some(BuiltinName::segment_arena),
+            RANGE_CHECK_96_BUILTIN_NAME => Some(BuiltinName::range_check96),
+            ADD_MOD_BUILTIN_NAME => Some(BuiltinName::add_mod),
+            MUL_MOD_BUILTIN_NAME => Some(BuiltinName::mul_mod),
+            _ => None,
+        }
+    }
 }
 
 impl core::fmt::Display for BuiltinName {
@@ -137,7 +155,7 @@ pub(crate) mod serde_generic_map_impl {
         let map = HashMap::<String, V>::deserialize(d)?;
         // Then match keys to BuiltinName and handle invalid names
         map.into_iter()
-            .map(|(k, v)| BuiltinName::from_suffixed_string(&k).map(|k| (k, v)))
+            .map(|(k, v)| BuiltinName::from_str_with_suffix(&k).map(|k| (k, v)))
             .collect::<Option<HashMap<_, _>>>()
             .ok_or(D::Error::custom("Invalid builtin name"))
     }