Properly unify result classes with list columns

drift_dev/lib/src/analysis/custom_result_class.dart

@@ -1,3 +1,5 @@
+import 'dart:collection';
+
 import 'results/results.dart';
 
 /// Transforms queries given in [inputs] so that their result sets respect
@@ -66,38 +68,71 @@ Map<SqlSelectQuery, SqlSelectQuery> transformCustomResultClasses(
     }
 
     final referenceResult = queries.first.resultSet;
-    final dartNames = {
-      for (final column in referenceResult.columns)
-        column: referenceResult.dartNameFor(column),
-    };
 
-    var isFirst = true;
     for (final query in queries) {
-      final newResultSet = InferredResultSet(
-        null,
-        query.resultSet.columns,
-        resultClassName: resultSetName,
-        // Only generate a result class for the first query in the group
-        dontGenerateResultClass: !isFirst,
-      );
-
-      // Make sure compatible columns in the two result sets have the same
-      // Dart name.
-      newResultSet.forceDartNames({
-        for (final entry in dartNames.entries)
-          newResultSet.columns.singleWhere((e) => e.isCompatibleTo(entry.key)):
-              entry.value,
-      });
-
+      final newResultSet =
+          _makeResultSetsCompatible(query.resultSet, referenceResult);
       final newQuery = query.replaceResultSet(newResultSet);
       replacements[query] = newQuery;
-      isFirst = false;
     }
   }
 
   return replacements;
 }
 
+InferredResultSet _makeResultSetsCompatible(
+    InferredResultSet target, InferredResultSet reference) {
+  var columns = target.columns;
+  Map<ResultColumn, String>? newNames;
+
+  if (target != reference) {
+    // Make sure the result sets agree on the Dart column names to use.
+    final remainingColumns = LinkedHashSet.of(target.columns);
+    newNames = <ResultColumn, String>{};
+    columns = [];
+
+    for (final column in reference.columns) {
+      var columnFromThisResultSet =
+          remainingColumns.firstWhere((e) => e.isCompatibleTo(column));
+      remainingColumns.remove(columnFromThisResultSet);
+
+      newNames[columnFromThisResultSet] = reference.dartNameFor(column);
+
+      // For list columns, we need to apply the same unification to the
+      // result set of this column as well.
+      if (columnFromThisResultSet is NestedResultQuery) {
+        final nested = columnFromThisResultSet.query.resultSet;
+        if (nested.needsOwnClass) {
+          final transformed = _makeResultSetsCompatible(
+            nested,
+            (column as NestedResultQuery).query.resultSet,
+          );
+
+          columnFromThisResultSet = NestedResultQuery(
+            from: columnFromThisResultSet.from,
+            query: columnFromThisResultSet.query.replaceResultSet(transformed),
+          );
+        }
+      }
+
+      columns.add(columnFromThisResultSet);
+    }
+  }
+
+  final newResultSet = InferredResultSet(
+    null,
+    columns,
+    resultClassName: reference.resultClassName,
+    // Only generate a result class for the first query in the group
+    dontGenerateResultClass: target != reference,
+  );
+  if (newNames != null) {
+    newResultSet.forceDartNames(newNames);
+  }
+
+  return newResultSet;
+}
+
 bool _resultSetsCompatible(Iterable<InferredResultSet> resultSets) {
   InferredResultSet? last;
 

drift_dev/test/analysis/resolver/queries/custom_class_name_test.dart

@@ -115,4 +115,36 @@ getTitlesWithGroupOther AS GroupWithTitles: SELECT group_name, LIST(SELECT title
       )),
     }, results.dartOutputs, results);
   });
+
+  test('supports query with two list columns', () async {
+    var queries = await analyzeQueries('''
+CREATE TABLE books (
+  id INT NOT NULL PRIMARY KEY,
+  title TEXT NULL,
+  group_name TEXT NULL
+);
+
+a AS TitleList: SELECT LIST(SELECT title FROM books), LIST(SELECT id, title FROM books) FROM books;
+b AS TitleList: SELECT LIST(SELECT title FROM books), LIST(SELECT id, title FROM books) FROM books;
+    ''');
+
+    final errors = <String>[];
+    queries = transformCustomResultClasses(queries, errors.add).values;
+    expect(errors, isEmpty);
+
+    final a = queries.first;
+    final b = queries.last;
+
+    expect(a.name, 'a');
+    expect(b.name, 'b');
+    expect(a.resultClassName, 'TitleList');
+    expect(b.resultClassName, 'TitleList');
+    expect(a.resultSet?.dontGenerateResultClass, isFalse);
+    expect(b.resultSet?.dontGenerateResultClass, isTrue);
+
+    final nestedA = a.resultSet?.nestedResults.last as NestedResultQuery;
+    final nestedB = b.resultSet?.nestedResults.last as NestedResultQuery;
+    expect(nestedA.query.resultSet.needsOwnClass, isTrue);
+    expect(nestedB.query.resultSet.needsOwnClass, isFalse);
+  });
 }