fix type finder

TauMC · Oct 21, 2024 · 65e919b · 65e919b
1 parent fa9ddb5
commit 65e919b
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/src/main/java/org/taumc/glsl/TypeFinder.java b/src/main/java/org/taumc/glsl/TypeFinder.java
@@ -7,7 +7,7 @@
 
 import java.util.concurrent.atomic.AtomicInteger;
 
-public class TypeFinder extends GLSLParserBaseListener {
+public class TypeFinder extends GLSLCancelableBaseListener {
 
     private final String name;
     private final AtomicInteger type;
@@ -27,6 +27,20 @@ public void enterSingle_declaration(GLSLParser.Single_declarationContext ctx) {
                 if (ctx.fully_specified_type().type_specifier().type_specifier_nonarray().getChild(0) instanceof TerminalNode node) {
                     if (node.getSymbol() instanceof CommonToken t) {
                         type.set(t.getType());
+                        keepWalking = false;
+                    }
+                }
+            }
+        } else if (ctx.getParent() instanceof GLSLParser.Init_declarator_listContext listContext) {
+            for (var entry : listContext.typeless_declaration()) {
+                if (entry.IDENTIFIER().getSymbol() instanceof CommonToken cToken) {
+                    if (cToken.getText().equals(name)) {
+                        if (ctx.fully_specified_type().type_specifier().type_specifier_nonarray().getChild(0) instanceof TerminalNode node) {
+                            if (node.getSymbol() instanceof CommonToken t) {
+                                type.set(t.getType());
+                                keepWalking = false;
+                            }
+                        }
                     }
                 }
             }

diff --git a/src/main/java/org/taumc/glsl/Util.java b/src/main/java/org/taumc/glsl/Util.java
@@ -82,7 +82,7 @@ public static void removeVariable(GLSLParser.Translation_unitContext root, Strin
 
     public static int findType(GLSLParser.Translation_unitContext root, String code) {
         AtomicInteger type = new AtomicInteger();
-        ParseTreeWalker.DEFAULT.walk(new TypeFinder(code, type), root);
+        FastTreeWalker.walk(new TypeFinder(code, type), root);
         return type.get();
     }