fix: for-of on class array filter/sort/slice results now correctly resolves element type (#402)

Co-authored-by: cs01 <cs01@users.noreply.github.com>

Author: cs01

src/codegen/infrastructure/type-inference.ts

@@ -1735,6 +1735,33 @@ export class TypeInference {
 
   getObjectArrayElementType(expr: Expression): string | null {
     const e = expr as ExprBase;
+    if (e.type === "variable") {
+      const varName = (expr as VariableNode).name;
+      const rawType = this.st.getRawInterfaceType(varName);
+      if (rawType && this.getClass(rawType)) return rawType;
+      const resolved = this.resolveExpressionType(expr);
+      if (resolved && resolved.arrayDepth > 0 && this.getClass(resolved.base)) {
+        return resolved.base;
+      }
+      return null;
+    }
+    if (e.type === "array") {
+      const arrayExpr = expr as ArrayNode;
+      const elements = arrayExpr.elements || [];
+      if (elements.length > 0) {
+        const firstElem = elements[0] as ExprBase;
+        if (firstElem.type === "new") {
+          const newExpr = elements[0] as NewNode;
+          const resolvedClassName = this.resolveClassAlias(newExpr.className);
+          if (this.getClass(resolvedClassName)) return resolvedClassName;
+        }
+        const elemResolved = this.resolveExpressionType(elements[0]);
+        if (elemResolved && elemResolved.arrayDepth === 0 && this.getClass(elemResolved.base)) {
+          return elemResolved.base;
+        }
+      }
+      return null;
+    }
     if (e.type === "binary") {
       const binExpr = expr as BinaryNode;
       if (binExpr.op === "||" || binExpr.op === "??") {
@@ -1744,6 +1771,11 @@ export class TypeInference {
     if (e.type === "method_call") {
       const methodExpr = expr as MethodCallNode;
       const methodObjBase = methodExpr.object as ExprBase;
+      const m = methodExpr.method;
+      if (m === "filter" || m === "sort" || m === "reverse" || m === "slice") {
+        const sourceElemType = this.getObjectArrayElementType(methodExpr.object);
+        if (sourceElemType) return sourceElemType;
+      }
       if (methodObjBase.type === "this") {
         const className = this.ctx.getCurrentClassName();
         if (className) {

src/codegen/llvm-generator.ts

@@ -2290,6 +2290,10 @@ export class LLVMGenerator extends BaseGenerator implements IGeneratorContext {
           if (!elementType && resolvedBase && this.isKnownClass(resolvedBase)) {
             elementType = resolvedBase;
           }
+          if (!elementType && stmt.value) {
+            const objArrElemType = this.typeInference.getObjectArrayElementType(stmt.value);
+            if (objArrElemType) elementType = objArrElemType;
+          }
           llvmType = "%ObjectArray*";
           kind = SymbolKind.ObjectArray;
           defaultValue = "null";

src/codegen/statements/control-flow.ts

@@ -1070,6 +1070,16 @@ export class ControlFlowGenerator {
     }
     if (iterBase.type === "method_call") {
       const mcNode = iterable as MethodCallNode;
+      const mcMethod = mcNode.method;
+      if (
+        mcMethod === "filter" ||
+        mcMethod === "sort" ||
+        mcMethod === "reverse" ||
+        mcMethod === "slice"
+      ) {
+        const sourceElemType = this.getIterableClassElementType(mcNode.object);
+        if (sourceElemType) return sourceElemType;
+      }
       const objBase = mcNode.object as ExprBase;
       if (objBase.type === "variable") {
         const objName = (mcNode.object as VariableNode).name;

tests/fixtures/arrays/class-array-for-of-filter.ts

@@ -0,0 +1,30 @@
+class Task {
+  name: string;
+  done: boolean;
+  constructor(n: string, d: boolean) {
+    this.name = n;
+    this.done = d;
+  }
+}
+
+const tasks = [new Task("a", true), new Task("b", false), new Task("c", true)];
+const done = tasks.filter((t: Task): boolean => t.done);
+
+const names: string[] = [];
+for (const t of done) {
+  names.push(t.name);
+}
+console.log(names.join(","));
+
+const sorted = tasks.sort((a: Task, b: Task): number => {
+  if (a.name < b.name) return -1;
+  if (a.name > b.name) return 1;
+  return 0;
+});
+const sortedNames: string[] = [];
+for (const t of sorted) {
+  sortedNames.push(t.name);
+}
+console.log(sortedNames.join(","));
+
+console.log("TEST_PASSED");