Feature: SET_TREE_STRUCTURE hint

src/hints/mod.rs

@@ -28,6 +28,7 @@ use crate::io::input::StarknetOsInput;
 pub mod block_context;
 pub mod builtins;
 pub mod execution;
+mod output;
 mod patricia;
 pub mod state;
 pub mod syscalls;
@@ -45,7 +46,7 @@ type HintImpl = fn(
     &HashMap<String, Felt252>,
 ) -> Result<(), HintError>;
 
-static HINTS: [(&str, HintImpl); 100] = [
+static HINTS: [(&str, HintImpl); 101] = [
     (INITIALIZE_CLASS_HASHES, initialize_class_hashes),
     (INITIALIZE_STATE_CHANGES, initialize_state_changes),
     (IS_N_GE_TWO, is_n_ge_two),
@@ -125,8 +126,9 @@ static HINTS: [(&str, HintImpl); 100] = [
     (execution::TX_RESOURCE_BOUNDS_LEN, execution::tx_resource_bounds_len),
     (execution::TX_TIP, execution::tx_tip),
     (execution::WRITE_SYSCALL_RESULT, execution::write_syscall_result),
-    (state::LOAD_EDGE, state::load_edge),
+    (output::SET_TREE_STRUCTURE, output::set_tree_structure),
     (patricia::SET_SIBLINGS, patricia::set_siblings),
+    (state::LOAD_EDGE, state::load_edge),
     (state::SET_PREIMAGE_FOR_CLASS_COMMITMENTS, state::set_preimage_for_class_commitments),
     (state::SET_PREIMAGE_FOR_CURRENT_COMMITMENT_INFO, state::set_preimage_for_current_commitment_info),
     (state::SET_PREIMAGE_FOR_STATE_COMMITMENTS, state::set_preimage_for_state_commitments),

src/hints/output.rs

@@ -0,0 +1,183 @@
+use std::cmp::min;
+use std::collections::HashMap;
+
+use cairo_vm::hint_processor::builtin_hint_processor::hint_utils::get_ptr_from_var_name;
+use cairo_vm::hint_processor::hint_processor_definition::HintReference;
+use cairo_vm::serde::deserialize_program::ApTracking;
+use cairo_vm::types::exec_scope::ExecutionScopes;
+use cairo_vm::vm::errors::hint_errors::HintError;
+use cairo_vm::vm::errors::vm_errors::VirtualMachineError;
+use cairo_vm::vm::runners::builtin_runner::{BuiltinRunner, OutputBuiltinRunner};
+use cairo_vm::vm::vm_core::VirtualMachine;
+use cairo_vm::Felt252;
+use indoc::indoc;
+use num_integer::div_ceil;
+
+use crate::hints::vars;
+
+const MAX_PAGE_SIZE: usize = 3800;
+
+// TODO: replace by vm.get_output_builtin() once it's merged upstream
+pub fn get_output_builtin(vm: &mut VirtualMachine) -> Result<&mut OutputBuiltinRunner, VirtualMachineError> {
+    for builtin in vm.get_builtin_runners_as_mut() {
+        if let BuiltinRunner::Output(output_builtin) = builtin {
+            return Ok(output_builtin);
+        };
+    }
+
+    Err(VirtualMachineError::NotImplemented)
+}
+
+pub const SET_TREE_STRUCTURE: &str = indoc! {r#"
+	from starkware.python.math_utils import div_ceil
+	onchain_data_start = ids.da_start
+	onchain_data_size = ids.output_ptr - onchain_data_start
+
+	max_page_size = 3800
+	n_pages = div_ceil(onchain_data_size, max_page_size)
+	for i in range(n_pages):
+	    start_offset = i * max_page_size
+	    output_builtin.add_page(
+	        page_id=1 + i,
+	        page_start=onchain_data_start + start_offset,
+	        page_size=min(onchain_data_size - start_offset, max_page_size),
+	    )
+	# Set the tree structure to a root with two children:
+	# * A leaf which represents the main part
+	# * An inner node for the onchain data part (which contains n_pages children).
+	#
+	# This is encoded using the following sequence:
+	output_builtin.add_attribute('gps_fact_topology', [
+	    # Push 1 + n_pages pages (all of the pages).
+	    1 + n_pages,
+	    # Create a parent node for the last n_pages.
+	    n_pages,
+	    # Don't push additional pages.
+	    0,
+	    # Take the first page (the main part) and the node that was created (onchain data)
+	    # and use them to construct the root of the fact tree.
+	    2,
+	])"#
+};
+
+pub fn set_tree_structure(
+    vm: &mut VirtualMachine,
+    _exec_scopes: &mut ExecutionScopes,
+    ids_data: &HashMap<String, HintReference>,
+    ap_tracking: &ApTracking,
+    _constants: &HashMap<String, Felt252>,
+) -> Result<(), HintError> {
+    let onchain_data_start = get_ptr_from_var_name(vars::ids::DA_START, vm, ids_data, ap_tracking)?;
+    let output_ptr = get_ptr_from_var_name(vars::ids::OUTPUT_PTR, vm, ids_data, ap_tracking)?;
+    let onchain_data_size = (output_ptr - onchain_data_start)?;
+
+    let output_builtin = get_output_builtin(vm)?;
+
+    let n_pages = div_ceil(onchain_data_size, MAX_PAGE_SIZE);
+    for i in 0..n_pages {
+        let start_offset = i * MAX_PAGE_SIZE;
+        let page_id = i + 1;
+        let page_start = (onchain_data_start + start_offset)?;
+        let page_size = min(onchain_data_size - start_offset, MAX_PAGE_SIZE);
+        output_builtin
+            .add_page(page_id, page_start, page_size)
+            .map_err(|e| HintError::CustomHint(e.to_string().into_boxed_str()))?;
+    }
+
+    // TODO: replace the get_state() / set_state() sequence with add_attribute()
+    //       once https://github.com/lambdaclass/cairo-vm/pull/1691 is merged.
+    let mut builtin_state = output_builtin.get_state();
+    // Set the tree structure to a root with two children:
+    // * A leaf which represents the main part
+    // * An inner node for the onchain data part (which contains n_pages children).
+    //
+    // This is encoded using the following sequence:
+    builtin_state.attributes.insert(
+        "gps_fact_topology".to_string(),
+        vec![
+            // Push 1 + n_pages pages (all of the pages).
+            1 + n_pages,
+            // Create a parent node for the last n_pages.
+            n_pages,
+            // Don't push additional pages.
+            0,
+            // Take the first page (the main part) and the node that was created (onchain data)
+            // and use them to construct the root of the fact tree.
+            2,
+        ],
+    );
+    output_builtin.set_state(builtin_state);
+
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use cairo_vm::hint_processor::builtin_hint_processor::hint_utils::insert_value_from_var_name;
+    use cairo_vm::types::relocatable::Relocatable;
+    use cairo_vm::vm::runners::cairo_pie::PublicMemoryPage;
+
+    use super::*;
+
+    #[test]
+    fn test_set_tree_structure() {
+        let mut vm = VirtualMachine::new(false);
+        vm.add_memory_segment();
+        vm.add_memory_segment();
+        vm.set_fp(2);
+
+        let mut output_builtin_runner = OutputBuiltinRunner::new(true);
+        output_builtin_runner.initialize_segments(&mut vm.segments);
+        let output_base = output_builtin_runner.base() as isize;
+        vm.builtin_runners.push(BuiltinRunner::Output(output_builtin_runner));
+
+        let ap_tracking = ApTracking::new();
+        let constants = HashMap::new();
+
+        let ids_data = HashMap::from([
+            (vars::ids::DA_START.to_string(), HintReference::new_simple(-2)),
+            (vars::ids::OUTPUT_PTR.to_string(), HintReference::new_simple(-1)),
+        ]);
+        insert_value_from_var_name(
+            vars::ids::DA_START,
+            Relocatable::from((output_base, 0)),
+            &mut vm,
+            &ids_data,
+            &ap_tracking,
+        )
+        .unwrap();
+        insert_value_from_var_name(
+            vars::ids::OUTPUT_PTR,
+            Relocatable::from((output_base, 10000)),
+            &mut vm,
+            &ids_data,
+            &ap_tracking,
+        )
+        .unwrap();
+
+        let mut exec_scopes: ExecutionScopes = Default::default();
+
+        set_tree_structure(&mut vm, &mut exec_scopes, &ids_data, &ap_tracking, &constants)
+            .expect("Hint should succeed");
+
+        let n_expected_pages: usize = 3;
+
+        let output_builtin = get_output_builtin(&mut vm).unwrap();
+        let builtin_state = output_builtin.get_state();
+
+        let pages = builtin_state.pages;
+        assert_eq!(pages.len(), n_expected_pages);
+        assert_eq!(
+            pages,
+            HashMap::from([
+                (1usize, PublicMemoryPage { start: 0, size: MAX_PAGE_SIZE }),
+                (2usize, PublicMemoryPage { start: MAX_PAGE_SIZE, size: MAX_PAGE_SIZE }),
+                (3usize, PublicMemoryPage { start: 2 * MAX_PAGE_SIZE, size: 2400 })
+            ])
+        );
+
+        let attributes = builtin_state.attributes;
+        let gps_fact_topology = attributes.get("gps_fact_topology").unwrap();
+        assert_eq!(gps_fact_topology, &vec![1 + n_expected_pages, n_expected_pages, 0, 2]);
+    }
+}

src/hints/unimplemented.rs

@@ -8,39 +8,6 @@ const CACHE_CONTRACT_STORAGE: &str = indoc! {r#"
 	assert ids.value == value, "Inconsistent storage value.""#
 };
 
-#[allow(unused)]
-const SET_TREE_STRUCTURE: &str = indoc! {r#"
-	from starkware.python.math_utils import div_ceil
-	onchain_data_start = ids.da_start
-	onchain_data_size = ids.output_ptr - onchain_data_start
-
-	max_page_size = 3800
-	n_pages = div_ceil(onchain_data_size, max_page_size)
-	for i in range(n_pages):
-	    start_offset = i * max_page_size
-	    output_builtin.add_page(
-	        page_id=1 + i,
-	        page_start=onchain_data_start + start_offset,
-	        page_size=min(onchain_data_size - start_offset, max_page_size),
-	    )
-	# Set the tree structure to a root with two children:
-	# * A leaf which represents the main part
-	# * An inner node for the onchain data part (which contains n_pages children).
-	#
-	# This is encoded using the following sequence:
-	output_builtin.add_attribute('gps_fact_topology', [
-	    # Push 1 + n_pages pages (all of the pages).
-	    1 + n_pages,
-	    # Create a parent node for the last n_pages.
-	    n_pages,
-	    # Don't push additional pages.
-	    0,
-	    # Take the first page (the main part) and the node that was created (onchain data)
-	    # and use them to construct the root of the fact tree.
-	    2,
-	])"#
-};
-
 #[allow(unused)]
 const ENTER_SCOPE_SYSCALL_HANDLER: &str = "vm_enter_scope({'syscall_handler': syscall_handler})";
 

src/hints/vars.rs

@@ -29,6 +29,7 @@ pub mod ids {
     pub const COMPILED_CLASS_FACT: &str = "compiled_class_fact";
     pub const CONTRACT_ADDRESS: &str = "contract_address";
     pub const CONTRACT_STATE_CHANGES: &str = "contract_state_changes";
+    pub const DA_START: &str = "da_start";
     pub const DEPRECATED_TX_INFO: &str = "deprecated_tx_info";
     pub const DEST_PTR: &str = "dest_ptr";
     pub const EDGE: &str = "edge";
@@ -42,6 +43,7 @@ pub mod ids {
     pub const NEW_STATE_ENTRY: &str = "new_state_entry";
     pub const NODE: &str = "node";
     pub const OS_CONTEXT: &str = "os_context";
+    pub const OUTPUT_PTR: &str = "output_ptr";
     pub const PREV_VALUE: &str = "prev_value";
     pub const REQUEST: &str = "request";
     pub const SECP_P: &str = "SECP_P";