transform/InlinePtrLets: Fix ICE for lets in for-loops

For loop initializers and continuing statements do not have a BlockStatement as their parent.
Handle removal of these statements with a new Transform::RemoveStatement() helper

Fixed: tint:990
Change-Id: I24e7b18dcf71d3ef0a4d3ee68b9f68518e0eb5e8
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/58063
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: James Price <jrprice@google.com>

src/clone_context.h

@@ -513,12 +513,15 @@ class CloneContext {
   /// Reports an internal compiler error if the cast failed.
   template <typename TO, typename FROM>
   TO* CheckedCast(FROM* obj) {
-    if (TO* cast = As<TO>(obj)) {
+    if (obj == nullptr) {
+      return nullptr;
+    }
+    if (TO* cast = obj->template As<TO>()) {
       return cast;
     }
     TINT_ICE(Clone, Diagnostics())
         << "Cloned object was not of the expected type\n"
-        << "got:      " << (obj ? obj->TypeInfo().name : "<null>") << "\n"
+        << "got:      " << obj->TypeInfo().name << "\n"
         << "expected: " << TypeInfo::Of<TO>().name;
     return nullptr;
   }

src/transform/inline_pointer_lets.cc

@@ -177,7 +177,7 @@ void InlinePointerLets::Run(CloneContext& ctx, const DataMap&, DataMap&) {
       ptr_lets.emplace(let->variable(), std::move(ptr_let));
       // As the original `let` declaration will be fully inlined, there's no
       // need for the original declaration to exist. Remove it.
-      ctx.Remove(block->statements(), let);
+      RemoveStatement(ctx, let);
     }
   }
 

src/transform/transform.cc

@@ -19,6 +19,8 @@
 
 #include "src/program_builder.h"
 #include "src/sem/atomic_type.h"
+#include "src/sem/block_statement.h"
+#include "src/sem/for_loop_statement.h"
 #include "src/sem/reference_type.h"
 
 TINT_INSTANTIATE_TYPEINFO(tint::transform::Transform);
@@ -84,6 +86,21 @@ ast::DecorationList Transform::RemoveDecorations(
   return new_decorations;
 }
 
+void Transform::RemoveStatement(CloneContext& ctx, ast::Statement* stmt) {
+  auto* sem = ctx.src->Sem().Get(stmt);
+  if (auto* block = tint::As<sem::BlockStatement>(sem->Parent())) {
+    ctx.Remove(block->Declaration()->statements(), stmt);
+    return;
+  }
+  if (tint::Is<sem::ForLoopStatement>(sem->Parent())) {
+    ctx.Replace(stmt, static_cast<ast::Expression*>(nullptr));
+    return;
+  }
+  TINT_ICE(Transform, ctx.dst->Diagnostics())
+      << "unable to remove statement from parent of type "
+      << sem->TypeInfo().name;
+}
+
 ast::Type* Transform::CreateASTTypeFor(CloneContext& ctx, const sem::Type* ty) {
   if (ty->Is<sem::Void>()) {
     return ctx.dst->create<ast::Void>();

src/transform/transform.h

@@ -195,6 +195,13 @@ class Transform : public Castable<Transform> {
       const ast::DecorationList& in,
       std::function<bool(const ast::Decoration*)> should_remove);
 
+  /// Removes the statement `stmt` from the transformed program.
+  /// RemoveStatement handles edge cases, like statements in the initializer and
+  /// continuing of for-loops.
+  /// @param ctx the clone context
+  /// @param stmt the statement to remove when the program is cloned
+  static void RemoveStatement(CloneContext& ctx, ast::Statement* stmt);
+
   /// CreateASTTypeFor constructs new ast::Type nodes that reconstructs the
   /// semantic type `ty`.
   /// @param ctx the clone context