Populate environment with assets

Signed-off-by: Paul-Elliot <[email protected]>

Author: panglesd

src/xref2/component.ml

@@ -527,6 +527,8 @@ module Element = struct
   (* No component for pages yet *)
   type page = [ `Page of Identifier.Page.t * Odoc_model.Lang.Page.t ]
 
+  type asset = [ `Asset of Identifier.AssetFile.t ]
+
   type label_parent = [ signature | type_ | page ]
 
   type fragment_type_parent = [ signature | datatype ]
@@ -543,7 +545,8 @@ module Element = struct
     | extension
     | extension_decl
     | field
-    | page ]
+    | page
+    | asset ]
 
   let identifier : [< any ] -> Odoc_model.Paths.Identifier.t =
     let open Odoc_model.Paths.Identifier in
@@ -561,6 +564,7 @@ module Element = struct
     | `Extension (id, _, _) -> (id :> t)
     | `ExtensionDecl (id, _) -> (id :> t)
     | `Page (id, _) -> (id :> t)
+    | `Asset id -> (id :> t)
 end
 
 module Fmt = struct

src/xref2/component.mli

@@ -491,6 +491,8 @@ module Element : sig
   (* No component for pages yet *)
   type page = [ `Page of Identifier.Page.t * Odoc_model.Lang.Page.t ]
 
+  type asset = [ `Asset of Identifier.AssetFile.t ]
+
   type label_parent = [ signature | type_ | page ]
 
   type fragment_type_parent = [ signature | datatype ]
@@ -507,7 +509,8 @@ module Element : sig
     | extension
     | extension_decl
     | field
-    | page ]
+    | page
+    | asset ]
 
   val identifier : [< any ] -> Identifier.t
 end

src/xref2/env.ml

@@ -84,6 +84,7 @@ type kind =
   | Kind_Exception
   | Kind_Extension
   | Kind_Field
+  | Kind_asset
 
 module ElementsByName : sig
   type t
@@ -163,7 +164,6 @@ type t = {
   resolver : resolver option;
   recorder : recorder option;
   fragmentroot : (int * Component.Signature.t) option;
-  parent_page : Identifier.Page.t option;  (** parent page *)
 }
 
 let is_linking env = env.linking
@@ -200,7 +200,6 @@ let empty =
     recorder = None;
     ambiguous_labels = Identifier.Maps.Label.empty;
     fragmentroot = None;
-    parent_page = None;
   }
 
 let add_fragment_root sg env =
@@ -358,6 +357,10 @@ let add_extension_constructor identifier
   add_to_elts Kind_Extension identifier (`Extension (identifier, ec, te)) env
   |> add_cdocs identifier ec.doc
 
+let add_asset identifier env =
+  if env.linking then add_to_elts Kind_asset identifier (`Asset identifier) env
+  else env
+
 let module_of_unit : Lang.Compilation_unit.t -> Component.Module.t =
  fun unit ->
   let id = (unit.id :> Paths.Identifier.Module.t) in
@@ -597,6 +600,9 @@ let s_fragment_type_parent : Component.Element.fragment_type_parent scope =
     | #Component.Element.fragment_type_parent as r -> Some r
     | _ -> None)
 
+let s_asset : Component.Element.asset scope =
+  make_scope (function #Component.Element.asset as r -> Some r | _ -> None)
+
 let len = ref 0
 
 let n = ref 0
@@ -806,6 +812,29 @@ let open_units resolver env =
       | _ -> env)
     env resolver.open_units
 
+let rec collect_assets env (page : Lang.Page.t) =
+  let env =
+    match page.name with
+    | { iv = `Page (Some parent, _); _ }
+    | { iv = `LeafPage (Some parent, _); _ } -> (
+        let parent_name = match parent.iv with `Page (_, name) -> name in
+        match lookup_page (PageName.to_string parent_name) env with
+        | None -> env
+        | Some parent_page -> collect_assets env parent_page)
+    | _ -> env
+  in
+  let env =
+    List.fold_left
+      (fun env new_asset ->
+        let id = Identifier.Mk.asset_file (page.name, new_asset) in
+        add_asset id env)
+      env
+      (List.filter_map
+         (function Lang.Page.Asset_child c -> Some c | _ -> None)
+         page.children)
+  in
+  env
+
 let env_of_unit t ~linking resolver =
   let open Lang.Compilation_unit in
   let initial_env =
@@ -814,20 +843,25 @@ let env_of_unit t ~linking resolver =
     let env = { empty with linking } in
     env |> add_module (t.id :> Identifier.Path.Module.t) dm m.doc
   in
-  let parent_page :> Identifier.Page.t option =
+  let initial_env = set_resolver initial_env resolver |> open_units resolver in
+  let initial_env =
     match t.id.iv with
-    | `Root (None, _) -> None
-    | `Root (Some parent, _) -> Some parent
+    | `Root (None, _) -> initial_env
+    | `Root (Some parent, _) -> (
+        let parent_name = match parent.iv with `Page (_, name) -> name in
+        match lookup_page (PageName.to_string parent_name) initial_env with
+        | None -> initial_env
+        | Some parent_page -> collect_assets initial_env parent_page)
   in
-  let initial_env = { initial_env with parent_page } in
-  set_resolver initial_env resolver |> open_units resolver
+  initial_env
 
 let open_page page env = add_docs page.Lang.Page.content env
 
 let env_of_page page resolver =
   let initial_env = open_page page empty in
-  let initial_env = { initial_env with parent_page = Some page.name } in
-  set_resolver initial_env resolver |> open_units resolver
+  let initial_env = set_resolver initial_env resolver |> open_units resolver in
+  let initial_env = collect_assets initial_env page in
+  initial_env
 
 let env_for_reference resolver =
   set_resolver empty resolver |> open_units resolver
@@ -889,5 +923,3 @@ let verify_lookups env lookups =
   | true, Some r -> r.lookups <- LookupTypeSet.union r.lookups lookups
   | _ -> ());
   result
-
-let parent_page env = env.parent_page

src/xref2/env.mli

@@ -146,6 +146,8 @@ val s_label_parent : Component.Element.label_parent scope
 
 val s_fragment_type_parent : Component.Element.fragment_type_parent scope
 
+val s_asset : Component.Element.asset scope
+
 (* val open_component_signature :
    Paths_types.Identifier.signature -> Component.Signature.t -> t -> t *)
 
@@ -184,5 +186,3 @@ val len : int ref
 val n : int ref
 
 val verify_lookups : t -> LookupTypeSet.t -> bool
-
-val parent_page : t -> Identifier.Page.t option

src/xref2/ref_tools.ml

@@ -74,6 +74,7 @@ let ref_kind_of_element = function
   | `ExtensionDecl _ -> "extension-decl"
   | `Field _ -> "field"
   | `Page _ -> "page"
+  | `Asset _ -> "asset"
 
 let ref_kind_of_find = function
   | `FModule _ | `FModule_subst _ -> "module"
@@ -580,6 +581,10 @@ end
 module A = struct
   (** Assets *)
 
+  let in_env env name =
+    env_lookup_by_name ~kind:`Asset Env.s_asset name env >>= fun (`Asset id) ->
+    Ok (`Identifier id)
+
   let rec in_page env (page : Odoc_model.Lang.Page.t) (asset_name : string) :
       (Reference.Resolved.Asset.t, _) result =
     let has_asset children asset =
@@ -854,12 +859,7 @@ let resolve_page_reference env (r : Reference.Page.t) =
 let resolve_asset_reference env (m : Reference.Asset.t) =
   match m with
   | `Resolved r -> Ok r
-  | `Root (name, _) -> (
-      match Env.parent_page env with
-      | None -> Error (`Lookup_by_name (`Asset, name))
-      | Some parent_page ->
-          Page.in_env_from_id env parent_page >>= fun (_, page) ->
-          A.in_page env page name)
+  | `Root (name, _) -> A.in_env env name
   | `Dot (parent, name) ->
       let x =
         resolve_label_parent_reference env parent >>= function
@@ -878,24 +878,21 @@ let resolve_reference =
     match r with
     | `Root (name, `TUnknown) -> (
         let identifier id = Ok (`Identifier (id :> Identifier.t)) in
-        match env_lookup_by_name Env.s_any name env with
-        | Ok (`Module (_, _) as e) -> resolved (M.of_element env e)
-        | Ok (`ModuleType (_, _) as e) -> resolved (MT.of_element env e)
-        | Ok (`Value (id, _)) -> identifier id
-        | Ok (`Type (id, _)) -> identifier id
-        | Ok (`Label (id, _)) -> identifier id
-        | Ok (`Class (id, _)) -> identifier id
-        | Ok (`ClassType (id, _)) -> identifier id
-        | Ok (`Constructor (id, _)) -> identifier id
-        | Ok (`Exception (id, _)) -> identifier id
-        | Ok (`Extension (id, _, _)) -> identifier id
-        | Ok (`ExtensionDecl (id, _)) -> identifier id
-        | Ok (`Field (id, _)) -> identifier id
-        | Ok (`Page (id, _)) -> identifier id
-        | Error _ as e -> (
-            match resolve_asset_reference env (`Root (name, `TAsset)) with
-            | Ok res -> resolved1 res
-            | Error _ -> e))
+        env_lookup_by_name Env.s_any name env >>= function
+        | `Module (_, _) as e -> resolved (M.of_element env e)
+        | `ModuleType (_, _) as e -> resolved (MT.of_element env e)
+        | `Value (id, _) -> identifier id
+        | `Type (id, _) -> identifier id
+        | `Label (id, _) -> identifier id
+        | `Class (id, _) -> identifier id
+        | `ClassType (id, _) -> identifier id
+        | `Constructor (id, _) -> identifier id
+        | `Exception (id, _) -> identifier id
+        | `Extension (id, _, _) -> identifier id
+        | `ExtensionDecl (id, _) -> identifier id
+        | `Field (id, _) -> identifier id
+        | `Page (id, _) -> identifier id
+        | `Asset id -> identifier id)
     | `Resolved r -> Ok r
     | `Root (name, (`TModule | `TChildModule)) -> M.in_env env name >>= resolved
     | `Module (parent, name) ->