[hxb] POC: lazy type restoring

Try reducing the amount of modules restored during display requests

src/compiler/hxb/hxbReader.ml

@@ -21,11 +21,19 @@ let create_field_reader_context p ts vars tthis = {
 type hxb_reader_stats = {
 	modules_fully_restored : int ref;
 	modules_partially_restored : int ref;
+
+	full_restore : int ref;
+	display_file_restore : int ref;
+	minimal_restore : int ref;
 }
 
 let create_hxb_reader_stats () = {
 	modules_fully_restored = ref 0;
 	modules_partially_restored = ref 0;
+
+	full_restore = ref 0;
+	minimal_restore = ref 0;
+	display_file_restore = ref 0;
 }
 
 module ClassFieldInfo = struct
@@ -144,6 +152,14 @@ let dump_stats name stats =
 	print_endline (Printf.sprintf "hxb_reader stats for %s" name);
 	print_endline (Printf.sprintf "  modules partially restored: %i" (!(stats.modules_partially_restored) - !(stats.modules_fully_restored)));
 	print_endline (Printf.sprintf "  modules fully restored: %i" !(stats.modules_fully_restored));
+	print_endline (Printf.sprintf "  full restore: %i" !(stats.full_restore));
+	print_endline (Printf.sprintf "  minimal restore: %i" !(stats.minimal_restore));
+	print_endline (Printf.sprintf "  display file restore: %i" !(stats.display_file_restore));
+
+type restore_level =
+	| Minimal
+	| DisplayFile
+	| Full
 
 class hxb_reader
 	(mpath : path)
@@ -152,7 +168,7 @@ class hxb_reader
 	(timers_enabled : bool)
 = object(self)
 	val mutable api = Obj.magic ""
-	val mutable minimal_restore = false
+	val mutable restore_level = Full
 	val mutable current_module = null_module
 
 	val mutable ch = BytesWithPosition.create (Bytes.create 0)
@@ -161,9 +177,9 @@ class hxb_reader
 	val mutable doc_pool = Array.make 0 ""
 
 	val mutable classes = Array.make 0 null_class
-	val mutable abstracts = Array.make 0 null_abstract
-	val mutable enums = Array.make 0 null_enum
-	val mutable typedefs = Array.make 0 null_typedef
+	val mutable abstracts = Array.make 0 (Lazy.from_val null_abstract)
+	val mutable enums = Array.make 0 (Lazy.from_val null_enum)
+	val mutable typedefs = Array.make 0 (Lazy.from_val null_typedef)
 	val mutable anons = Array.make 0 null_tanon
 	val mutable anon_fields = Array.make 0 null_field
 	val mutable tmonos = Array.make 0 (mk_mono())
@@ -709,6 +725,9 @@ class hxb_reader
 
 	(* Type instances *)
 
+	method should_lazy_wrap =
+		restore_level = Minimal
+
 	method resolve_ttp_ref = function
 		| 1 ->
 			let i = read_uleb128 ch in
@@ -752,10 +771,20 @@ class hxb_reader
 			c.cl_type
 		| 11 ->
 			let en = self#read_enum_ref in
-			en.e_type
+			if self#should_lazy_wrap then
+				TLazy (ref (LWait (fun () ->
+					(Lazy.force en).e_type
+				)))
+			else
+				(Lazy.force en).e_type
 		| 12 ->
 			let a = self#read_abstract_ref in
-			TType(abstract_module_type a [],[])
+			if self#should_lazy_wrap then
+				TLazy (ref (LWait (fun () ->
+					TType(abstract_module_type (Lazy.force a) [],[])
+				)))
+			else
+				TType(abstract_module_type (Lazy.force a) [],[])
 		| 13 ->
 			let e = self#read_expr in
 			let c = {null_class with cl_kind = KExpr e; cl_module = current_module } in
@@ -830,52 +859,112 @@ class hxb_reader
 			TInst(c,tl)
 		| 50 ->
 			let en = self#read_enum_ref in
-			TEnum(en,[])
+			if self#should_lazy_wrap then
+				TLazy (ref (LWait (fun () ->
+					TEnum(Lazy.force en,[])
+				)))
+			else
+				TEnum(Lazy.force en,[])
 		| 51 ->
 			let en = self#read_enum_ref in
 			let t1 = self#read_type_instance in
-			TEnum(en,[t1])
+			if self#should_lazy_wrap then
+				TLazy (ref (LWait (fun () ->
+					TEnum(Lazy.force en,[t1])
+				)))
+			else
+				TEnum(Lazy.force en,[t1])
 		| 52 ->
 			let en = self#read_enum_ref in
 			let t1 = self#read_type_instance in
 			let t2 = self#read_type_instance in
-			TEnum(en,[t1;t2])
+			if self#should_lazy_wrap then
+				TLazy (ref (LWait (fun () ->
+					TEnum(Lazy.force en,[t1;t2])
+				)))
+			else
+				TEnum(Lazy.force en,[t1;t2])
 		| 59 ->
 			let e = self#read_enum_ref in
 			let tl = self#read_types in
-			TEnum(e,tl)
+			if self#should_lazy_wrap then
+				TLazy (ref (LWait (fun () ->
+					TEnum(Lazy.force e,tl)
+				)))
+			else
+				TEnum(Lazy.force e,tl)
 		| 60 ->
 			let td = self#read_typedef_ref in
-			TType(td,[])
+			if self#should_lazy_wrap then
+				TLazy (ref (LWait (fun () ->
+					TType(Lazy.force td,[])
+				)))
+			else
+				TType(Lazy.force td,[])
 		| 61 ->
 			let td = self#read_typedef_ref in
 			let t1 = self#read_type_instance in
-			TType(td,[t1])
+			if self#should_lazy_wrap then
+				TLazy (ref (LWait (fun () ->
+					TType(Lazy.force td,[t1])
+				)))
+			else
+				TType(Lazy.force td,[t1])
 		| 62 ->
 			let td = self#read_typedef_ref in
 			let t1 = self#read_type_instance in
 			let t2 = self#read_type_instance in
-			TType(td,[t1;t2])
+			if self#should_lazy_wrap then
+				TLazy (ref (LWait (fun () ->
+					TType(Lazy.force td,[t1;t2])
+				)))
+			else
+				TType(Lazy.force td,[t1;t2])
 		| 69 ->
 			let t = self#read_typedef_ref in
 			let tl = self#read_types in
-			TType(t,tl)
+			if self#should_lazy_wrap then
+				TLazy (ref (LWait (fun () ->
+					TType(Lazy.force t,tl)
+				)))
+			else
+				TType(Lazy.force t,tl)
 		| 70 ->
 			let a = self#read_abstract_ref in
-			TAbstract(a,[])
+			if self#should_lazy_wrap then
+				TLazy (ref (LWait (fun () ->
+					TAbstract(Lazy.force a,[])
+				)))
+			else
+				TAbstract(Lazy.force a,[])
 		| 71 ->
 			let a = self#read_abstract_ref in
 			let t1 = self#read_type_instance in
-			TAbstract(a,[t1])
+			if self#should_lazy_wrap then
+				TLazy (ref (LWait (fun () ->
+					TAbstract(Lazy.force a,[t1])
+				)))
+			else
+				TAbstract(Lazy.force a,[t1])
 		| 72 ->
 			let a = self#read_abstract_ref in
 			let t1 = self#read_type_instance in
 			let t2 = self#read_type_instance in
-			TAbstract(a,[t1;t2])
+			if self#should_lazy_wrap then
+				TLazy (ref (LWait (fun () ->
+					TAbstract(Lazy.force a,[t1;t2])
+				)))
+			else
+				TAbstract(Lazy.force a,[t1;t2])
 		| 79 ->
 			let a = self#read_abstract_ref in
 			let tl = self#read_types in
-			TAbstract(a,tl)
+			if self#should_lazy_wrap then
+				TLazy (ref (LWait (fun () ->
+					TAbstract(Lazy.force a,tl)
+				)))
+			else
+				TAbstract(Lazy.force a,tl)
 		| 80 ->
 			empty_anon
 		| 81 ->
@@ -1244,6 +1333,7 @@ class hxb_reader
 						let e1 = loop () in
 						let en = self#read_enum_ref in
 						let ef = self#read_enum_field_ref in
+						let en = Lazy.force en in
 						TField(e1,FEnum(en,ef)),None
 					| 108 ->
 						let e1 = loop () in
@@ -1270,11 +1360,12 @@ class hxb_reader
 						TTypeExpr (TClassDecl c),(Some c.cl_type)
 					| 121 ->
 						let en = self#read_enum_ref in
+						let en = Lazy.force en in
 						TTypeExpr (TEnumDecl en),(Some en.e_type)
 					| 122 ->
-						TTypeExpr (TAbstractDecl self#read_abstract_ref),None
+						TTypeExpr (TAbstractDecl (Lazy.force self#read_abstract_ref)),None
 					| 123 ->
-						TTypeExpr (TTypeDecl self#read_typedef_ref),None
+						TTypeExpr (TTypeDecl (Lazy.force self#read_typedef_ref)),None
 					| 124 ->
 						TCast(loop (),None),None
 					| 125 ->
@@ -1498,7 +1589,7 @@ class hxb_reader
 			KGenericInstance(c,tl)
 		| 5 -> KMacroType
 		| 6 -> KGenericBuild (self#read_list (fun () -> self#read_cfield))
-		| 7 -> KAbstractImpl self#read_abstract_ref
+		| 7 -> KAbstractImpl (Lazy.force self#read_abstract_ref)
 		| 8 -> KModuleFields current_module
 		| i ->
 			error (Printf.sprintf "Invalid class kind id: %i" i)
@@ -1600,6 +1691,7 @@ class hxb_reader
 		let a = Array.init l (fun i ->
 			let en = self#read_enum_ref in
 			let name = self#read_string in
+			let en = Lazy.force en in
 			PMap.find name en.e_constrs
 		) in
 		enum_fields <- a
@@ -1723,7 +1815,7 @@ class hxb_reader
 	method read_afd =
 		let l = read_uleb128 ch in
 		for i = 0 to l - 1 do
-			let a = abstracts.(i) in
+			let a = Lazy.force abstracts.(i) in
 			self#read_abstract_fields a;
 		done
 
@@ -1737,21 +1829,21 @@ class hxb_reader
 	method read_abd =
 		let l = read_uleb128 ch in
 		for i = 0 to l - 1 do
-			let a = abstracts.(i) in
+			let a = Lazy.force abstracts.(i) in
 			self#read_abstract a;
 		done
 
 	method read_end =
 		let l = read_uleb128 ch in
 		for i = 0 to l - 1 do
-			let en = enums.(i) in
+			let en = Lazy.force enums.(i) in
 			self#read_enum en;
 		done
 
 	method read_efd =
 		let l = read_uleb128 ch in
 		for i = 0 to l - 1 do
-			let e = enums.(i) in
+			let e = Lazy.force enums.(i) in
 			self#read_enum_fields e;
 			Type.unify (TType(enum_module_type e,[])) e.e_type
 		done
@@ -1785,7 +1877,7 @@ class hxb_reader
 	method read_tdd =
 		let l = read_uleb128 ch in
 		for i = 0 to l - 1 do
-			let t = typedefs.(i) in
+			let t = Lazy.force typedefs.(i) in
 			self#read_typedef t;
 		done
 
@@ -1804,33 +1896,39 @@ class hxb_reader
 		let l = read_uleb128 ch in
 		abstracts <- (Array.init l (fun i ->
 			let (pack,mname,tname) = self#read_full_path in
-			match self#resolve_type pack mname tname with
-			| TAbstractDecl a ->
-				a
-			| _ ->
-				error ("Unexpected type where abstract was expected: " ^ (s_type_path (pack,tname)))
+			Lazy.from_fun (fun () ->
+				match self#resolve_type pack mname tname with
+				| TAbstractDecl a ->
+					a
+				| _ ->
+					error ("Unexpected type where abstract was expected: " ^ (s_type_path (pack,tname)))
+			)
 		))
 
 	method read_enr =
 		let l = read_uleb128 ch in
 		enums <- (Array.init l (fun i ->
 			let (pack,mname,tname) = self#read_full_path in
-			match self#resolve_type pack mname tname with
-			| TEnumDecl en ->
-				en
-			| _ ->
-				error ("Unexpected type where enum was expected: " ^ (s_type_path (pack,tname)))
+			Lazy.from_fun (fun () ->
+				match self#resolve_type pack mname tname with
+				| TEnumDecl en ->
+					en
+				| _ ->
+					error ("Unexpected type where enum was expected: " ^ (s_type_path (pack,tname)))
+			)
 		))
 
 	method read_tdr =
 		let l = read_uleb128 ch in
 		typedefs <- (Array.init l (fun i ->
 			let (pack,mname,tname) = self#read_full_path in
-			match self#resolve_type pack mname tname with
-			| TTypeDecl tpd ->
-				tpd
-			| _ ->
-				error ("Unexpected type where typedef was expected: " ^ (s_type_path (pack,tname)))
+			Lazy.from_fun (fun () ->
+				match self#resolve_type pack mname tname with
+				| TTypeDecl tpd ->
+					tpd
+				| _ ->
+					error ("Unexpected type where typedef was expected: " ^ (s_type_path (pack,tname)))
+			)
 		))
 
 	method read_imports =
@@ -1908,12 +2006,12 @@ class hxb_reader
 			| 2 ->
 				let td = mk_typedef current_module path pos name_pos (mk_mono()) in
 				td.t_params <- Array.to_list params;
-				typedefs <- Array.append typedefs (Array.make 1 td);
+				typedefs <- Array.append typedefs (Array.make 1 (Lazy.from_val td));
 				TTypeDecl td
 			| 3 ->
 				let a = mk_abstract current_module path pos name_pos in
 				a.a_params <- Array.to_list params;
-				abstracts <- Array.append abstracts (Array.make 1 a);
+				abstracts <- Array.append abstracts (Array.make 1 (Lazy.from_val a));
 				TAbstractDecl a
 			| _ ->
 				error ("Invalid type kind: " ^ (string_of_int kind));
@@ -1949,8 +2047,14 @@ class hxb_reader
 		| MTF ->
 			current_module.m_types <- self#read_mtf;
 			api#add_module current_module;
+			incr stats.modules_partially_restored;
+			incr (match restore_level with
+				| Full -> stats.full_restore
+				| DisplayFile -> stats.display_file_restore
+				| Minimal -> stats.minimal_restore
+			);
 		| IMP ->
-			if not minimal_restore then self#read_imports;
+			if restore_level = Full then self#read_imports;
 		| CLR ->
 			self#read_clr;
 		| ENR ->
@@ -2018,11 +2122,11 @@ class hxb_reader
 		close()
 
 	method read_chunks (new_api : hxb_reader_api) (chunks : cached_chunks) =
-		fst (self#read_chunks_until new_api chunks EOM false)
+		fst (self#read_chunks_until new_api chunks EOM Full)
 
-	method read_chunks_until (new_api : hxb_reader_api) (chunks : cached_chunks) end_chunk minimal_restore' =
+	method read_chunks_until (new_api : hxb_reader_api) (chunks : cached_chunks) end_chunk level =
 		api <- new_api;
-		minimal_restore <- minimal_restore';
+		restore_level <- level;
 		let rec loop = function
 			| (kind,data) :: chunks ->
 				ch <- BytesWithPosition.create data;
@@ -2035,7 +2139,7 @@ class hxb_reader
 
 	method read (new_api : hxb_reader_api) (bytes : bytes) =
 		api <- new_api;
-		minimal_restore <- false;
+		restore_level <- Full;
 		ch <- BytesWithPosition.create bytes;
 		if (Bytes.to_string (read_bytes ch 3)) <> "hxb" then
 			raise (HxbFailure "magic");

src/compiler/server.ml

@@ -311,8 +311,7 @@ let check_module sctx com m_path m_extra p =
 		in
 		let check_dependencies () =
 			let full_restore =
-				com.is_macro_context
-				|| com.display.dms_full_typing
+				com.display.dms_full_typing
 				|| DisplayPosition.display_position#is_in_file (Path.UniqueKey.lazy_key m_extra.m_file)
 			in
 			PMap.iter (fun _ mdep ->
@@ -394,7 +393,7 @@ let check_module sctx com m_path m_extra p =
 let get_hxb_module com cc path =
 	try
 		let mc = cc#get_hxb_module path in
-		if not com.is_macro_context && not com.display.dms_full_typing && not (DisplayPosition.display_position#is_in_file (Path.UniqueKey.lazy_key mc.mc_extra.m_file)) then begin
+		if not com.display.dms_full_typing && not (DisplayPosition.display_position#is_in_file (Path.UniqueKey.lazy_key mc.mc_extra.m_file)) then begin
 			mc.mc_extra.m_cache_state <- MSGood;
 			BinaryModule mc
 		end else
@@ -437,11 +436,14 @@ class hxb_reader_api_server
 			m
 		| BinaryModule mc ->
 			let reader = new HxbReader.hxb_reader path com.hxb_reader_stats (Some cc#get_string_pool_arr) (Common.defined com Define.HxbTimes) in
-			let is_display_file = DisplayPosition.display_position#is_in_file (Path.UniqueKey.lazy_key mc.mc_extra.m_file) in
-			let full_restore = com.is_macro_context || com.display.dms_full_typing || is_display_file in
+			let restore_level:HxbReader.restore_level =
+				if com.display.dms_full_typing then Full
+				else if DisplayPosition.display_position#is_in_file (Path.UniqueKey.lazy_key mc.mc_extra.m_file) then DisplayFile
+				else Minimal
+			in
 			let f_next chunks until =
 				let t_hxb = Timer.timer ["server";"module cache";"hxb read";"until " ^ (string_of_chunk_kind until)] in
-				let r = reader#read_chunks_until (self :> HxbReaderApi.hxb_reader_api) chunks until (not full_restore) in
+				let r = reader#read_chunks_until (self :> HxbReaderApi.hxb_reader_api) chunks until restore_level in
 				t_hxb();
 				r
 			in
@@ -449,7 +451,7 @@ class hxb_reader_api_server
 
 			(* We try to avoid reading expressions as much as possible, so we only do this for
 				 our current display file if we're in display mode. *)
-			if full_restore then ignore(f_next chunks EOM)
+			if restore_level <> Minimal then ignore(f_next chunks EOM)
 			else delay (fun () -> ignore(f_next chunks EOF));
 			m
 		| BadModule reason ->
@@ -504,28 +506,28 @@ let rec add_modules sctx com delay (m : module_def) (from_binary : bool) (p : po
 					com.module_lut#add m.m_path m;
 				handle_cache_bound_objects com m.m_extra.m_cache_bound_objects;
 				let full_restore =
-					com.is_macro_context
-					|| com.display.dms_full_typing
+					com.display.dms_full_typing
 					|| DisplayPosition.display_position#is_in_file (Path.UniqueKey.lazy_key m.m_extra.m_file)
 				in
 				PMap.iter (fun _ mdep ->
 					let mpath = mdep.md_path in
 					if mdep.md_sign = own_sign then begin
-						let m2 = try
-							com.module_lut#find mpath
-						with Not_found ->
-							match type_module sctx com delay mpath p with
-							| GoodModule m ->
-								m
-							| BinaryModule mc ->
-								failwith (Printf.sprintf "Unexpectedly found unresolved binary module %s as a dependency of %s" (s_type_path mpath) (s_type_path m0.m_path))
-							| NoModule ->
-								failwith (Printf.sprintf "Unexpectedly could not find module %s as a dependency of %s" (s_type_path mpath) (s_type_path m0.m_path))
-							| BadModule reason ->
-								failwith (Printf.sprintf "Unexpected bad module %s (%s) as a dependency of %s" (s_type_path mpath) (Printer.s_module_skip_reason reason) (s_type_path m0.m_path))
-						in
-						add_modules (tabs ^ "  ") m0 m2
+						try begin
+							let m2 = com.module_lut#find mpath in
+							add_modules (tabs ^ "  ") m0 m2
+						end with Not_found ->
+							if full_restore then
+								match type_module sctx com delay mpath p with
+								| GoodModule m ->
+									add_modules (tabs ^ "  ") m0 m
+								| BinaryModule mc ->
+									failwith (Printf.sprintf "Unexpectedly found unresolved binary module %s as a dependency of %s" (s_type_path mpath) (s_type_path m0.m_path))
+								| NoModule ->
+									failwith (Printf.sprintf "Unexpectedly could not find module %s as a dependency of %s" (s_type_path mpath) (s_type_path m0.m_path))
+								| BadModule reason ->
+									failwith (Printf.sprintf "Unexpected bad module %s (%s) as a dependency of %s" (s_type_path mpath) (Printer.s_module_skip_reason reason) (s_type_path m0.m_path))
 					end
+				(* ) (if full_restore then m.m_extra.m_deps else PMap.empty) *)
 				) (if full_restore then m.m_extra.m_deps else Option.default m.m_extra.m_deps m.m_extra.m_sig_deps)
 			)
 		end
@@ -579,8 +581,11 @@ and type_module sctx com delay mpath p =
 			begin match check_module sctx mpath mc.mc_extra p with
 				| None ->
 					let reader = new HxbReader.hxb_reader mpath com.hxb_reader_stats (Some cc#get_string_pool_arr) (Common.defined com Define.HxbTimes) in
-					let is_display_file = DisplayPosition.display_position#is_in_file (Path.UniqueKey.lazy_key mc.mc_extra.m_file) in
-					let full_restore = com.is_macro_context || com.display.dms_full_typing || is_display_file in
+					let restore_level:HxbReader.restore_level =
+						if com.display.dms_full_typing then Full
+						else if DisplayPosition.display_position#is_in_file (Path.UniqueKey.lazy_key mc.mc_extra.m_file) then DisplayFile
+						else Minimal
+					in
 					let api = match com.hxb_reader_api with
 						| Some api ->
 							api
@@ -591,14 +596,14 @@ and type_module sctx com delay mpath p =
 					in
 					let f_next chunks until =
 						let t_hxb = Timer.timer ["server";"module cache";"hxb read";"until " ^ (string_of_chunk_kind until)] in
-						let r = reader#read_chunks_until api chunks until (not full_restore) in
+						let r = reader#read_chunks_until api chunks until restore_level in
 						t_hxb();
 						r
 					in
 					let m,chunks = f_next mc.mc_chunks EOT in
 					(* We try to avoid reading expressions as much as possible, so we only do this for
 					   our current display file if we're in display mode. *)
-					if full_restore then ignore(f_next chunks EOM)
+					if restore_level <> Minimal then ignore(f_next chunks EOM)
 					else delay (fun () -> ignore(f_next chunks EOF));
 					add_modules true m;
 				| Some reason ->
@@ -641,6 +646,8 @@ let after_save sctx ctx =
 
 let after_compilation sctx ctx =
 	sctx.cs#clear_temp_cache;
+	if Define.raw_defined ctx.com.defines "hxb.stats" then
+		HxbReader.dump_stats (platform_name ctx.com.platform) ctx.com.hxb_reader_stats;
 	()
 
 let mk_length_prefixed_communication allow_nonblock chin chout =

src/context/display/displayJson.ml

@@ -140,7 +140,7 @@ class hxb_reader_api_com
 		with Not_found ->
 			let mc = cc#get_hxb_module m_path in
 			let reader = new HxbReader.hxb_reader mc.mc_path com.hxb_reader_stats (Some cc#get_string_pool_arr) (Common.defined com Define.HxbTimes) in
-			fst (reader#read_chunks_until (self :> HxbReaderApi.hxb_reader_api) mc.mc_chunks (if minimal_restore then MTF else EOM) minimal_restore)
+			fst (reader#read_chunks_until (self :> HxbReaderApi.hxb_reader_api) mc.mc_chunks (if minimal_restore then MTF else EOM) Minimal)
 
 	method basic_types =
 		com.basic
@@ -153,6 +153,7 @@ class hxb_reader_api_com
 end
 
 let find_module ~(minimal_restore : bool) com cc path =
+	(* TODO: check all this... *)
 	(new hxb_reader_api_com ~minimal_restore com cc)#find_module path
 
 type handler_context = {