Fix snapshot clones

src/libfuncs/dup.rs

@@ -41,21 +41,18 @@ pub fn build<'ctx, 'this>(
         .and_then(|meta| meta.wrap_invoke(&info.signature.param_signatures[0].ty))
     {
         Some(clone_fn) => {
+            let original_value = entry.argument(0)?.into();
             let (entry, cloned_value) = clone_fn(
                 context,
                 registry,
                 entry,
                 location,
                 helper,
                 metadata,
-                entry.argument(0)?.into(),
+                original_value,
             )?;
 
-            entry.append_operation(helper.br(
-                0,
-                &[entry.argument(0)?.into(), cloned_value],
-                location,
-            ));
+            entry.append_operation(helper.br(0, &[original_value, cloned_value], location));
         }
         None => {
             entry.append_operation(helper.br(

src/metadata/enum_snapshot_variants.rs

@@ -9,14 +9,8 @@ pub struct EnumSnapshotVariantsMeta {
 }
 
 impl EnumSnapshotVariantsMeta {
-    pub fn set_mapping(
-        &mut self,
-        snapshot_id: &ConcreteTypeId,
-        enum_variants: Option<&[ConcreteTypeId]>,
-    ) {
-        if let Some(variants) = enum_variants {
-            self.map.insert(snapshot_id.clone(), variants.to_vec());
-        }
+    pub fn set_mapping(&mut self, snapshot_id: &ConcreteTypeId, variants: &[ConcreteTypeId]) {
+        self.map.insert(snapshot_id.clone(), variants.to_vec());
     }
 
     pub fn get_variants(&self, snapshot_id: &ConcreteTypeId) -> Option<&Vec<ConcreteTypeId>> {

src/metadata/snapshot_clones.rs

@@ -64,6 +64,12 @@ impl SnapshotClonesMeta {
         );
     }
 
+    pub(crate) fn register_dup(&mut self, id: ConcreteTypeId, from_id: &ConcreteTypeId) {
+        if let Some(clone_fn) = self.mappings.get(from_id) {
+            self.mappings.insert(id, clone_fn.clone());
+        }
+    }
+
     pub(crate) fn wrap_invoke(&self, id: &ConcreteTypeId) -> Option<CloneFnWrapper> {
         self.mappings.get(id).cloned()
     }

src/types/snapshot.rs

@@ -15,7 +15,10 @@
 use super::{TypeBuilder, WithSelf};
 use crate::{
     error::Result,
-    metadata::{enum_snapshot_variants::EnumSnapshotVariantsMeta, MetadataStorage},
+    metadata::{
+        enum_snapshot_variants::EnumSnapshotVariantsMeta, snapshot_clones::SnapshotClonesMeta,
+        MetadataStorage,
+    },
     utils::ProgramRegistryExt,
 };
 use cairo_lang_sierra::{
@@ -43,13 +46,19 @@ pub fn build<'ctx>(
     // This type is like a `Cow<T>` that clones whenever the original type is modified to keep the
     // original data. Since implementing that is complicated we can just clone the entire value for
     // now.
-    match metadata.get_mut::<EnumSnapshotVariantsMeta>() {
-        Some(x) => x,
-        None => metadata
-            .insert(EnumSnapshotVariantsMeta::default())
-            .expect("should not fail because we checked there is no metadata beforehand"),
+
+    // Register enum variants for the snapshot.
+    if let Some(variants) = registry.get_type(&info.ty)?.variants() {
+        metadata
+            .get_or_insert_with(EnumSnapshotVariantsMeta::default)
+            .set_mapping(info.self_ty, variants);
+    }
+
+    // Ensure the inner type is built and register the snapshot clone logic builder.
+    let self_ty = registry.build_type(context, module, registry, metadata, &info.ty)?;
+    if let Some(snapshot_clones_meta) = metadata.get_mut::<SnapshotClonesMeta>() {
+        snapshot_clones_meta.register_dup(info.self_ty.clone(), &info.ty);
     }
-    .set_mapping(info.self_ty, registry.get_type(&info.ty)?.variants());
 
-    registry.build_type(context, module, registry, metadata, &info.ty)
+    Ok(self_ty)
 }