re_viewport 0.32.0

use std::collections::BTreeMap;

use ahash::HashMap;
use nohash_hasher::IntMap;
use rayon::prelude::*;
use re_viewer_context::{
    MissingChunkReporter, PerVisualizerTypeInViewClass, SystemExecutionOutput,
    ViewContextCollection, ViewContextSystemOncePerFrameResult, ViewId, ViewQuery, ViewState,
    ViewStates, ViewSystemExecutionError, ViewSystemIdentifier, ViewSystemState, ViewerContext,
    VisualizerCollection, VisualizerExecutionOutput, VisualizerInstructionsPerType,
};
use re_viewport_blueprint::ViewBlueprint;

use crate::view_highlights::highlights_for_view;

fn run_view_systems(
    ctx: &ViewerContext<'_>,
    view: &ViewBlueprint,
    query: &ViewQuery<'_>,
    view_state: &dyn ViewState,
    context_system_once_per_frame_results: &IntMap<
        ViewSystemIdentifier,
        ViewContextSystemOncePerFrameResult,
    >,
    context_systems: &mut ViewContextCollection,
    view_systems: &VisualizerCollection,
) -> PerVisualizerTypeInViewClass<Result<VisualizerExecutionOutput, ViewSystemExecutionError>> {
    re_tracing::profile_function!(view.class_identifier().as_str());

    let view_ctx = view.bundle_context_with_state(ctx, view_state);

    {
        re_tracing::profile_wait!("ViewContextSystem::execute");
        context_systems
            .systems
            .par_iter_mut()
            .for_each(|(name, (view_ctx_system, state))| {
                re_tracing::profile_scope!("ViewContextSystem::execute", name.as_str());
                let missing_chunk_reporter = MissingChunkReporter::default();
                let once_per_frame_result = context_system_once_per_frame_results
                    .get(name)
                    .expect("Context system execution result didn't occur");
                view_ctx_system.execute(
                    &view_ctx,
                    &missing_chunk_reporter,
                    query,
                    once_per_frame_result,
                );
                *state = ViewSystemState {
                    any_missing_chunks: missing_chunk_reporter.any_missing(),
                };
            });
    };

    re_tracing::profile_wait!("VisualizerSystem::execute");
    let per_visualizer_type_results = view_systems
        .systems
        .par_iter()
        .map(|(name, vis_system)| {
            re_tracing::profile_scope!("VisualizerSystem::execute", name.as_str());
            let affinity = vis_system.affinity();
            let result = vis_system
                .execute(&view_ctx, query, context_systems)
                .map(|mut output| {
                    output.affinity = affinity;
                    output
                });
            (*name, result)
        })
        .collect();

    PerVisualizerTypeInViewClass {
        view_class_identifier: view.class_identifier(),
        per_visualizer: per_visualizer_type_results,
    }
}

/// Creates a new [`ViewQuery`] for the given view.
pub fn new_view_query<'a>(ctx: &'a ViewerContext<'a>, view: &'a ViewBlueprint) -> ViewQuery<'a> {
    let highlights = highlights_for_view(ctx, view.id);

    let query_result = ctx.lookup_query_result(view.id);

    let mut active_visualizer_instructions_per_type = VisualizerInstructionsPerType::default();
    {
        re_tracing::profile_scope!("active_visualizer_instructions_per_type");

        for data_result in query_result.tree.iter_data_results() {
            if !data_result.visible {
                continue;
            }

            for instruction in &data_result.visualizer_instructions {
                active_visualizer_instructions_per_type
                    .entry(instruction.visualizer_type)
                    .or_default()
                    .push((data_result, instruction));
            }
        }
    }

    let current_query = ctx.time_ctrl.current_query();
    re_viewer_context::ViewQuery {
        view_id: view.id,
        space_origin: &view.space_origin,
        active_visualizer_instructions_per_type,
        timeline: current_query.timeline(),
        latest_at: current_query.at(),
        highlights,
    }
}

pub fn execute_systems_for_view<'a>(
    ctx: &'a ViewerContext<'_>,
    view: &'a ViewBlueprint,
    view_state: &dyn ViewState,
    context_system_once_per_frame_results: &IntMap<
        ViewSystemIdentifier,
        ViewContextSystemOncePerFrameResult,
    >,
) -> (ViewQuery<'a>, SystemExecutionOutput) {
    re_tracing::profile_function!(view.class_identifier().as_str());

    let query = new_view_query(ctx, view);

    let mut context_systems = ctx
        .view_class_registry()
        .new_context_collection(view.class_identifier());
    let view_systems = ctx
        .view_class_registry()
        .new_visualizer_collection(view.class_identifier());

    let visualizer_execution_output = run_view_systems(
        ctx,
        view,
        &query,
        view_state,
        context_system_once_per_frame_results,
        &mut context_systems,
        &view_systems,
    );

    (
        query,
        SystemExecutionOutput {
            context_systems,
            visualizer_execution_output,
        },
    )
}

pub fn execute_systems_for_all_views<'a>(
    ctx: &'a ViewerContext<'a>,
    tree: &egui_tiles::Tree<ViewId>,
    views: &'a BTreeMap<ViewId, ViewBlueprint>,
    view_states: &mut ViewStates,
) -> HashMap<ViewId, (ViewQuery<'a>, SystemExecutionOutput)> {
    re_tracing::profile_wait!("execute_systems");

    let store_id = ctx.store_id();

    // During system execution we only have read access to the view states, so we need to ensure they exist ahead of time.
    for (view_id, view) in views {
        view_states.ensure_state_exists(store_id, *view_id, view.class(ctx.view_class_registry()));
    }

    // Once-per-frame context system execution.
    // The same context system class may be used by several view classes, so we have to do this before
    // running anything per-view.
    let context_system_once_per_frame_results = ctx
        .view_class_registry()
        .run_once_per_frame_context_systems(
            ctx,
            views.values().map(|view| view.class_identifier()),
        );

    tree.active_tiles()
        .into_par_iter()
        .filter_map(|tile_id| {
            let tile = tree.tiles.get(tile_id)?;
            match tile {
                egui_tiles::Tile::Pane(view_id) => {
                    let view = views.get(view_id)?;
                    let Some(view_state) = view_states.get(store_id, *view_id) else {
                        re_log::debug_panic!("View state for view {view_id:?} not found. That shouldn't be possible since we just ensured they exist above.");
                        return None;
                    };

                    let result = execute_systems_for_view(ctx, view, view_state, &context_system_once_per_frame_results);
                    Some((*view_id, result))
                },
                egui_tiles::Tile::Container(_) => None,
            }
        })
        .collect::<HashMap<_, _>>()
}