Fix type checking around comprehensions with empty iteration range

PiperOrigin-RevId: 979965644

Author: l46kok

checker/internal/type_checker_impl.cc

@@ -838,7 +838,8 @@ void ResolveVisitor::PostVisitComprehensionSubexpression(
                            GetDeducedType(&comprehension.accu_init())));
       break;
     case ComprehensionArg::ITER_RANGE: {
-      Type range_type = GetDeducedType(&comprehension.iter_range());
+      Type range_type = inference_context_->FullySubstitute(
+          GetDeducedType(&comprehension.iter_range()), /*free_to_dyn=*/false);
       Type iter_type = DynType();   // iter_var for non comprehensions v2.
       Type iter_type1 = DynType();  // iter_var for comprehensions v2.
       Type iter_type2 = DynType();  // iter_var2 for comprehensions v2.
@@ -851,7 +852,14 @@ void ResolveVisitor::PostVisitComprehensionSubexpression(
           iter_type = iter_type1 = range_type.GetMap().key();
           iter_type2 = range_type.GetMap().value();
           break;
+        case TypeKind::kTypeParam:
+          // Set the range type to DYN to prevent assignment to a potentially
+          // incorrect type at a later point in type-checking. The IsAssignable
+          // call will update the type substitutions for the type param.
+          inference_context_->IsAssignable(DynType(), range_type);
+          break;
         case TypeKind::kDyn:
+        case TypeKind::kError:
           break;
         default:
           ReportIssue(TypeCheckIssue::CreateError(

checker/internal/type_checker_impl_test.cc

@@ -1287,6 +1287,76 @@ TEST(TypeCheckerImplTest, ComprehensionDynRange) {
   EXPECT_THAT(result.GetIssues(), IsEmpty());
 }
 
+TEST(TypeCheckerImplTest, EmptyListRangeNestedListComprehension) {
+  TypeCheckEnv env(GetSharedTestingDescriptorPool());
+  google::protobuf::Arena arena;
+  ASSERT_THAT(RegisterMinimalBuiltins(&arena, env), IsOk());
+
+  TypeCheckerImpl impl(std::move(env));
+  ASSERT_OK_AND_ASSIGN(auto ast, MakeTestParsedAst("[].map(x, x.map(y, y))"));
+  ASSERT_OK_AND_ASSIGN(ValidationResult result, impl.Check(std::move(ast)));
+
+  EXPECT_TRUE(result.IsValid()) << result.FormatError();
+  EXPECT_THAT(result.GetIssues(), IsEmpty());
+
+  ASSERT_OK_AND_ASSIGN(auto checked_ast, result.ReleaseAst());
+  EXPECT_EQ(checked_ast->GetReturnType(),
+            AstType(ListTypeSpec(std::make_unique<AstType>(
+                ListTypeSpec(std::make_unique<AstType>(DynTypeSpec()))))));
+}
+
+TEST(TypeCheckerImplTest, EmptyMapRangeNestedComprehension) {
+  TypeCheckEnv env(GetSharedTestingDescriptorPool());
+  google::protobuf::Arena arena;
+  ASSERT_THAT(RegisterMinimalBuiltins(&arena, env), IsOk());
+
+  TypeCheckerImpl impl(std::move(env));
+  ASSERT_OK_AND_ASSIGN(auto ast, MakeTestParsedAst("{}.map(k, k.map(y, y))"));
+  ASSERT_OK_AND_ASSIGN(ValidationResult result, impl.Check(std::move(ast)));
+
+  EXPECT_TRUE(result.IsValid()) << result.FormatError();
+  EXPECT_THAT(result.GetIssues(), IsEmpty());
+
+  ASSERT_OK_AND_ASSIGN(auto checked_ast, result.ReleaseAst());
+  EXPECT_EQ(checked_ast->GetReturnType(),
+            AstType(ListTypeSpec(std::make_unique<AstType>(
+                ListTypeSpec(std::make_unique<AstType>(DynTypeSpec()))))));
+}
+
+TEST(TypeCheckerImplTest, EmptyRangeNestedComprehensionPredicate) {
+  TypeCheckEnv env(GetSharedTestingDescriptorPool());
+  google::protobuf::Arena arena;
+  ASSERT_THAT(RegisterMinimalBuiltins(&arena, env), IsOk());
+
+  TypeCheckerImpl impl(std::move(env));
+  ASSERT_OK_AND_ASSIGN(auto ast,
+                       MakeTestParsedAst("[].all(x, x.all(y, y == 1))"));
+  ASSERT_OK_AND_ASSIGN(ValidationResult result, impl.Check(std::move(ast)));
+
+  EXPECT_TRUE(result.IsValid()) << result.FormatError();
+  EXPECT_THAT(result.GetIssues(), IsEmpty());
+
+  ASSERT_OK_AND_ASSIGN(auto checked_ast, result.ReleaseAst());
+  EXPECT_EQ(checked_ast->GetReturnType(), AstType(PrimitiveType::kBool));
+}
+
+TEST(TypeCheckerImplTest, EmptyRangeNestedComprehensionBeforeConstraint) {
+  TypeCheckEnv env(GetSharedTestingDescriptorPool());
+  google::protobuf::Arena arena;
+  ASSERT_THAT(RegisterMinimalBuiltins(&arena, env), IsOk());
+
+  TypeCheckerImpl impl(std::move(env));
+  ASSERT_OK_AND_ASSIGN(
+      auto ast, MakeTestParsedAst("[].all(x, x.all(y, y == 1) && x == 1)"));
+  ASSERT_OK_AND_ASSIGN(ValidationResult result, impl.Check(std::move(ast)));
+
+  EXPECT_TRUE(result.IsValid()) << result.FormatError();
+  EXPECT_THAT(result.GetIssues(), IsEmpty());
+
+  ASSERT_OK_AND_ASSIGN(auto checked_ast, result.ReleaseAst());
+  EXPECT_EQ(checked_ast->GetReturnType(), AstType(PrimitiveType::kBool));
+}
+
 TEST(TypeCheckerImplTest, BasicOvlResolution) {
   TypeCheckEnv env(GetSharedTestingDescriptorPool());
   google::protobuf::Arena arena;