Repository: MystenLabs/sui
File: external-crates/move/crates/move-vm-runtime/src/jit/execution/translate.rs (Rust)

Code before:
    }

    let input_modules: BTreeMap<ModuleId, &input::Module> = package_modules
        .iter()
        .map(|m| (m.compiled_module.self_id(), m))
        .collect();

    let mut state: BTreeMap<ModuleId, State> = BTreeMap::new();

    for root_id in input_modules.keys() {
        let root_key = package_context.interner.intern_ident_str(root_id.name())?;

        // Skip if we already fully processed this module.
        if matches!(state.get(root_id), Some(State::Visited)) {
            debug_assert!(package_context.loaded_modules.contains_key(&root_key));
            continue;
        }

        let mut stack: Vec<ModuleId> = Vec::new();
        stack.push(root_id.clone());

        while let Some(cur_id) = stack.pop() {
            let cur_state = *state.get(&cur_id).unwrap_or(&State::NotVisited);

            match cur_state {
                State::Visited => continue,
                State::Visiting => {
                    // all deps done, now load if needed
                    let key = package_context.interner.intern_ident_str(cur_id.name())?;
                    if !package_context.loaded_modules.contains_key(&key) {
                        let input_module = input_modules.get(&cur_id).ok_or_else(|| {
                            make_invariant_violation!(format!(
                                "Module {} not found in initial modules",
                                cur_id
                            ))
                        })?;
                        let loaded_module =
                            module(package_context, package_context.version_id, input_module)?;
                        assert!(
                            package_context
                                .loaded_modules
                                .insert(key, loaded_module)
                                .is_none()
                        );
                    }
                    state.insert(cur_id, State::Visited);
                }
                State::NotVisited => {
                    state.insert(cur_id.clone(), State::Visiting);
                    stack.push(cur_id.clone());


Review comment (bug):
Nit: may be worthwhile adding the debug assert that the module is in `loaded_modules` like you have above. 

Code after:
                        }
                    }
                    state.insert(cur_id, State::Visited);
                }
                State::NotVisited => {
                    if state.insert(cur_id.clone(), State::Visiting).is_some() {
                        return Err(make_invariant_violation!(format!(
                            "Module {} added to load queue as unvisited twice",
                            cur_id
                        )));
                    }
                    stack.push(cur_id.clone());

                    let input_module = input_modules.get(&cur_id).ok_or_else(|| {
                        make_invariant_violation!(format!(
                            "Module {} not found in initial modules",
                            cur_id
                        ))
                    })?;

                    for dep in input_module
                        .compiled_module
                        .immediate_dependencies()
                        .iter()
                        .filter(|dep| pkg_module_ids.contains(dep) && *dep != &cur_id)
                    {
                        match state.get(dep).copied().unwrap_or(State::NotVisited) {
                            State::Visited => { /* nothing */ }
                            State::Visiting => {
                                return Err(make_invariant_violation!(format!(
                                    "Cycle detected when loading module for package: {}",
                                    dep
                                )));
                            }
                            State::NotVisited => {
                                stack.push(dep.clone());
                            }
                        }
                    }
                }
            }
        }
    }

    Ok(())
}

// -------------------------------------------------------------------------------------------------
// Module Translation

fn module(