Fix null-safety issues

SomeRanDev · Dec 1, 2023 · e1864a0 · e1864a0
1 parent 8f7d1e1
commit e1864a0
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 4 deletions.
diff --git a/src/reflaxe/ReflectCompiler.hx b/src/reflaxe/ReflectCompiler.hx
@@ -435,7 +435,9 @@ class ReflectCompiler {
 		if(compiler.options.normalizeEIE) {
 			final eiec = new EverythingIsExprSanitizer(data.expr, compiler);
 			data.setExpr(eiec.convertedExpr());
-			data.setVariableUsageCount(eiec.variableUsageCount);
+			if(eiec.variableUsageCount != null) {
+				data.setVariableUsageCount(eiec.variableUsageCount);
+			}
 		}
 		if(compiler.options.processAvoidTemporaries) {
 			final tvr = new TemporaryVarRemover(data.expr, data.getOrFindVariableUsageCount());

diff --git a/src/reflaxe/compiler/EverythingIsExprSanitizer.hx b/src/reflaxe/compiler/EverythingIsExprSanitizer.hx
@@ -139,12 +139,12 @@ class EverythingIsExprSanitizer {
 	/**
 		Search through parents to find valid `variableUsageCount`.
 	**/
-	function getVariableUsageCount() {
+	function getVariableUsageCount(): Map<Int, Int> {
 		var p = this;
 		while(p.parent != null) {
 			p = p.parent;
 		}
-		return p.variableUsageCount;
+		return p.variableUsageCount.trustMe(); // The top-most should always exist.
 	}
 
 	public function convertedExpr(): TypedExpr {

diff --git a/src/reflaxe/compiler/TemporaryVarRemover.hx b/src/reflaxe/compiler/TemporaryVarRemover.hx
@@ -126,7 +126,7 @@ class TemporaryVarRemover {
 	**/
 	function getVariableUsageCount(variableId: Int): Int {
 		return if(varUsageCount != null && varUsageCount.exists(variableId)) {
-			varUsageCount.get(variableId);
+			varUsageCount.get(variableId) ?? 0;
 		} else {
 			0;
 		}

diff --git a/src/reflaxe/data/ClassFuncData.hx b/src/reflaxe/data/ClassFuncData.hx
@@ -92,6 +92,10 @@ class ClassFuncData {
 		If it has not been calculated yet, it is calculated here.
 	**/
 	public function getOrFindVariableUsageCount(): Map<Int, Int> {
+		if(expr == null) {
+			return [];
+		}
+
 		final map: Map<Int, Int> = [];
 		function count(e: TypedExpr) {
 			switch(e.expr) {