From aff2d0766d93dd9b136ee020db8f3a68476ec27c Mon Sep 17 00:00:00 2001
From: lunarcleint <88809950+lunarcleint@users.noreply.github.com>
Date: Thu, 25 Apr 2024 18:20:17 -0700
Subject: [PATCH] clean up EImport

---
 hscript/Interp.hx            | 74 ++++++++++++++++++++----------------
 tests/src/tests/FinalCase.hx |  8 ++++
 2 files changed, 49 insertions(+), 33 deletions(-)

diff --git a/hscript/Interp.hx b/hscript/Interp.hx
index c2efb6f..0131f2d 100644
--- a/hscript/Interp.hx
+++ b/hscript/Interp.hx
@@ -565,68 +565,76 @@ class Interp {
 				}
 				customClasses.set(name, new CustomClassHandler(this, name, fields, importVar(extend), [for (i in interfaces) importVar(i)]));
 			case EImport(c, mode):
-				if (!importEnabled)
-					return null;
+				if (!importEnabled) return null;
 
 				if(mode == IAll) {
 					throw "TODO";
 					return null;
 				}
 
-				var splitClassName = c.split(".");
-				var origSplitClassName = splitClassName.copy();
-				var field = null;
-
-				var varName = switch(mode) {
+				var splitClassName:Array<String> = c.split(".");
+				if (splitClassName.length <= 0) return null;
+				
+				var varName:String = switch(mode) {
 					case IAs(name): name;
-					default: splitClassName[splitClassName.length - 1];
-				}
+					default: splitClassName[splitClassName.length-1];
+				};
 
-				if (variables.exists(varName)) // class is already imported
+				// Class is already imported
+				if (variables.exists(varName))
 					return null;
 
-				var cl = getClass(c);
+				// Orginal class
+				var importedClass = getClass(c);
+				if (importedClass != null) {
+					variables.set(varName, importedClass);
+					return importedClass;
+				}
 
 				// Allow for flixel.ui.FlxBar.FlxBarFillDirection;
-				if(cl == null && origSplitClassName.length > 1) {
-					splitClassName = origSplitClassName.copy();
+				var newClassName:Array<String> = splitClassName.copy();
+				newClassName.splice(-2, 1); // Remove the last last item
 
-					splitClassName.splice(-2, 1); // Remove the last last item
-					cl = getClass(splitClassName.join("."));
+				importedClass = getClass(newClassName.join("."));
+				if (importedClass != null) {
+					variables.set(varName, importedClass);
+					return importedClass;
 				}
 
 				// Allow for Std.isOfType;
-				if(cl == null && origSplitClassName.length > 1) {
-					splitClassName = origSplitClassName.copy();
+				var importField:String = null;
+				var newClassName:Array<String> = splitClassName.copy();
 
-					field = splitClassName.pop();
-					cl = getClass(splitClassName.join("."));
-				}
+				importField = newClassName.pop();
+				importedClass = getClass(newClassName.join("."));
 
-				if (cl != null) {
-					var value:Dynamic = switch(cl) {
+				// Import the .isOfType
+				if (importedClass != null) {
+					var classOrEnum:Dynamic = switch(importedClass) {
 						case Left(e): e;
 						case Right(e): Tools.getEnum(e);
-					}
-					if(field != null) { // import Std.isOfType;
+					};
+					if(importField != null) {
 						var v:Dynamic = null;
-						if(v == null)
-							v = UnsafeReflect.getProperty(value, field);
-						if(v == null)
-							v = UnsafeReflect.field(value, field);
+						if(v == null) v = UnsafeReflect.getProperty(classOrEnum, importField);
+						if(v == null) v = UnsafeReflect.field(classOrEnum, importField);
 
 						if(v == null)
-							error(EInvalidAccess(field, c));
-						value = v;
+							error(EInvalidAccess(importField, c));
+						classOrEnum = v;
 					}
-					variables.set(varName, value);
-					return value; // If someone wants to import and use the value as a return, they can
+
+					variables.set(newClassName[newClassName.length-1], importedClass);
+					variables.set(switch(mode) {
+						case IAs(name): name;
+						default: importField;
+					}, classOrEnum);
+					return classOrEnum;
 				}
 
 				if (importFailedCallback == null || !importFailedCallback(c.split("."))) // Incase of custom import
 					error(EInvalidClass(c));
 				return null;
-
 			case EConst(c):
 				switch (c) {
 					case CInt(v): return v;
diff --git a/tests/src/tests/FinalCase.hx b/tests/src/tests/FinalCase.hx
index 726475b..38b391d 100644
--- a/tests/src/tests/FinalCase.hx
+++ b/tests/src/tests/FinalCase.hx
@@ -112,6 +112,14 @@ class FinalCase extends TestCase {
 		Std.isOfType('', String);
 		", Std.isOfType('', String));
 
+		assertEq("
+		import Math.round as F;
+		F(2.9);", Math.round(2.9));
+
+		assertEq("
+		import Math.round;
+		round(1.9);", Math.round(1.9));
+
 		// Test EOF with preprocessor
 	}