Add cairo1-run pretty printing

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 #### Upcoming Changes
 
+* feat: Add cairo1-run output pretty-printing for felts, arrays/spans and dicts [#1630](https://github.com/lambdaclass/cairo-vm/pull/1630)
+
 #### [1.0.0-rc1] - 2024-02-23
 
 * Bump `starknet-types-core` dependency version to 0.0.9 [#1628](https://github.com/lambdaclass/cairo-vm/pull/1628)

cairo1-run/src/main.rs

@@ -6,27 +6,24 @@
 use cairo_vm::{
     air_public_input::PublicInputError,
     cairo_run::EncodeTraceError,
-    types::{errors::program_errors::ProgramError, relocatable::MaybeRelocatable},
-    vm::{
-        errors::{
-            memory_errors::MemoryError, runner_errors::RunnerError, trace_errors::TraceError,
-            vm_errors::VirtualMachineError,
-        },
-        vm_core::VirtualMachine,
+    types::errors::program_errors::ProgramError,
+    vm::errors::{
+        memory_errors::MemoryError, runner_errors::RunnerError, trace_errors::TraceError,
+        vm_errors::VirtualMachineError,
     },
     Felt252,
 };
 use clap::{Parser, ValueHint};
 use itertools::Itertools;
+use serialize_output::serialize_output;
 use std::{
     io::{self, Write},
-    iter::Peekable,
     path::PathBuf,
-    slice::Iter,
 };
 use thiserror::Error;
 
 pub mod cairo_run;
+pub mod serialize_output;
 
 #[derive(Parser, Debug)]
 #[clap(author, version, about, long_about = None)]
@@ -297,7 +294,7 @@
         Err(Error::Cli(err)) => err.exit(),
         Ok(output) => {
             if let Some(output_string) = output {
-                println!("Program Output : {}", output_string);
+                println!("{}", output_string);
             }
             Ok(())
         }
@@ -323,48 +320,6 @@
     }
 }
 
-pub fn serialize_output(vm: &VirtualMachine, return_values: &[MaybeRelocatable]) -> String {
-    let mut output_string = String::new();
-    let mut return_values_iter: Peekable<Iter<MaybeRelocatable>> = return_values.iter().peekable();
-    serialize_output_inner(&mut return_values_iter, &mut output_string, vm);
-    fn serialize_output_inner(
-        iter: &mut Peekable<Iter<MaybeRelocatable>>,
-        output_string: &mut String,
-        vm: &VirtualMachine,
-    ) {
-        while let Some(val) = iter.next() {
-            if let MaybeRelocatable::RelocatableValue(x) = val {
-                // Check if the next value is a relocatable of the same index
-                if let Some(MaybeRelocatable::RelocatableValue(y)) = iter.peek() {
-                    // Check if the two relocatable values represent a valid array in memory
-                    if x.segment_index == y.segment_index && x.offset <= y.offset {
-                        // Fetch the y value from the iterator so we don't serialize it twice
-                        iter.next();
-                        // Fetch array
-                        maybe_add_whitespace(output_string);
-                        output_string.push('[');
-                        let array = vm.get_continuous_range(*x, y.offset - x.offset).unwrap();
-                        let mut array_iter: Peekable<Iter<MaybeRelocatable>> =
-                            array.iter().peekable();
-                        serialize_output_inner(&mut array_iter, output_string, vm);
-                        output_string.push(']');
-                        continue;
-                    }
-                }
-            }
-            maybe_add_whitespace(output_string);
-            output_string.push_str(&val.to_string());
-        }
-    }
-
-    fn maybe_add_whitespace(string: &mut String) {
-        if !string.is_empty() && !string.ends_with('[') {
-            string.push(' ');
-        }
-    }
-    output_string
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -562,4 +517,22 @@
         let args = args.iter().cloned().map(String::from);
         assert_matches!(run(args), Ok(Some(res)) if res.is_empty());
     }
+
+    #[rstest]
+    #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/felt_dict.cairo", "--print_output", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo", "--cairo_pie_output", "/dev/null"].as_slice())]
+    #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/felt_dict.cairo", "--print_output", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo", "--proof_mode", "--air_public_input", "/dev/null", "--air_private_input", "/dev/null"].as_slice())]
+    fn test_run_felt_dict(#[case] args: &[&str]) {
+        let args = args.iter().cloned().map(String::from);
+        let expected_output = "{\n\t0x10473: [0x8,0x9,0xa,0xb,],\n\t0x10474: [0x1,0x2,0x3,],\n}\n";
+        assert_matches!(run(args), Ok(Some(res)) if res == expected_output);
+    }
+
+    #[rstest]
+    #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/felt_span.cairo", "--print_output", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo", "--cairo_pie_output", "/dev/null"].as_slice())]
+    #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/felt_span.cairo", "--print_output", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo", "--proof_mode", "--air_public_input", "/dev/null", "--air_private_input", "/dev/null"].as_slice())]
+    fn test_run_felt_span(#[case] args: &[&str]) {
+        let args = args.iter().cloned().map(String::from);
+        let expected_output = "[0x8,0x9,0xa,0xb,]";
+        assert_matches!(run(args), Ok(Some(res)) if res == expected_output);
+    }
 }

cairo1-run/src/serialize_output.rs

@@ -0,0 +1,166 @@
+use cairo_vm::{
+    types::relocatable::{MaybeRelocatable, Relocatable},
+    vm::{errors::memory_errors::MemoryError, vm_core::VirtualMachine},
+    Felt252,
+};
+use itertools::Itertools;
+use std::{collections::HashMap, iter::Peekable, slice::Iter};
+use thiserror::Error;
+
+#[derive(Debug)]
+pub(crate) enum Output {
+    Felt(Felt252),
+    FeltSpan(Vec<Output>),
+    FeltDict(HashMap<Felt252, Output>),
+}
+
+#[derive(Debug, Error)]
+pub struct FormatError;
+
+impl std::fmt::Display for FormatError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "Format error")
+    }
+}
+
+impl Output {
+    pub fn from_memory(
+        vm: &VirtualMachine,
+        relocatable: &Relocatable,
+    ) -> Result<Self, FormatError> {
+        match vm.get_relocatable(*relocatable) {
+            Ok(relocatable_value) => {
+                let segment_size = vm
+                    .get_segment_size(relocatable_value.segment_index as usize)
+                    .ok_or(FormatError)?;
+                let segment_data = vm
+                    .get_continuous_range(relocatable_value, segment_size)
+                    .map_err(|_| FormatError)?;
+
+                // check if the segment data is a valid array of felts
+                if segment_data
+                    .iter()
+                    .all(|v| matches!(v, MaybeRelocatable::Int(_)))
+                {
+                    let span_segment: Vec<Output> = segment_data
+                        .iter()
+                        .map(|v| Output::Felt(v.get_int().unwrap()))
+                        .collect();
+                    Ok(Output::FeltSpan(span_segment))
+                } else {
+                    Err(FormatError)
+                }
+            }
+            Err(MemoryError::UnknownMemoryCell(relocatable_value)) => {
+                // here we assume that the value is a dictionary
+                let mut felt252dict: HashMap<Felt252, Output> = HashMap::new();
+
+                let segment_size = vm
+                    .get_segment_size(relocatable_value.segment_index as usize)
+                    .ok_or(FormatError)?;
+                let mut segment_start = relocatable_value.clone();
+                segment_start.offset = 0;
+                let segment_data = vm
+                    .get_continuous_range(*segment_start, segment_size)
+                    .map_err(|_| FormatError)?;
+
+                for (dict_key, _, value_relocatable) in segment_data.iter().tuples() {
+                    let key = dict_key.get_int().ok_or(FormatError)?;
+                    let value_segment = value_relocatable.get_relocatable().ok_or(FormatError)?;
+                    let value = Output::from_memory(vm, &value_segment)?;
+                    felt252dict.insert(key, value);
+                }
+                Ok(Output::FeltDict(felt252dict))
+            }
+            _ => Err(FormatError),
+        }
+    }
+}
+
+impl std::fmt::Display for Output {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Output::Felt(felt) => write!(f, "{}", felt.to_hex_string()),
+            Output::FeltSpan(span) => {
+                write!(f, "[")?;
+                for elem in span {
+                    write!(f, "{}", elem)?;
+                    write!(f, ",")?;
+                }
+                write!(f, "]")?;
+                Ok(())
+            }
+            Output::FeltDict(felt_dict) => {
+                let mut keys: Vec<_> = felt_dict.keys().collect();
+                keys.sort();
+                writeln!(f, "{{")?;
+                for key in keys {
+                    writeln!(f, "\t{}: {},", key.to_hex_string(), felt_dict[key])?;
+                }
+                writeln!(f, "}}")?;
+                Ok(())
+            }
+        }
+    }
+}
+
+pub(crate) fn serialize_output(vm: &VirtualMachine, return_values: &[MaybeRelocatable]) -> String {
+    let mut output_string = String::new();
+    let mut return_values_iter: Peekable<Iter<MaybeRelocatable>> = return_values.iter().peekable();
+    serialize_output_inner(&mut return_values_iter, &mut output_string, vm);
+
+    output_string
+}
+
+fn maybe_add_whitespace(string: &mut String) {
+    if !string.is_empty() && !string.ends_with('[') {
+        string.push(' ');
+    }
+}
+
+fn serialize_output_inner(
+    iter: &mut Peekable<Iter<MaybeRelocatable>>,
+    output_string: &mut String,
+    vm: &VirtualMachine,
+) -> Result<(), FormatError> {
+    while let Some(val) = iter.next() {
+        match val {
+            MaybeRelocatable::Int(x) => {
+                maybe_add_whitespace(output_string);
+                output_string.push_str(&x.to_string());
+                continue;
+            }
+            MaybeRelocatable::RelocatableValue(x) if ((iter.len() + 1) % 2) == 0 /* felt array */ => {
+                // Check if the next value is a relocatable of the same index
+                let y = iter.next().unwrap().get_relocatable().ok_or(FormatError)?;
+                // Check if the two relocatable values represent a valid array in memory
+                if x.segment_index == y.segment_index && x.offset <= y.offset {
+                        // Fetch array
+                        maybe_add_whitespace(output_string);
+                        output_string.push('[');
+                        let array = vm.get_continuous_range(*x, y.offset - x.offset).map_err(|_| FormatError)?;
+                        let mut array_iter: Peekable<Iter<MaybeRelocatable>> =
+                            array.iter().peekable();
+                        serialize_output_inner(&mut array_iter, output_string, vm)?;
+                        output_string.push(']');
+                        continue;
+                }
+            },
+            MaybeRelocatable::RelocatableValue(x) if iter.len() > 1 => {
+                let mut segment_start = *x;
+                segment_start.offset = 0;
+                for elem in iter.into_iter() {
+                    let output_value = Output::from_memory(vm, &elem.get_relocatable().ok_or(FormatError)?)?;
+                    output_string.push_str(output_value.to_string().as_str())
+                }
+            }
+            MaybeRelocatable::RelocatableValue(x) => {
+                match Output::from_memory(vm, x) {
+                    Ok(output_value) => output_string.push_str(format!("{}", output_value).as_str()),
+                    Err(_) => output_string.push_str("The output could not be formatted"),
+                }
+            }
+        }
+    }
+    Ok(())
+}

cairo_programs/cairo-1-programs/felt_dict.cairo

@@ -0,0 +1,15 @@
+use core::nullable::{nullable_from_box, match_nullable, FromNullableResult};
+
+fn main() -> Felt252Dict<Nullable<Span<felt252>>> {
+    // Create the dictionary
+    let mut d: Felt252Dict<Nullable<Span<felt252>>> = Default::default();
+
+    // Create the array to insert
+    let a = array![8, 9, 10, 11];
+    let b = array![1, 2, 3];
+
+    // Insert it as a `Span`
+    d.insert(66675, nullable_from_box(BoxTrait::new(a.span())));
+    d.insert(66676, nullable_from_box(BoxTrait::new(b.span())));
+    d
+}

cairo_programs/cairo-1-programs/felt_span.cairo

@@ -0,0 +1,6 @@
+use core::nullable::{nullable_from_box, match_nullable, FromNullableResult};
+
+fn main() -> Nullable<Span<felt252>> {
+   let a = array![8, 9, 10, 11];
+   nullable_from_box(BoxTrait::new(a.span()))
+}