Add support for .custom.hx handling (#12397)

Author: kLabz

src/compiler/args.ml

@@ -277,6 +277,9 @@ let parse_args com =
 			let pack, target = (try ExtString.String.split s ":" with _ -> raise (Arg.Bad "Invalid remap format, expected source:target")) in
 			com.package_rules <- PMap.add pack (Remap target) com.package_rules;
 		),"<package:target>","remap a package to another one");
+		("Compilation",["--custom-extension"],[],Arg.String (fun ext ->
+			com.custom_ext <- Some ext;
+		),"<extension>","custom extension used to shadow modules with Module.custom.hx files");
 		("Compilation",["--macro"],[], Arg.String (fun e ->
 			actx.force_typing <- true;
 			actx.config_macros <- e :: actx.config_macros

src/compiler/compiler.ml

@@ -79,7 +79,7 @@ let run_command ctx cmd =
 module Setup = struct
 	let initialize_target ctx com actx =
 		init_platform com;
-		com.class_paths#lock_context (platform_name com.platform) false;
+		com.class_paths#lock_context com.custom_ext (platform_name com.platform) false;
 		let add_std dir =
 			com.class_paths#modify_inplace (fun cp -> match cp#scope with
 				| Std ->

src/context/common.ml

@@ -282,6 +282,7 @@ type context = {
 	mutable doinline : bool;
 	mutable platform : platform;
 	mutable config : PlatformConfig.platform_config;
+	mutable custom_ext : string option;
 	empty_class_path : ClassPath.class_path;
 	class_paths : ClassPaths.class_paths;
 	main : Gctx.context_main;
@@ -746,6 +747,7 @@ let create timer_ctx compilation_step cs version args display_mode =
 		features = Hashtbl.create 0;
 		platform = Cross;
 		config = default_config;
+		custom_ext = None;
 		print = (fun s -> print_string s; flush stdout);
 		run_command = Sys.command;
 		run_command_args = (fun s args -> com.run_command (Printf.sprintf "%s %s" s (String.concat " " args)));
@@ -854,6 +856,7 @@ let clone com is_macro_context =
 		doinline = com.doinline;
 		platform = com.platform;
 		config = com.config;
+		custom_ext = com.custom_ext;
 		print = com.print;
 		run_command = com.run_command;
 		run_command_args = com.run_command_args;

src/context/display/displayToplevel.ml

@@ -89,7 +89,8 @@ class explore_class_path_task com checked recursive f_pack f_module dir pack = o
 									let name = String.sub file 0 (l - 3) in
 									try
 										let dot_pos = String.rindex name '.' in
-										if platform_str = String.sub file dot_pos (String.length name - dot_pos) then
+										let second_ext = String.sub file dot_pos (String.length name - dot_pos) in
+										if (Option.map_default (fun custom_ext -> custom_ext = second_ext) false com.custom_ext) || platform_str = second_ext then
 											String.sub file 0 dot_pos
 										else
 											raise Exit

src/context/lookup.ml

@@ -3,6 +3,7 @@ class virtual ['key,'value] lookup = object(self)
 	method virtual add : 'key -> 'value -> unit
 	method virtual remove : 'key -> unit
 	method virtual find : 'key -> 'value
+	method virtual find_opt : 'key -> 'value option
 	method virtual iter : ('key -> 'value -> unit) -> unit
 	method virtual fold : 'acc . ('key -> 'value -> 'acc -> 'acc) -> 'acc -> 'acc
 	method virtual mem : 'key -> bool
@@ -22,6 +23,9 @@ class ['key,'value] pmap_lookup = object(self)
 	method find (key : 'key) : 'value =
 		PMap.find key lut
 
+	method find_opt (key : 'key) : 'value option =
+		try Some (PMap.find key lut) with Not_found -> None
+
 	method iter (f : 'key -> 'value -> unit) =
 		PMap.iter f lut
 
@@ -48,6 +52,9 @@ class ['key,'value] hashtbl_lookup = object(self)
 	method find (key : 'key) : 'value =
 		Hashtbl.find lut key
 
+	method find_opt (key : 'key) : 'value option =
+		Hashtbl.find_opt lut key
+
 	method iter (f : 'key -> 'value -> unit) =
 		Hashtbl.iter f lut
 

src/core/classPaths.ml

@@ -11,16 +11,25 @@ let create_resolved_file file class_path = {
 	class_path;
 }
 
+type file_resolution_specificity =
+	| SpecificityNormal            (* Standard Module.hx file *)
+	| SpecificityPlatformSpecific  (* Module.[platform].hx file matching current platform *)
+	| SpecificityCustomExtension   (* Module.[custom].hx file matching --custom-extension config *)
+	| SpecificityMacroSpecific     (* Module.macro.hx file while in macro context *)
+	| SpecificityCoreApi           (* Module.hx takes priority when loading @:coreApi types *)
+
 (* We need to clean-up absolute ("") vs. cwd ("."). *)
 let absolute_class_path = new directory_class_path "" User
 
 class class_paths = object(self)
 	val mutable l = []
 	val file_lookup_cache = new Lookup.hashtbl_lookup;
+	val mutable custom_ext = None
 	val mutable platform_ext = ""
 	val mutable is_loading_core_api = false
 
-	method lock_context (platform_name : string) (core_api : bool) : unit =
+	method lock_context (custom_extension : string option) (platform_name : string) (core_api : bool) : unit =
+		custom_ext <- Option.map (fun ext -> "." ^ ext) custom_extension;
 		platform_ext <- "." ^ platform_name;
 		is_loading_core_api <- core_api;
 		self#clear_cache
@@ -76,59 +85,60 @@ class class_paths = object(self)
 		(*
 			This function is invoked for each file in the `dir`.
 			Each file is checked if it's specific for current platform
-			(e.g. ends with `.js.hx` while compiling for JS).
-			If it's not platform-specific:
-				Check the lookup cache and if the file is not there store full file path in the cache.
-			If the file is platform-specific:
-				Store the full file path in the lookup cache probably replacing the cached path to a
-				non-platform-specific file.
+			(e.g. ends with `.js.hx` while compiling for JS) or current
+			custom extension (e.g. ends with `.custom.hx` while compiling
+			with `--custom-extension custom`)
+
+			The lookup cache will store the full file path which is the more
+			specific in current context (see `file_resolution_specificity` type)
 		*)
 		let found = ref None in
 		let f_dir = Filename.dirname f_search in
 		let prepare_file file_own_name =
 			let relative_to_classpath = if f_dir = "." then file_own_name else f_dir ^ "/" ^ file_own_name in
 			(* `representation` is how the file is referenced to. E.g. when it's deduced from a module path. *)
-			let is_platform_specific,representation =
-				(* Platform specific file extensions are not allowed for loading @:coreApi types. *)
+			let specificity,representation =
 				if is_loading_core_api then
-					false,relative_to_classpath
+					SpecificityCoreApi,relative_to_classpath
 				else begin
 					let ext = extension relative_to_classpath in
 					let second_ext = extension (remove_extension relative_to_classpath) in
+					(* The file contains double extension and the secondary one matches current custom extension *)
+					if (Option.map_default (fun custom_ext -> custom_ext = second_ext) false custom_ext) then
+						SpecificityCustomExtension,(remove_extension (remove_extension relative_to_classpath)) ^ ext
+					(* The file contains ".macro.hx" double extension and we are in macro context *)
+					else if platform_ext = second_ext && second_ext = ".macro" then
+						SpecificityMacroSpecific,(remove_extension (remove_extension relative_to_classpath)) ^ ext
 					(* The file contains double extension and the secondary one matches current platform *)
-					if platform_ext = second_ext then
-						true,(remove_extension (remove_extension relative_to_classpath)) ^ ext
+					else if platform_ext = second_ext then
+						SpecificityPlatformSpecific,(remove_extension (remove_extension relative_to_classpath)) ^ ext
 					else
-						false,relative_to_classpath
+						SpecificityNormal,relative_to_classpath
 				end
 			in
-			(*
-				Store current full path for `representation` if
-				- we're loading @:coreApi
-				- or this is a platform-specific file for `representation`
-				- this `representation` was never found before
-			*)
-			if is_loading_core_api || is_platform_specific || not (file_lookup_cache#mem representation) then begin
-				let full_path = if dir = "." then file_own_name else dir ^ "/" ^ file_own_name in
-				let full_path = Some(create_resolved_file full_path cp) in
-				file_lookup_cache#add representation full_path;
-				if representation = f_search then found := full_path
-			end
+
+			let full_path = if dir = "." then file_own_name else dir ^ "/" ^ file_own_name in
+			let full_path = Some(create_resolved_file full_path cp, specificity) in
+
+			match file_lookup_cache#find_opt representation with
+			| Some (Some (_, old_specificity)) when (old_specificity >= specificity)-> ()
+			| _ -> file_lookup_cache#add representation full_path;
+
+			if representation = f_search then
+				match !found with
+				| Some (_, old_specificity) when (old_specificity >= specificity) -> ()
+				| _ -> found := full_path;
 		in
 		Array.iter prepare_file dir_listing;
 		!found
 
 	method find_file_noraise (f : string) =
 		try
-			match file_lookup_cache#find f with
-			| None ->
-				None
-			| Some f ->
-				Some f
+			file_lookup_cache#find f
 		with
 		| Not_found when Path.is_absolute_path f ->
 			let r = if Sys.file_exists f then
-				Some (create_resolved_file f absolute_class_path)
+				Some (create_resolved_file f absolute_class_path, SpecificityNormal)
 			else
 				None
 			in
@@ -144,8 +154,8 @@ class class_paths = object(self)
 							loop l
 						| Some(dir,dir_listing) ->
 							match self#cache_directory cp dir f dir_listing with
-								| Some f ->
-									Some f
+								| Some (f, specificity) ->
+									Some (f, specificity)
 								| None ->
 									loop l
 					end
@@ -157,7 +167,7 @@ class class_paths = object(self)
 	method find_file (f : string) =
 		match self#find_file_noraise f with
 		| None -> raise Not_found
-		| Some f -> f
+		| Some (f, _) -> f
 
 	method relative_path file =
 		let slashes path = String.concat "/" (ExtString.String.nsplit path "\\") in

src/typing/macroContext.ml

@@ -716,7 +716,7 @@ let create_macro_context com =
 	com2.package_rules <- PMap.empty;
 	(* Inherit most display settings, but require normal typing. *)
 	com2.display <- {com.display with dms_kind = DMNone; dms_full_typing = true; dms_force_macro_typing = true; dms_inline = true; };
-	com2.class_paths#lock_context "macro" false;
+	com2.class_paths#lock_context com2.custom_ext "macro" false;
 	let name = platform_name Eval in
 	let eval_std = ref None in
 	com2.class_paths#modify (fun cp -> match cp#scope with

src/typing/typeload.ml

@@ -799,7 +799,7 @@ let load_core_class ctx c =
 			allow_package com2 "sys";
 			Define.raw_define_value com2.defines "target.threaded" "true"; (* hack because we check this in sys.thread classes *)
 			if ctx.com.is_macro_context then Common.define com2 Define.Macro;
-			com2.class_paths#lock_context (platform_name_macro ctx.com) true;
+			com2.class_paths#lock_context ctx.com.custom_ext (platform_name_macro ctx.com) true;
 			com2.class_paths#modify (fun cp -> match cp#scope with
 				| Std ->
 					[cp#clone]

tests/misc/customExt/Bar.api.hx

@@ -0,0 +1,8 @@
+class Bar {
+	public static function bar() {
+		trace("Hello from Bar.hx");
+		test();
+	}
+
+	static macro function test() return macro trace("Hello from Bar.hx macro");
+}

tests/misc/customExt/Bar.hx

@@ -0,0 +1,5 @@
+class Bar {
+	public static function bar() {
+		trace("Hello from Bar.hx");
+	}
+}