larsiusprime · lemz1 · Feb 17, 2025
diff --git a/polymod/Polymod.hx b/polymod/Polymod.hx
@@ -655,6 +655,7 @@ class Polymod
 		#if hscript
 		@:privateAccess
 		polymod.hscript._internal.PolymodScriptClass.clearScriptedClasses();
+		polymod.hscript._internal.PolymodEnum.clearScriptedEnums();
 		polymod.hscript.HScriptable.ScriptRunner.clearScripts();
 		#else
 		Polymod.warning(SCRIPT_HSCRIPT_NOT_INSTALLED, "Cannot register script classes, HScript is not available.");
@@ -691,6 +692,8 @@ class Polymod
 					polymod.hscript._internal.PolymodScriptClass.registerScriptClassByPath(path);
 				}
 			}
+
+			polymod.hscript._internal.PolymodInterpEx.validateImports(); 
 		}
 		#else
 		Polymod.warning(SCRIPT_HSCRIPT_NOT_INSTALLED, "Cannot register script classes, HScript is not available.");
@@ -728,6 +731,9 @@ class Polymod
 				if (future != null) futures.push(future);
 			}
 		}
+
+		polymod.hscript._internal.PolymodInterpEx.validateImports();
+
 		return futures;
 		#else
 		Polymod.warning(SCRIPT_HSCRIPT_NOT_INSTALLED, "Cannot register script classes, HScript is not available.");
@@ -1324,6 +1330,13 @@ enum abstract PolymodErrorCode(String) from String to String
 	 */
 	var SCRIPT_CLASS_MODULE_BLACKLISTED:String = 'script_class_module_blacklisted';
 
+	/**
+	 * You attempted to register a new enum with a name that is already in use.
+	 * - Rename the enum to one that is unique and will not conflict with other enums.
+	 * - If you need to clear all existing enum descriptors, call `Polymod.clearScripts()`.
+	 */
+	 var SCRIPT_ENUM_ALREADY_REGISTERED:String = 'script_enum_already_registered';
+
 	/**
 	 * One or more scripts are about to be parsed.
 	 * - This is an info message. You can log it or ignore it if you like.

diff --git a/polymod/hscript/_internal/PolymodClassDeclEx.hx b/polymod/hscript/_internal/PolymodClassDeclEx.hx
@@ -14,6 +14,7 @@ typedef PolymodClassDeclEx =
 	 * Save performance and improve sandboxing by resolving imports at interpretation time.
 	 */
 	@:optional var imports:Map<String, PolymodClassImport>;
+	@:optional var importsToValidate:Map<String, PolymodClassImport>;
 	@:optional var pkg:Array<String>;
 
 	@:optional var staticFields:Array<FieldDecl>;

diff --git a/polymod/hscript/_internal/PolymodEnum.hx b/polymod/hscript/_internal/PolymodEnum.hx
@@ -0,0 +1,66 @@
+package polymod.hscript._internal;
+
+#if hscript
+import hscript.Expr;
+
+@:access(hscript.Interp)
+@:allow(polymod.Polymod)
+class PolymodEnum 
+{
+  private static final scriptInterp = new PolymodInterpEx(null, null);
+
+  private var _e:PolymodEnumDeclEx;
+
+  private var _value:String;
+
+  private var _args:Array<Dynamic>;
+
+  public function new(e:PolymodEnumDeclEx, value:String, args:Array<Dynamic>) 
+  {
+    this._e = e;
+
+    var field = getField(value);
+
+    if (field == null) 
+    {
+      Polymod.error(SCRIPT_PARSE_ERROR, '${e.name}.${value} does not exist.');
+      return;
+    }
+
+    this._value = value;
+
+    if (args.length != field.args.length) 
+    {
+      Polymod.error(SCRIPT_PARSE_ERROR, '${e.name}.${value} got the wrong number of arguments.');
+      return;
+    }
+
+    this._args = args;
+  }
+
+  public static function clearScriptedEnums():Void 
+  {
+    scriptInterp.clearScriptEnumDescriptors();
+  }
+
+  private function getField(name:String):Null<EnumFieldDecl> 
+  {
+    for (field in _e.fields) 
+      {
+        if (field.name == name) 
+        {
+          return field;
+        }
+      }
+      return null;
+  }
+
+  public function toString():String 
+  {
+    var result:String = '${_e.name}.${_value}';
+    if(_args.length > 0)
+      result += '(${_args.join(',')})';
+    return result;
+  }
+}
+#end
diff --git a/polymod/hscript/_internal/PolymodEnumDeclEx.hx b/polymod/hscript/_internal/PolymodEnumDeclEx.hx
@@ -0,0 +1,13 @@
+package polymod.hscript._internal;
+
+#if hscript
+import hscript.Expr;
+
+typedef PolymodEnumDeclEx = 
+{
+    > EnumDecl,
+
+    @:optional var pkg:Array<String>;
+}
+
+#end
diff --git a/polymod/hscript/_internal/PolymodInterpEx.hx b/polymod/hscript/_internal/PolymodInterpEx.hx
@@ -14,6 +14,7 @@ import polymod.hscript._internal.PolymodClassDeclEx.PolymodStaticClassReference;
  */
 @:access(polymod.hscript._internal.PolymodScriptClass)
 @:access(polymod.hscript._internal.PolymodAbstractScriptClass)
+@:access(polymod.hscript._internal.PolymodEnum)
 class PolymodInterpEx extends Interp
 {
 	var targetCls:Class<Dynamic>;
@@ -194,6 +195,53 @@ class PolymodInterpEx extends Interp
 		return _scriptClassDescriptors.get(name);
 	}
 
+	private static var _scriptEnumDescriptors:Map<String, PolymodEnumDeclEx> = new Map<String, PolymodEnumDeclEx>();
+
+	private static function registerScriptEnum(e:PolymodEnumDeclEx)
+	{
+		var name = e.name;
+		if (e.pkg != null)
+		{
+			name = e.pkg.join(".") + "." + name;
+		}
+
+		if (_scriptEnumDescriptors.exists(name)) {
+			Polymod.error(SCRIPT_ENUM_ALREADY_REGISTERED, 'An enum with the fully qualified name "$name" has already been defined. Please change the enum name to ensure a unique name.');
+			return;
+		} else {
+			Polymod.debug('Registering enum $name');
+			_scriptEnumDescriptors.set(name, e);
+		}
+	}
+
+	public function clearScriptEnumDescriptors():Void {
+		// Clear the script enum descriptors.
+		_scriptEnumDescriptors.clear();
+
+		// Also destroy local variable scope.
+		this.resetVariables();
+	}
+
+	public static function validateImports():Void 
+	{
+		for (cls in _scriptClassDescriptors) 
+		{
+			var clsPath = cls.pkg != null ? (cls.pkg.join(".") + ".") : "";
+			clsPath += cls.name;
+
+			for (key => imp in cls.importsToValidate) 
+			{
+				if (_scriptEnumDescriptors.exists(imp.fullPath))
+				{
+					cls.imports.set(key, imp);
+					continue;
+				}
+
+				Polymod.error(SCRIPT_CLASS_MODULE_NOT_FOUND, 'Could not import ${imp.fullPath}', clsPath);
+			}
+		}
+	}
+
 	override function setVar(id:String, v:Dynamic)
 	{
 		if (_proxy != null && _proxy.superClass != null)
@@ -463,6 +511,89 @@ class PolymodInterpEx extends Interp
 				// If there is a try/catch block, the error will be caught.
 				// If there is no try/catch block, the error will be reported.
 				errorEx(EScriptThrow(str));
+
+			// Enums
+			case EField(e,f):
+				var name = getIdent(e);
+				name = getClassDecl().imports.get(name)?.fullPath ?? name;
+				if (name != null && _scriptEnumDescriptors.exists(name))
+				{
+					return new PolymodEnum(_scriptEnumDescriptors.get(name), f, []);
+				}
+			case ECall(e,params):
+				var args = new Array();
+				for (p in params)
+					args.push(expr(p));
+
+				switch(Tools.expr(e)) {
+					case EField(e,f):
+						var name = getIdent(e);
+						name = getClassDecl().imports.get(name)?.fullPath ?? name;
+						if (name != null && _scriptEnumDescriptors.exists(name))
+						{
+							return new PolymodEnum(_scriptEnumDescriptors.get(name), f, args);
+						}
+					default:
+				}
+				case ESwitch(e, cases, def):
+					var val:Dynamic = expr(e);
+
+					if (Std.isOfType(val, PolymodEnum))
+					{
+						var old:Int = declared.length;
+						var match = false;
+						for(c in cases) 
+						{
+							for(v in c.values) 
+							{
+								switch (Tools.expr(v)) 
+								{
+									case ECall(e, params):
+										switch (Tools.expr(e)) 
+										{
+											case EField(_, f):
+												if (val._value == f) 
+												{
+													for (i => p in params) 
+													{
+														switch (Tools.expr(p)) 
+														{
+															case EIdent(n):
+																declared.push({
+																	n: n,
+																	old: {r: locals.get(n)}
+																});
+																locals.set(n, {r: val._args[i]});
+															default:
+														}
+													}
+													match = true;
+													break;
+												}
+											default:
+										}
+									case EField(_, f):
+										if (val._value == f) 
+										{
+											match = true;
+											break;
+										}
+									default:
+								}
+							}
+							if(match) 
+							{
+								val = expr(c.expr);
+								break;
+							}
+						}
+						if (!match)
+						{
+							val = def == null ? null : expr(def);
+						}
+						restore(old);
+						return val;
+					}
 			default:
 				// Do nothing.
 		}
@@ -601,6 +732,21 @@ class PolymodInterpEx extends Interp
 		return a;
 	}
 
+	function getIdent(e:Expr):Null<String> {
+		#if hscriptPos
+		switch (e.e)
+		{
+		#else
+		switch (e)
+		{
+		#end
+			case EIdent(v):
+				return v;
+			default:
+				return null;
+		}
+	}
+
 	override function makeIterator(v:Dynamic):Iterator<Dynamic>
 	{
 		if (v.iterator != null)
@@ -1255,6 +1401,7 @@ class PolymodInterpEx extends Interp
 	{
 		var pkg:Array<String> = null;
 		var imports:Map<String, PolymodClassImport> = [];
+		var importsToValidate:Map<String, PolymodClassImport> = [];
 
 		for (importPath in PolymodScriptClass.defaultImports.keys())
 		{
@@ -1306,6 +1453,8 @@ class PolymodInterpEx extends Interp
 						importedClass.cls = PolymodScriptClass.abstractClassImpls.get(importedClass.fullPath);
 						trace('RESOLVED ABSTRACT CLASS ${importedClass.fullPath} -> ${Type.getClassName(importedClass.cls)}');
 						trace(Type.getClassFields(importedClass.cls));
+					} else if (_scriptEnumDescriptors.exists(importedClass.fullPath)) {
+						// do nothing
 					} else {
 						var resultCls:Class<Dynamic> = Type.resolveClass(importedClass.fullPath);
 
@@ -1316,7 +1465,9 @@ class PolymodInterpEx extends Interp
 
 						// If the class is still not found, skip this import entirely.
 						if (resultCls == null && resultEnm == null) {
-							Polymod.error(SCRIPT_CLASS_MODULE_NOT_FOUND, 'Could not import class ${importedClass.fullPath}', origin);
+							//Polymod.error(SCRIPT_CLASS_MODULE_NOT_FOUND, 'Could not import class ${importedClass.fullPath}', origin);
+							// this could be a scripted class or enum that hasn't been registered yet
+							importsToValidate.set(importedClass.name, importedClass);
 							continue;
 						} else if (resultCls != null) {
 							importedClass.cls = resultCls;
@@ -1372,6 +1523,7 @@ class PolymodInterpEx extends Interp
 
 					var classDecl:PolymodClassDeclEx = {
 						imports: imports,
+						importsToValidate: importsToValidate,
 						pkg: pkg,
 						name: c.name,
 						params: c.params,
@@ -1384,6 +1536,25 @@ class PolymodInterpEx extends Interp
 						staticFields: staticFields,
 					};
 					registerScriptClass(classDecl);
+				case DEnum(e):
+					if (pkg != null) 
+					{
+						imports.set(e.name, {
+							name: e.name,
+							pkg: pkg,
+							fullPath: pkg.join(".") + "." + e.name,
+							cls: null,
+							enm: null,
+						});
+					}
+
+					var enumDecl:PolymodEnumDeclEx = {
+						pkg: pkg,
+						name: e.name,
+						fields: e.fields,
+					};
+
+					registerScriptEnum(enumDecl);
 				case DTypedef(_):
 			}
 		}