serde: implement get_offset for KakarotSerde

crates/exex/src/serde.rs

@@ -14,7 +14,7 @@ use thiserror::Error;
 
 /// Represents errors that can occur during the serialization and deserialization processes between
 /// Cairo VM programs and Rust representations.
-#[derive(Debug, Error)]
+#[derive(Debug, Error, PartialEq)]
 pub enum KakarotSerdeError {
     /// Error variant indicating that no identifier matching the specified name was found.
     #[error("Expected one struct named '{struct_name}', found 0 matches. Expected type: {expected_type:?}")]
@@ -285,6 +285,39 @@ impl KakarotSerde {
         // Creates a `U256` value from the concatenated big-endian byte array.
         Ok(U256::from_be_slice(&bytes))
     }
+
+    /// Calculates the offset size for a given Cairo type.
+    ///
+    /// This function computes the offset based on the provided Cairo type:
+    /// - For a tuple, it returns the number of elements in the tuple.
+    /// - For a struct, it finds the identifier, then counts its members.
+    ///   - If no members are found (`None`, which should never happen), it returns an offset of `0`
+    ///     (an empty struct).
+    /// - For all other types, it defaults to an offset of `1`.
+    pub fn get_offset(&self, cairo_type: CairoType) -> Result<usize, KakarotSerdeError> {
+        // Determine offset based on the Cairo type.
+        match cairo_type {
+            // For a tuple type, return the number of elements in the tuple.
+            CairoType::Tuple { members, .. } => Ok(members.len()),
+
+            // For a struct type, attempt to retrieve its identifier using the scope (struct's name
+            // path).
+            CairoType::Struct { scope, .. } => {
+                // Retrieve the struct identifier and verify its presence in the program.
+                let identifier =
+                    self.get_identifier(&scope.path.join("."), Some("struct".to_string()))?;
+
+                // Return the count of struct members.
+                //
+                // If members are `None`, return an offset of `0`, indicating an empty struct.
+                // We should never encounter a struct with no members, so this is a safeguard.
+                Ok(identifier.members.unwrap_or_default().len())
+            }
+
+            // For all other Cairo types, return a default offset of 1.
+            _ => Ok(1),
+        }
+    }
 }
 
 #[cfg(test)]
@@ -667,6 +700,99 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_get_offset_tuple() {
+        // Setup the KakarotSerde instance
+        let kakarot_serde = setup_kakarot_serde();
+
+        // Create Cairo types for Tuple members.
+        let member1 = TupleItem::new(Some("a".to_string()), CairoType::felt_type(None), None);
+        let member2 = TupleItem::new(
+            Some("b".to_string()),
+            CairoType::pointer_type(CairoType::felt_type(None), None),
+            None,
+        );
+
+        // Create a Cairo type for a Tuple.
+        let cairo_type =
+            CairoType::tuple_from_members(vec![member1.clone(), member2.clone()], true, None);
+
+        // Assert that the offset of the Tuple is equal to the number of members.
+        assert_eq!(kakarot_serde.get_offset(cairo_type).unwrap(), 2);
+    }
+
+    #[test]
+    fn test_get_offset_struct_invalid_identifier() {
+        // Setup the KakarotSerde instance
+        let kakarot_serde = setup_kakarot_serde();
+
+        // Create a Cairo type for a Struct with an invalid identifier.
+        let cairo_type = CairoType::Struct {
+            scope: ScopedName {
+                path: vec!["an".to_string(), "invalid".to_string(), "path".to_string()],
+            },
+            location: None,
+        };
+
+        // Assert that the offset of the struct is an error due to the invalid identifier.
+        assert_eq!(
+            kakarot_serde.get_offset(cairo_type),
+            Err(KakarotSerdeError::IdentifierNotFound {
+                struct_name: "an.invalid.path".to_string(),
+                expected_type: Some("struct".to_string()),
+            })
+        );
+    }
+
+    #[test]
+    fn test_get_offset_struct_valid_identifier() {
+        // Setup the KakarotSerde instance
+        let kakarot_serde = setup_kakarot_serde();
+
+        // Create a Cairo type for a Struct with a valid identifier (3 members).
+        let cairo_type = CairoType::Struct {
+            scope: ScopedName { path: vec!["main".to_string(), "ImplicitArgs".to_string()] },
+            location: None,
+        };
+
+        // Assert that the offset of the struct is equal to the number of members (3).
+        assert_eq!(kakarot_serde.get_offset(cairo_type).unwrap(), 3);
+    }
+
+    #[test]
+    fn test_get_offset_struct_valid_identifier_without_members() {
+        // Setup the KakarotSerde instance
+        let kakarot_serde = setup_kakarot_serde();
+
+        // Create a Cairo type for a Struct with a valid identifier (no members).
+        let cairo_type = CairoType::Struct {
+            scope: ScopedName { path: vec!["alloc".to_string(), "ImplicitArgs".to_string()] },
+            location: None,
+        };
+
+        // Assert that the offset of the struct is 0 since the identifier has no members.
+        assert_eq!(kakarot_serde.get_offset(cairo_type).unwrap(), 0);
+    }
+
+    #[test]
+    fn test_get_offset_felt_pointer() {
+        // Setup the KakarotSerde instance
+        let kakarot_serde = setup_kakarot_serde();
+
+        // Create a Cairo type for a Felt.
+        let cairo_type = CairoType::felt_type(None);
+
+        // Assert that the offset of the Felt is 1 (value by default).
+        assert_eq!(kakarot_serde.get_offset(cairo_type).unwrap(), 1);
+
+        // Create a Cairo type for a Pointer.
+        let pointee_type = CairoType::felt_type(None);
+        let cairo_type = CairoType::pointer_type(pointee_type.clone(), None);
+
+        // Assert that the offset of the Pointer is 1 (value by default).
+        assert_eq!(kakarot_serde.get_offset(cairo_type).unwrap(), 1);
+    }
+
     #[test]
     fn test_cairo_type_struct_type() {
         // A dummy scope name for the struct type.