use ref count in the allocation pointer

Author: FrancoGiachetta

Cargo.toml

@@ -20,6 +20,7 @@ repository.workspace = true
 resolver = "2"
 
 [features]
+default = ["with-debug-utils"]
 with-cheatcode = []
 with-debug-utils = []
 with-mem-tracing = []

src/libfuncs/circuit.rs

@@ -8,11 +8,13 @@ use crate::{
     execution_result::{ADD_MOD_BUILTIN_SIZE, MUL_MOD_BUILTIN_SIZE, RANGE_CHECK96_BUILTIN_SIZE},
     libfuncs::r#struct::build_struct_value,
     metadata::{
-        drop_overrides::DropOverridesMeta, realloc_bindings::ReallocBindingsMeta,
-        runtime_bindings::RuntimeBindingsMeta, MetadataStorage,
+        debug_utils::DebugUtils, drop_overrides::DropOverridesMeta, realloc_bindings::ReallocBindingsMeta, runtime_bindings::RuntimeBindingsMeta, MetadataStorage
     },
     native_panic,
-    types::{circuit::build_u384_struct_type, TypeBuilder},
+    types::{
+        circuit::{build_u384_struct_type, calc_circuit_output_prefix_layout},
+        TypeBuilder,
+    },
     utils::{get_integer_layout, layout_repeat, ProgramRegistryExt},
 };
 use cairo_lang_sierra::{
@@ -363,13 +365,15 @@ fn build_eval<'ctx, 'this>(
             circuit_info.mul_offsets.len() * MUL_MOD_BUILTIN_SIZE,
         )?;
 
-        // Calculate capacity for array.
         let outputs_capacity = circuit_info.values.len();
         let u384_integer_layout = get_integer_layout(384);
-        let outputs_layout = layout_repeat(&u384_integer_layout, outputs_capacity)?.0;
-        let outputs_capacity_bytes = outputs_layout.pad_to_align().size();
+        let outputs_prefix_layout = calc_circuit_output_prefix_layout();
+        let outputs_layout = outputs_prefix_layout
+            .extend(layout_repeat(&u384_integer_layout, outputs_capacity)?.0)?
+            .0
+            .pad_to_align();
         let outputs_capacity_bytes_value =
-            ok_block.const_int(context, location, outputs_capacity_bytes, 64)?;
+            ok_block.const_int(context, location, outputs_layout.size(), 64)?;
 
         // Alloc memory for array.
         let ptr_ty = llvm::r#type::pointer(context, 0);
@@ -381,13 +385,19 @@ fn build_eval<'ctx, 'this>(
             location,
         )?)?;
 
+        // Insert initial reference count, equal to 1.
+        let k1 = ok_block.const_int(context, location, 1, 32)?;
+        ok_block.store(context, location, outputs_ptr, k1)?;
+
         // Insert evaluated gates into the array.
         for (i, gate) in gates.into_iter().enumerate() {
             let value_ptr = ok_block.gep(
                 context,
                 location,
                 outputs_ptr,
-                &[GepIndex::Const(i as i32)],
+                &[GepIndex::Const(
+                    outputs_prefix_layout.size() as i32 + i as i32,
+                )],
                 IntegerType::new(context, 384).into(),
             )?;
             ok_block.store(context, location, value_ptr, gate)?;
@@ -397,7 +407,6 @@ fn build_eval<'ctx, 'this>(
 
         // Build output struct
         let outputs_type_id = &info.branch_signatures()[0].vars[2].ty;
-        let ref_count = ok_block.const_int(context, location, 1, 8)?;
         let outputs = build_struct_value(
             context,
             registry,
@@ -406,7 +415,7 @@ fn build_eval<'ctx, 'this>(
             helper,
             metadata,
             outputs_type_id,
-            &[ref_count, outputs_ptr, modulus_struct],
+            &[outputs_ptr, modulus_struct],
         )?;
 
         helper.br(ok_block, 0, &[add_mod, mul_mod, outputs], location)?;
@@ -920,23 +929,34 @@ fn build_get_output<'ctx, 'this>(
         location,
         outputs,
         llvm::r#type::pointer(context, 0),
-        1,
+        0,
     )?;
     let modulus_struct = entry.extract_value(
         context,
         location,
         outputs,
         build_u384_struct_type(context),
-        2,
+        1,
     )?;
 
+    let circuit_output_prefix_offset = calc_circuit_output_prefix_layout().size();
+    metadata
+                .get_mut::<DebugUtils>()
+                .unwrap()
+                .debug_print(context, helper.module, &entry, "AFTER OUTPUT", location)?;
     let output_integer_ptr = entry.gep(
         context,
         location,
         circuit_ptr,
-        &[GepIndex::Const(output_idx as i32)],
+        &[GepIndex::Const(
+           circuit_output_prefix_offset as i32 + output_idx as i32,
+        )],
         u384_type,
     )?;
+    metadata
+                .get_mut::<DebugUtils>()
+                .unwrap()
+                .debug_print(context, helper.module, &entry, "BEFORE OUTPUT", location)?;
     let output_integer = entry.load(context, location, output_integer_ptr, u384_type)?;
     let output_struct = u384_integer_to_struct(context, entry, location, output_integer)?;
 

src/types/circuit.rs

@@ -6,8 +6,8 @@ use super::WithSelf;
 use crate::{
     error::{Result, SierraAssertError},
     metadata::{
-        drop_overrides::DropOverridesMeta, dup_overrides::DupOverridesMeta,
-        realloc_bindings::ReallocBindingsMeta, MetadataStorage,
+        debug_utils::DebugUtils, drop_overrides::DropOverridesMeta,
+        dup_overrides::DupOverridesMeta, realloc_bindings::ReallocBindingsMeta, MetadataStorage,
     },
     utils::{get_integer_layout, layout_repeat, ProgramRegistryExt},
 };
@@ -22,7 +22,7 @@ use cairo_lang_sierra::{
 };
 use melior::{
     dialect::{arith::CmpiPredicate, func, llvm, scf},
-    helpers::{ArithBlockExt, BuiltinBlockExt, LlvmBlockExt},
+    helpers::{ArithBlockExt, BuiltinBlockExt, GepIndex, LlvmBlockExt},
     ir::{r#type::IntegerType, Block, BlockLike, Location, Module, Region, Type, Value},
     Context,
 };
@@ -305,8 +305,6 @@ pub fn build_circuit_outputs<'ctx>(
     metadata: &mut MetadataStorage,
     info: WithSelf<InfoOnlyConcreteType>,
 ) -> Result<Type<'ctx>> {
-    let u8_ty = IntegerType::new(context, 8).into();
-
     DupOverridesMeta::register_with(
         context,
         module,
@@ -318,15 +316,42 @@ pub fn build_circuit_outputs<'ctx>(
             let region = Region::new();
             let value_ty = registry.build_type(context, module, metadata, info.self_ty())?;
             let entry = region.append_block(Block::new(&[(value_ty, location)]));
-            let k1 = entry.const_int(context, location, 1, 8)?;
+            let k1 = entry.const_int(context, location, 1, 32)?;
 
             let outputs = entry.arg(0)?;
-            let ref_count = entry.extract_value(context, location, outputs, u8_ty, 0)?;
+            let gates_ptr = entry.extract_value(
+                context,
+                location,
+                outputs,
+                llvm::r#type::pointer(context, 0),
+                0,
+            )?;
+            let ref_count_ptr = entry.gep(
+                context,
+                location,
+                gates_ptr,
+                &[GepIndex::Const(0)],
+                IntegerType::new(context, 8).into(),
+            )?;
+            let ref_count = entry.load(
+                context,
+                location,
+                ref_count_ptr,
+                IntegerType::new(context, 32).into(),
+            )?;
             let ref_count_inc = entry.addi(ref_count, k1, location)?;
 
-            let output_new = entry.insert_value(context, location, outputs, ref_count_inc, 0)?;
+            metadata
+                .get_mut::<DebugUtils>()
+                .unwrap()
+                .debug_print(context, module, &entry, "DUP", location)?;
+            metadata
+                .get_mut::<DebugUtils>()
+                .unwrap()
+                .print_i32(context, module, &entry, ref_count, location)?;
 
-            entry.append_operation(func::r#return(&[output_new, output_new], location));
+            entry.store(context, location, ref_count_ptr, ref_count_inc)?;
+            entry.append_operation(func::r#return(&[outputs, outputs], location));
 
             Ok(Some(region))
         },
@@ -342,14 +367,43 @@ pub fn build_circuit_outputs<'ctx>(
             let region = Region::new();
             let value_ty = registry.build_type(context, module, metadata, info.self_ty())?;
             let entry = region.append_block(Block::new(&[(value_ty, location)]));
-            let k1 = entry.const_int(context, location, 1, 8)?;
+            let k1 = entry.const_int(context, location, 1, 32)?;
 
             let outputs = entry.arg(0)?;
-            let ref_count = entry.extract_value(context, location, outputs, u8_ty, 0)?;
+            let gates_ptr = entry.extract_value(
+                context,
+                location,
+                outputs,
+                llvm::r#type::pointer(context, 0),
+                0,
+            )?;
+
+            let ref_count_ptr = entry.gep(
+                context,
+                location,
+                gates_ptr,
+                &[GepIndex::Const(0)],
+                IntegerType::new(context, 8).into(),
+            )?;
+            let ref_count = entry.load(
+                context,
+                location,
+                ref_count_ptr,
+                IntegerType::new(context, 32).into(),
+            )?;
 
             // Check that the reference counting is different from 1. If it is equeal to 1, then it is shared.
             let is_shared = entry.cmpi(context, CmpiPredicate::Ne, ref_count, k1, location)?;
 
+            metadata
+                .get_mut::<DebugUtils>()
+                .unwrap()
+                .debug_print(context, module, &entry, "DROP", location)?;
+            metadata
+                .get_mut::<DebugUtils>()
+                .unwrap()
+                .print_i32(context, module, &entry, ref_count, location)?;
+
             entry.append_operation(scf::r#if(
                 is_shared,
                 &[],
@@ -359,8 +413,7 @@ pub fn build_circuit_outputs<'ctx>(
                     let entry = region.append_block(Block::new(&[]));
                     let ref_count_dec = entry.subi(ref_count, k1, location)?;
 
-                    entry.insert_value(context, location, ref_count_dec, outputs, 0)?;
-
+                    entry.store(context, location, ref_count_ptr, ref_count_dec)?;
                     entry.append_operation(scf::r#yield(&[], location));
 
                     region
@@ -370,14 +423,6 @@ pub fn build_circuit_outputs<'ctx>(
                     let region = Region::new();
                     let entry = region.append_block(Block::new(&[]));
 
-                    let gates_ptr = entry.extract_value(
-                        context,
-                        location,
-                        outputs,
-                        llvm::r#type::pointer(context, 0),
-                        1,
-                    )?;
-
                     entry
                         .append_operation(ReallocBindingsMeta::free(context, gates_ptr, location)?);
                     entry.append_operation(scf::r#yield(&[], location));
@@ -396,7 +441,6 @@ pub fn build_circuit_outputs<'ctx>(
     Ok(llvm::r#type::r#struct(
         context,
         &[
-            u8_ty,
             llvm::r#type::pointer(context, 0),
             build_u384_struct_type(context),
         ],
@@ -546,3 +590,11 @@ pub fn build_u384_struct_type(context: &Context) -> Type<'_> {
         false,
     )
 }
+
+pub fn calc_circuit_output_prefix_layout() -> Layout {
+    let u384_layout = get_integer_layout(384);
+    get_integer_layout(32)
+        .align_to(u384_layout.align())
+        .expect("layout size rounded up to the next multiple of layout alignment should never be greater than ISIZE::MAX")
+        .pad_to_align()
+}