[display] rewrite find references

Not quite done yet

src/compiler/displayOutput.ml

@@ -606,26 +606,7 @@ let process_display_file com classes =
 
 let process_global_display_mode com tctx = match com.display.dms_kind with
 	| DMUsage with_definition ->
-		let symbols,relations = Statistics.collect_statistics tctx in
-		let rec loop acc relations = match relations with
-			| (Statistics.Referenced,p) :: relations -> loop (p :: acc) relations
-			| _ :: relations -> loop acc relations
-			| [] -> acc
-		in
-		let usages = Hashtbl.fold (fun p sym acc ->
-			if p = !Display.reference_position then begin
-				let acc = if with_definition then p :: acc else acc in
-				(try loop acc (Hashtbl.find relations p)
-				with Not_found -> acc)
-			end else
-				acc
-		) symbols [] in
-		let usages = List.sort (fun p1 p2 ->
-			let c = compare p1.pfile p2.pfile in
-			if c <> 0 then c else compare p1.pmin p2.pmin
-		) usages in
-		Display.reference_position := null_pos;
-		raise_position usages
+		FindReferences.find_references tctx com with_definition
 	| DMDiagnostics global ->
 		let dctx = Diagnostics.prepare com global in
 		(* Option.may (fun cs -> CompilationServer.cache_context cs com) (CompilationServer.get()); *)

src/compiler/main.ml

@@ -833,6 +833,12 @@ try
 		List.iter (fun f -> f ()) (List.rev com.callbacks.after_init_macros);
 		List.iter (fun cpath -> ignore(tctx.Typecore.g.Typecore.do_load_module tctx cpath null_pos)) (List.rev !classes);
 		Finalization.finalize tctx;
+		(* If we are trying to find references, let's syntax-explore everything we know to check for the
+		   identifier we are interested in. We then type only those modules that contain the identifier. *)
+		begin match !CompilationServer.instance,com.display.dms_kind with
+			| Some cs,DMUsage _ -> FindReferences.find_possible_references tctx cs;
+			| _ -> ()
+		end;
 		t();
 		if not ctx.com.display.dms_display && ctx.has_error then raise Abort;
 		let load_display_module_in_macro clear = match display_file_dot_path with

src/context/compilationServer.ml

@@ -124,6 +124,14 @@ let get_file_list cs com =
 	let sign = Define.get_signature com.defines in
 	Hashtbl.fold (fun (file,sign') decls acc -> if sign = sign' then (file,decls) :: acc else acc) cs.cache.c_files []
 
+let get_module_name_of_cfile file cfile = match cfile.c_module_name with
+	| None ->
+		let name = Path.module_name_of_file file in
+		cfile.c_module_name <- Some name;
+		name
+	| Some name ->
+		name
+
 (* haxelibs *)
 
 let find_haxelib cs key =

src/context/display/display.ml

@@ -10,7 +10,7 @@ open Globals
 open Genjson
 open DisplayPosition
 
-let reference_position = ref null_pos
+let reference_position = ref ("",null_pos,KVar)
 
 module ExprPreprocessing = struct
 	let find_before_pos com is_completion e =

src/context/display/displayEmitter.ml

@@ -3,7 +3,8 @@ open Ast
 open Type
 open Typecore
 open DisplayException
-open DisplayTypes.DisplayMode
+open DisplayTypes
+open DisplayMode
 open CompletionItem
 open ClassFieldOrigin
 open DisplayTypes.CompletionResultKind
@@ -38,7 +39,9 @@ let patch_type ctx t =
 
 let display_module_type ctx mt p = match ctx.com.display.dms_kind with
 	| DMDefinition -> raise_position [(t_infos mt).mt_name_pos];
-	| DMUsage _ -> reference_position := (t_infos mt).mt_name_pos
+	| DMUsage _ ->
+		let infos = t_infos mt in
+		reference_position := (snd infos.mt_path,infos.mt_name_pos,KModuleType)
 	| DMHover ->
 		let t = patch_type ctx (type_of_module_type mt) in
 		raise_hover (make_ci_type (CompletionModuleType.of_module_type mt) ImportStatus.Imported (Some t)) p
@@ -66,14 +69,12 @@ let check_display_type ctx t p =
 		if ctx.is_display_file && encloses_display_position p then
 			display_type ctx t p
 	in
-	match ctx.com.display.dms_kind with
-	| DMStatistics -> add_type_hint()
-	| DMUsage _ -> add_type_hint(); maybe_display_type()
-	| _ -> maybe_display_type()
+	add_type_hint();
+	maybe_display_type()
 
 let display_variable ctx v p = match ctx.com.display.dms_kind with
 	| DMDefinition -> raise_position [v.v_pos]
-	| DMUsage _ -> reference_position := v.v_pos
+	| DMUsage _ -> reference_position := (v.v_name,v.v_pos,KVar)
 	| DMHover ->
 		let t = patch_type ctx v.v_type in
 		raise_hover (make_ci_local v t) p
@@ -81,7 +82,7 @@ let display_variable ctx v p = match ctx.com.display.dms_kind with
 
 let display_field ctx origin scope cf p = match ctx.com.display.dms_kind with
 	| DMDefinition -> raise_position [cf.cf_name_pos]
-	| DMUsage _ -> reference_position := cf.cf_name_pos
+	| DMUsage _ -> reference_position := (cf.cf_name,cf.cf_name_pos,KClassField)
 	| DMHover ->
 		let cf = if Meta.has Meta.Impl cf.cf_meta then
 			prepare_using_field cf
@@ -101,7 +102,7 @@ let maybe_display_field ctx origin scope cf p =
 
 let display_enum_field ctx en ef p = match ctx.com.display.dms_kind with
 	| DMDefinition -> raise_position [ef.ef_name_pos]
-	| DMUsage _ -> reference_position := ef.ef_name_pos
+	| DMUsage _ -> reference_position := (ef.ef_name,ef.ef_name_pos,KEnumField)
 	| DMHover ->
 		let t = patch_type ctx ef.ef_type in
 		raise_hover (make_ci_enum_field (CompletionEnumField.make ef (Self (TEnumDecl en)) true) t) p

src/context/display/displayToplevel.ml

@@ -347,14 +347,7 @@ let collect ctx epos with_type =
 		) files in
 		let files = List.sort (fun (_,i1) (_,i2) -> -compare i1 i2) files in
 		List.iter (fun ((file,cfile),_) ->
-			let module_name = match cfile.c_module_name with
-			| None ->
-				let name = Path.module_name_of_file file in
-				cfile.c_module_name <- Some name;
-				name
-			| Some name ->
-				name
-			in
+			let module_name = CompilationServer.get_module_name_of_cfile file cfile in
 			begin match List.rev cfile.c_package with
 				| [] -> ()
 				| s :: sl -> add_package (List.rev sl,s)

src/context/display/findReferences.ml

@@ -0,0 +1,190 @@
+open Globals
+open Ast
+open DisplayTypes
+open Common
+open Typecore
+open CompilationServer
+open ImportHandling
+
+let find_possible_references kind name (pack,decls) =
+	(* Employ some heuristics: We know what kind of symbol we are looking for, so let's
+	   filter where we can. *)
+	let check kind' name' =
+		if name = name' then match kind',kind with
+			| KIdent,_
+			| KAnyField,(KAnyField | KClassField | KEnumField)
+			| KClassField,KClassField
+			| KEnumField,KEnumField
+			| KModuleType,KModuleType
+			| KConstructor,(KConstructor | KModuleType) ->
+				raise Exit
+			| _ ->
+				()
+	in
+	let rec type_path kind path =
+		check KModuleType path.tname;
+		Option.may (check KModuleType) path.tsub;
+		List.iter (function
+			| TPType th -> type_hint th
+			| TPExpr e -> expr e
+		) path.tparams
+	and type_hint th = match fst th with
+		| CTPath path -> type_path KModuleType path
+		| CTParent th | CTOptional th | CTNamed(_,th) -> type_hint th
+		| CTFunction(thl,th) ->
+			List.iter type_hint thl;
+			type_hint th;
+		| CTAnonymous cffl ->
+			List.iter field cffl
+		| CTExtend(tl,cffl) ->
+			List.iter (fun (path,_) -> type_path KModuleType path) tl;
+			List.iter field cffl;
+	and type_param tp =
+		List.iter type_param tp.tp_params;
+		List.iter type_hint tp.tp_constraints
+	and expr (e,p) =
+		begin match e with
+		| EConst(Ident s) ->
+			check KIdent s
+		| EField(e1,s) ->
+			expr e1;
+			check KAnyField s;
+		| EVars vl ->
+			List.iter (fun (_,tho,eo) ->
+				Option.may type_hint tho;
+				expr_opt eo
+			) vl;
+		| ECast(e1,tho) ->
+			expr e1;
+			Option.may type_hint tho;
+		| ENew((path,_),el) ->
+			type_path KConstructor path;
+			List.iter expr el;
+		| EFunction(_,f) ->
+			func f
+		| ETry(e1,catches) ->
+			expr e1;
+			List.iter (fun (_,th,e,_) ->
+				type_hint th;
+				expr e
+			) catches;
+		| ECheckType(e1,th) ->
+			expr e1;
+			type_hint th;
+		| _ ->
+			iter_expr expr (e,p)
+		end
+	and expr_opt eo = match eo with
+		| None -> ()
+		| Some e -> expr e
+	and func f =
+		List.iter (fun ((s,p),_,_,tho,eo) ->
+			Option.may type_hint tho;
+			expr_opt eo
+		) f.f_args;
+		List.iter type_param f.f_params;
+		Option.may type_hint f.f_type;
+		expr_opt f.f_expr
+	and field cff =
+		check KClassField (fst cff.cff_name);
+		match cff.cff_kind with
+		| FVar(tho,eo) ->
+			Option.may type_hint tho;
+			expr_opt eo
+		| FFun f ->
+			func f
+		| FProp(_,_,tho,eo) ->
+			Option.may type_hint tho;
+			expr_opt eo
+	in
+	List.iter (fun (td,p) -> match td with
+		| EImport(path,_) | EUsing path ->
+			begin match fst (ImportHandling.convert_import_to_something_usable null_pos path) with
+			| IDKModule(_,s) -> check KModuleType s
+			| IDKSubType(_,s1,s2) ->
+				check KModuleType s1;
+				check KModuleType s2;
+			| IDKSubTypeField(_,s1,s2,s3) ->
+				check KModuleType s1;
+				check KModuleType s2;
+				check KAnyField s3;
+			| IDKModuleField(_,s1,s2) ->
+				check KModuleType s1;
+				check KAnyField s2;
+			| IDKPackage _ | IDK ->
+				()
+			end;
+		| EClass d ->
+			check KModuleType (fst d.d_name);
+			List.iter (function
+				| HExtends(path,_) | HImplements(path,_) -> type_path KModuleType path
+				| _ -> ()
+			) d.d_flags;
+			List.iter type_param d.d_params;
+			List.iter field d.d_data
+		| EEnum d ->
+			check KModuleType (fst d.d_name);
+			List.iter (fun ef ->
+				Option.may type_hint ef.ec_type;
+				check KEnumField (fst ef.ec_name);
+				List.iter type_param ef.ec_params;
+			) d.d_data;
+			List.iter type_param d.d_params;
+		| ETypedef d ->
+			check KModuleType (fst d.d_name);
+			List.iter type_param d.d_params;
+			type_hint d.d_data;
+		| EAbstract d ->
+			check KModuleType (fst d.d_name);
+			List.iter field d.d_data;
+			List.iter type_param d.d_params;
+			List.iter (function
+				| AbFrom th | AbTo th | AbOver th -> type_hint th
+				| _ -> ()
+			) d.d_flags;
+	) decls
+
+let find_possible_references tctx cs =
+	let name,pos,kind = !Display.reference_position in
+	if not (CompilationServer.is_initialized cs) then begin
+		CompilationServer.set_initialized cs;
+		DisplayToplevel.read_class_paths tctx.com ["display";"references"];
+	end;
+	let files = CompilationServer.get_file_list cs tctx.com in
+	let _ = List.iter (fun (file,cfile) ->
+		try
+			find_possible_references kind name (cfile.c_package,cfile.c_decls);
+		with Exit ->
+			let module_name = CompilationServer.get_module_name_of_cfile file cfile in
+			begin try
+				ignore(tctx.g.do_load_module tctx (cfile.c_package,module_name) null_pos);
+				(* We have to flush immediately so we catch exceptions from weird modules *)
+				Typecore.flush_pass tctx Typecore.PFinal "final";
+			with _ ->
+				()
+			end
+	) files in
+	()
+
+let find_references tctx com with_definition =
+	let name,pos,kind = !Display.reference_position in
+	let symbols,relations = Statistics.collect_statistics tctx in
+	let rec loop acc relations = match relations with
+		| (Statistics.Referenced,p) :: relations -> loop (p :: acc) relations
+		| _ :: relations -> loop acc relations
+		| [] -> acc
+	in
+	let usages = Hashtbl.fold (fun p sym acc ->
+		if p = pos then begin
+			let acc = if with_definition then p :: acc else acc in
+			(try loop acc (Hashtbl.find relations p)
+			with Not_found -> acc)
+		end else
+			acc
+	) symbols [] in
+	let usages = List.sort (fun p1 p2 ->
+		let c = compare p1.pfile p2.pfile in
+		if c <> 0 then c else compare p1.pmin p2.pmin
+	) usages in
+	Display.reference_position := ("",null_pos,KVar);
+	DisplayException.raise_position usages

src/context/display/statistics.ml

@@ -164,21 +164,11 @@ let collect_statistics ctx =
 			check_module a.a_module;
 			declare (SKAbstract a) a.a_name_pos
 	in
-	begin match CompilationServer.get () with
-		| None ->
-			let rec loop com =
-				List.iter f com.types;
-				Option.may loop (com.get_macros())
-			in
-			loop ctx.com
-		| Some cs ->
-			let rec loop com =
-				(* CompilationServer.cache_context cs com; *)
-				CompilationServer.iter_modules cs com (fun m -> List.iter f m.m_types);
-				Option.may loop (com.get_macros())
-			in
-			loop ctx.com
-	end;
+	let rec loop com =
+		List.iter f com.types;
+		Option.may loop (com.get_macros())
+	in
+	loop ctx.com;
 	let l = List.fold_left (fun acc (_,cfi,_,cfo) -> match cfo with
 		| Some cf -> if List.mem_assoc cf.cf_name_pos acc then acc else (cf.cf_name_pos,cfi.cf_name_pos) :: acc
 		| None -> acc

src/core/displayTypes.ml

@@ -236,4 +236,13 @@ module DisplayMode = struct
 		| DMDiagnostics b -> (if b then "global " else "") ^ "diagnostics"
 		| DMStatistics -> "statistics"
 		| DMSignature -> "signature"
-end
+end
+
+type reference_kind =
+	| KVar
+	| KIdent
+	| KAnyField
+	| KClassField
+	| KEnumField
+	| KModuleType
+	| KConstructor

src/typing/typerDisplay.ml

@@ -228,32 +228,32 @@ and display_expr ctx e_ast e dk with_type p =
 	| DMUsage _ ->
 		let rec loop e = match e.eexpr with
 		| TField(_,FEnum(_,ef)) ->
-			Display.reference_position := ef.ef_name_pos;
+			Display.reference_position := (ef.ef_name,ef.ef_name_pos,KEnumField);
 		| TField(_,(FAnon cf | FInstance (_,_,cf) | FStatic (_,cf) | FClosure (_,cf))) ->
-			Display.reference_position := cf.cf_name_pos;
+			Display.reference_position := (cf.cf_name,cf.cf_name_pos,KClassField);
 		| TLocal v | TVar(v,_) ->
-			Display.reference_position := v.v_pos;
+			Display.reference_position := (v.v_name,v.v_pos,KVar);
 		| TTypeExpr mt ->
 			let ti = t_infos mt in
-			Display.reference_position := ti.mt_name_pos;
+			Display.reference_position := (snd ti.mt_path,ti.mt_name_pos,KModuleType);
 		| TNew(c,tl,_) ->
 			begin try
 				let _,cf = get_constructor ctx c tl p in
-				Display.reference_position := cf.cf_name_pos;
+				Display.reference_position := (snd c.cl_path,cf.cf_name_pos,KConstructor);
 			with Not_found ->
 				()
 			end
 		| TCall({eexpr = TConst TSuper},_) ->
 			begin try
 				let cf = get_super_constructor() in
-				Display.reference_position := cf.cf_name_pos;
+				Display.reference_position := (cf.cf_name,cf.cf_name_pos,KClassField);
 			with Not_found ->
 				()
 			end
 		| TConst TSuper ->
 			begin match ctx.curclass.cl_super with
 				| None -> ()
-				| Some (c,_) -> Display.reference_position := c.cl_name_pos;
+				| Some (c,_) -> Display.reference_position := (snd c.cl_path,c.cl_name_pos,KModuleType);
 			end
 		| TCall(e1,_) ->
 			loop e1