From 9cac8d27e5aa0bea54824254258a28c3875f539a Mon Sep 17 00:00:00 2001
From: Mingun <alexander_sergey@mail.ru>
Date: Thu, 26 Sep 2024 23:34:18 +0500
Subject: [PATCH] Fix enum resolution in `enum` keys

This commit also unify enum resolution in expression language and in `enum` keys

Fixes https://github.com/kaitai-io/kaitai_struct/issues/1028
---
 .../struct/precompile/ResolveTypes.scala      | 89 +++++--------------
 1 file changed, 23 insertions(+), 66 deletions(-)

diff --git a/shared/src/main/scala/io/kaitai/struct/precompile/ResolveTypes.scala b/shared/src/main/scala/io/kaitai/struct/precompile/ResolveTypes.scala
index 4fb8e27e0..9111ebd2a 100644
--- a/shared/src/main/scala/io/kaitai/struct/precompile/ResolveTypes.scala
+++ b/shared/src/main/scala/io/kaitai/struct/precompile/ResolveTypes.scala
@@ -1,8 +1,9 @@
 package io.kaitai.struct.precompile
 
-import io.kaitai.struct.Log
+import io.kaitai.struct.{ClassTypeProvider, Log}
 import io.kaitai.struct.datatype.DataType
 import io.kaitai.struct.datatype.DataType.{ArrayType, EnumType, SwitchType, UserType}
+import io.kaitai.struct.exprlang.Ast
 import io.kaitai.struct.format._
 import io.kaitai.struct.problems._
 
@@ -50,11 +51,27 @@ class ResolveTypes(specs: ClassSpecs, topClass: ClassSpec, opaqueTypes: Boolean)
         ut.classSpec = resClassSpec
         problems
       case et: EnumType =>
-        et.enumSpec = resolveEnumSpec(curClass, et.name)
-        if (et.enumSpec.isEmpty) {
-          Some(EnumNotFoundErr(et.name, curClass, path ++ List("enum")))
-        } else {
-          None
+        et.name match {
+          case typePath :+ name =>
+            try {
+              val resolver = new ClassTypeProvider(specs, curClass)
+              val ty = resolver.resolveEnum(Ast.typeId(false, typePath), name)
+              Log.enumResolve.info(() => s"    => ${ty.nameAsStr}")
+              et.enumSpec = Some(ty)
+              None
+            } catch {
+              case ex: TypeNotFoundError =>
+                Log.typeResolve.info(() => s"    => ??? (while resolving enum '${et.name}'): $ex")
+                Log.enumResolve.info(() => s"    => ??? (enclosing type not found, enum '${et.name}'): $ex")
+                Some(TypeNotFoundErr(typePath, curClass, path :+ "enum"))
+              case ex: EnumNotFoundError =>
+                Log.enumResolve.info(() => s"    => ??? (enum '${et.name}'): $ex")
+                Some(EnumNotFoundErr(et.name, curClass, path :+ "enum"))
+            }
+          case _ =>
+            Log.enumResolve.info(() => s"    => ??? (enum '${et.name}' without name)")
+            // TODO: Maybe more specific error about empty name?
+            Some(EnumNotFoundErr(et.name, curClass, path :+ "enum"))
         }
       case st: SwitchType =>
         st.cases.flatMap { case (caseName, ut) =>
@@ -135,64 +152,4 @@ class ResolveTypes(specs: ClassSpecs, topClass: ClassSpec, opaqueTypes: Boolean)
         }
     }
   }
-
-  private def resolveEnumSpec(curClass: ClassSpec, typeName: List[String]): Option[EnumSpec] = {
-    Log.enumResolve.info(() => s"resolveEnumSpec: at ${curClass.name} doing ${typeName.mkString("|")}")
-
-    val res = realResolveEnumSpec(curClass, typeName)
-    res match {
-      case None => {
-        Log.enumResolve.info(() => s"    => ???")
-        res
-      }
-      case Some(x) => {
-        Log.enumResolve.info(() => s"    => ${x.nameAsStr}")
-        res
-      }
-    }
-  }
-
-  private def realResolveEnumSpec(curClass: ClassSpec, typeName: List[String]): Option[EnumSpec] = {
-    // First, try to do it in current class
-
-    // If we're seeking composite name, we only have to resolve the very first
-    // part of it at this stage
-    val firstName :: restNames = typeName
-
-    val resolvedHere = if (restNames.isEmpty) {
-      curClass.enums.get(firstName)
-    } else {
-      curClass.types.get(firstName).flatMap((nestedClass) =>
-        resolveEnumSpec(nestedClass, restNames)
-      )
-    }
-
-    resolvedHere match {
-      case Some(_) => resolvedHere
-      case None =>
-        // No luck resolving here, let's try upper levels, if they exist
-        curClass.upClass match {
-          case Some(upClass) =>
-            resolveEnumSpec(upClass, typeName)
-          case None =>
-            // Check this class if it's top-level class
-            if (curClass.name.head == firstName) {
-              resolveEnumSpec(curClass, restNames)
-            } else {
-              // Check if top-level specs has this name
-              // If there's None => no luck at all
-              val resolvedTop = specs.get(firstName)
-              resolvedTop match {
-                case None => None
-                case Some(classSpec) => if (restNames.isEmpty) {
-                  // resolved everything, but this points to a type name, not enum name
-                  None
-                } else {
-                  resolveEnumSpec(classSpec, restNames)
-                }
-              }
-            }
-        }
-    }
-  }
 }