next_core/next_client_reference/

visit_client_reference.rs

use std::future::Future;

use anyhow::Result;
use bincode::{Decode, Encode};
use tracing::{Instrument, Level, Span};
use turbo_rcstr::RcStr;
use turbo_tasks::{
    NonLocalValue, ReadRef, ResolvedVc, TryJoinIterExt, Vc,
    debug::ValueDebugFormat,
    graph::{AdjacencyMap, GraphTraversal, Visit, VisitControlFlow},
};
use turbopack_core::{
    chunk::ChunkingType, module::Module, reference::primary_chunkable_referenced_modules,
};
use turbopack_css::chunk::CssChunkPlaceable;

use crate::{
    next_client_reference::{
        CssClientReferenceModule,
        ecmascript_client_reference::ecmascript_client_reference_module::EcmascriptClientReferenceModule,
    },
    next_server_component::server_component_module::NextServerComponentModule,
    next_server_utility::server_utility_module::NextServerUtilityModule,
};

#[derive(
    Copy, Clone, Eq, PartialEq, Hash, Debug, ValueDebugFormat, NonLocalValue, Encode, Decode,
)]
pub struct ClientReference {
    /// The module whose reference creates the client/SSR chunk group.
    pub parent_module: ResolvedVc<Box<dyn Module>>,
    /// The server component used to order this reference in the output.
    pub server_component: Option<ResolvedVc<NextServerComponentModule>>,
    pub ty: ClientReferenceType,
}

#[derive(
    Copy, Clone, Eq, PartialEq, Hash, Debug, ValueDebugFormat, NonLocalValue, Encode, Decode,
)]
pub enum ClientReferenceType {
    EcmascriptClientReference(ResolvedVc<EcmascriptClientReferenceModule>),
    CssClientReference(ResolvedVc<Box<dyn CssChunkPlaceable>>),
}

#[turbo_tasks::value(shared)]
#[derive(Clone, Debug, Default)]
pub struct ClientReferenceGraphResult {
    pub client_references: Vec<ClientReference>,
    pub server_component_entries: Vec<ResolvedVc<NextServerComponentModule>>,
    pub server_utils: Vec<ResolvedVc<NextServerUtilityModule>>,
}

#[turbo_tasks::value(shared)]
#[derive(Clone, Debug)]
pub struct ServerEntries {
    pub server_component_entries: Vec<ResolvedVc<NextServerComponentModule>>,
    pub server_utils: Vec<ResolvedVc<NextServerUtilityModule>>,
}

/// For a given RSC entry, finds all server components (i.e. layout segments) and server utils that
/// are referenced by the entry.
#[turbo_tasks::function]
pub async fn find_server_entries(
    entry: ResolvedVc<Box<dyn Module>>,
    include_traced: bool,
    include_binding_usage: bool,
) -> Result<Vc<ServerEntries>> {
    async move {
        let emit_spans = tracing::enabled!(Level::INFO);
        let graph = AdjacencyMap::new()
            .visit(
                vec![FindServerEntriesNode::Internal(
                    entry,
                    if emit_spans {
                        // INVALIDATION: we don't need to invalidate when the span name changes
                        Some(entry.ident_string().untracked().await?)
                    } else {
                        None
                    },
                )],
                FindServerEntries {
                    emit_spans,
                    include_traced,
                    include_binding_usage,
                },
            )
            .await
            .completed()?;

        let mut server_component_entries = vec![];
        let mut server_utils = vec![];
        for node in graph.postorder_topological() {
            match node {
                FindServerEntriesNode::ServerUtilEntry(server_util, _) => {
                    server_utils.push(*server_util);
                }
                FindServerEntriesNode::ServerComponentEntry(server_component, _) => {
                    server_component_entries.push(*server_component);
                }
                FindServerEntriesNode::Internal(_, _) | FindServerEntriesNode::ClientReference => {}
            }
        }

        Ok(ServerEntries {
            server_component_entries,
            server_utils,
        }
        .cell())
    }
    .instrument(tracing::info_span!("find server entries"))
    .await
}

struct FindServerEntries {
    emit_spans: bool,
    /// Whether to walk ChunkingType::Traced references
    include_traced: bool,
    /// Whether to read the binding usage information from modules
    include_binding_usage: bool,
}

#[derive(Clone, Eq, PartialEq, Hash, Debug, ValueDebugFormat, NonLocalValue)]
enum FindServerEntriesNode {
    ClientReference,
    ServerComponentEntry(
        ResolvedVc<NextServerComponentModule>,
        Option<ReadRef<RcStr>>,
    ),
    ServerUtilEntry(ResolvedVc<NextServerUtilityModule>, Option<ReadRef<RcStr>>),
    Internal(ResolvedVc<Box<dyn Module>>, Option<ReadRef<RcStr>>),
}

impl Visit<FindServerEntriesNode> for FindServerEntries {
    type EdgesIntoIter = Vec<(FindServerEntriesNode, ())>;
    type EdgesFuture = impl Future<Output = Result<Self::EdgesIntoIter>>;

    fn visit(&mut self, node: &FindServerEntriesNode, _edge: Option<&()>) -> VisitControlFlow {
        match node {
            FindServerEntriesNode::Internal(..) => VisitControlFlow::Continue,
            FindServerEntriesNode::ClientReference
            | FindServerEntriesNode::ServerUtilEntry(..)
            | FindServerEntriesNode::ServerComponentEntry(..) => VisitControlFlow::Skip,
        }
    }

    fn edges(&mut self, node: &FindServerEntriesNode) -> Self::EdgesFuture {
        let include_traced = self.include_traced;
        let include_binding_usage = self.include_binding_usage;
        let parent_module = match node {
            // This should never occur since we always skip visiting these
            // nodes' edges.
            FindServerEntriesNode::ClientReference => {
                unreachable!("ClientReference node should not be visited")
            }
            FindServerEntriesNode::Internal(module, _) => **module,
            FindServerEntriesNode::ServerUtilEntry(module, _) => Vc::upcast(**module),
            FindServerEntriesNode::ServerComponentEntry(module, _) => Vc::upcast(**module),
        };
        let emit_spans = self.emit_spans;
        async move {
            // Pass include_traced and include_binding_usage to reuse the same cached
            // `primary_chunkable_referenced_modules` task result, but the traced references will be
            // filtered out again afterwards.
            let referenced_modules = primary_chunkable_referenced_modules(
                parent_module,
                include_traced,
                include_binding_usage,
            )
            .await?;

            let referenced_modules = referenced_modules
                .iter()
                .flat_map(|(_, resolved)| match resolved.chunking_type {
                    ChunkingType::Traced { .. } => None,
                    _ => Some(resolved.modules.iter()),
                })
                .flatten()
                .map(async |module| {
                    if ResolvedVc::try_downcast_type::<EcmascriptClientReferenceModule>(*module)
                        .is_some()
                        || ResolvedVc::try_downcast_type::<CssClientReferenceModule>(*module)
                            .is_some()
                    {
                        return Ok((FindServerEntriesNode::ClientReference, ()));
                    }

                    if let Some(server_component_asset) =
                        ResolvedVc::try_downcast_type::<NextServerComponentModule>(*module)
                    {
                        return Ok((
                            FindServerEntriesNode::ServerComponentEntry(
                                server_component_asset,
                                if emit_spans {
                                    // INVALIDATION: we don't need to invalidate when the span name
                                    // changes
                                    Some(server_component_asset.ident_string().untracked().await?)
                                } else {
                                    None
                                },
                            ),
                            (),
                        ));
                    }

                    if let Some(server_util_module) =
                        ResolvedVc::try_downcast_type::<NextServerUtilityModule>(*module)
                    {
                        return Ok((
                            FindServerEntriesNode::ServerUtilEntry(
                                server_util_module,
                                if emit_spans {
                                    // INVALIDATION: we don't need to invalidate when the span name
                                    // changes
                                    Some(module.ident_string().untracked().await?)
                                } else {
                                    None
                                },
                            ),
                            (),
                        ));
                    }

                    Ok((
                        FindServerEntriesNode::Internal(
                            *module,
                            if emit_spans {
                                // INVALIDATION: we don't need to invalidate when the span name
                                // changes
                                Some(module.ident_string().untracked().await?)
                            } else {
                                None
                            },
                        ),
                        (),
                    ))
                });

            let assets = referenced_modules.try_join().await?;

            Ok(assets)
        }
    }

    fn span(&mut self, node: &FindServerEntriesNode, _edge: Option<&()>) -> tracing::Span {
        if !self.emit_spans {
            return Span::none();
        }
        match node {
            FindServerEntriesNode::ClientReference => {
                tracing::info_span!("client reference")
            }
            FindServerEntriesNode::Internal(_, name) => {
                tracing::info_span!("module", name = display(name.as_ref().unwrap()))
            }
            FindServerEntriesNode::ServerUtilEntry(_, name) => {
                tracing::info_span!("server util", name = display(name.as_ref().unwrap()))
            }
            FindServerEntriesNode::ServerComponentEntry(_, name) => {
                tracing::info_span!("layout segment", name = display(name.as_ref().unwrap()))
            }
        }
    }
}