next_core/next_client_reference/

visit_client_reference.rs

use std::future::Future;

use anyhow::Result;
use rustc_hash::FxHashSet;
use serde::{Deserialize, Serialize};
use tracing::Instrument;
use turbo_rcstr::RcStr;
use turbo_tasks::{
    FxIndexMap, FxIndexSet, NonLocalValue, ReadRef, ResolvedVc, TryJoinIterExt, ValueToString, Vc,
    debug::ValueDebugFormat,
    graph::{AdjacencyMap, GraphTraversal, Visit, VisitControlFlow},
    trace::TraceRawVcs,
};
use turbopack::css::chunk::CssChunkPlaceable;
use turbopack_core::{
    chunk::ChunkingType, module::Module, reference::primary_chunkable_referenced_modules,
};

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,
    Serialize,
    Deserialize,
    Debug,
    ValueDebugFormat,
    TraceRawVcs,
    NonLocalValue,
)]
pub struct ClientReference {
    pub server_component: Option<ResolvedVc<NextServerComponentModule>>,
    pub ty: ClientReferenceType,
}

impl ClientReference {
    pub fn server_component(&self) -> Option<ResolvedVc<NextServerComponentModule>> {
        self.server_component
    }

    pub fn ty(&self) -> ClientReferenceType {
        self.ty
    }
}

#[derive(
    Copy,
    Clone,
    Eq,
    PartialEq,
    Hash,
    Serialize,
    Deserialize,
    Debug,
    ValueDebugFormat,
    TraceRawVcs,
    NonLocalValue,
)]
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>,
    /// Only the [`ClientReferenceType::EcmascriptClientReference`]s are listed in this map.
    #[allow(clippy::type_complexity)]
    pub client_references_by_server_component:
        FxIndexMap<Option<ResolvedVc<NextServerComponentModule>>, Vec<ResolvedVc<Box<dyn Module>>>>,
    pub server_component_entries: Vec<ResolvedVc<NextServerComponentModule>>,
    pub server_utils: Vec<ResolvedVc<NextServerUtilityModule>>,
}

#[turbo_tasks::value(shared)]
pub struct VisitedClientReferenceGraphNodes(FxHashSet<FindServerEntriesNode>);

#[turbo_tasks::value_impl]
impl VisitedClientReferenceGraphNodes {
    #[turbo_tasks::function]
    pub fn empty() -> Vc<Self> {
        VisitedClientReferenceGraphNodes(Default::default()).cell()
    }
}

#[turbo_tasks::value(transparent)]
pub struct ClientReferenceTypes(FxIndexSet<ClientReferenceType>);

#[turbo_tasks::value_impl]
impl ClientReferenceGraphResult {
    #[turbo_tasks::function]
    pub fn types(&self) -> Vc<ClientReferenceTypes> {
        Vc::cell(
            self.client_references
                .iter()
                .map(|r| r.ty())
                .collect::<FxIndexSet<_>>(),
        )
    }
}

impl ClientReferenceGraphResult {
    /// Merges multiple return values of client_reference_graph together.
    pub fn extend(&mut self, other: &Self) {
        self.client_references
            .extend(other.client_references.iter().copied());
        for (k, v) in other.client_references_by_server_component.iter() {
            self.client_references_by_server_component
                .entry(*k)
                .or_insert_with(Vec::new)
                .extend(v);
        }
        self.server_component_entries
            .extend(other.server_component_entries.iter().copied());
        self.server_utils.extend(other.server_utils.iter().copied());
    }
}

#[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,
) -> Result<Vc<ServerEntries>> {
    async move {
        let graph = AdjacencyMap::new()
            .skip_duplicates()
            .visit(
                vec![FindServerEntriesNode::Internal(
                    entry,
                    entry.ident().to_string().await?,
                )],
                FindServerEntries { include_traced },
            )
            .await
            .completed()?
            .into_inner();

        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 {
    /// Whether to walk ChunkingType::Traced references
    include_traced: bool,
}

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

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

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

    fn edges(&mut self, node: &FindServerEntriesNode) -> Self::EdgesFuture {
        let node = node.clone();
        let include_traced = self.include_traced;
        async move {
            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, _) => ResolvedVc::upcast(module),
                FindServerEntriesNode::ServerComponentEntry(module, _) => {
                    ResolvedVc::upcast(module)
                }
            };

            // Pass include_traced 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).await?;

            let referenced_modules = referenced_modules
                .iter()
                .flat_map(|(chunking_type, _, modules)| match chunking_type {
                    ChunkingType::Traced => None,
                    _ => Some(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,
                            server_component_asset.ident().to_string().await?,
                        ));
                    }

                    if let Some(server_util_module) =
                        ResolvedVc::try_downcast_type::<NextServerUtilityModule>(*module)
                    {
                        return Ok(FindServerEntriesNode::ServerUtilEntry(
                            server_util_module,
                            module.ident().to_string().await?,
                        ));
                    }

                    Ok(FindServerEntriesNode::Internal(
                        *module,
                        module.ident().to_string().await?,
                    ))
                });

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

            Ok(assets)
        }
    }

    fn span(&mut self, node: &FindServerEntriesNode) -> tracing::Span {
        match node {
            FindServerEntriesNode::ClientReference => {
                tracing::info_span!("client reference")
            }
            FindServerEntriesNode::Internal(_, name) => {
                tracing::info_span!("module", name = name.to_string())
            }
            FindServerEntriesNode::ServerUtilEntry(_, name) => {
                tracing::info_span!("server util", name = name.to_string())
            }
            FindServerEntriesNode::ServerComponentEntry(_, name) => {
                tracing::info_span!("layout segment", name = name.to_string())
            }
        }
    }
}