CloneContext: Assert objects are owned by the program

Check that pre-clone objects are owned by the source program.
Check that post-clone object are owned by the target builder.

Fixed: tint:469
Change-Id: Idd0eeb8dfb386e295b66b4b6621cc13dc1a30786
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/48260
Commit-Queue: Ben Clayton <bclayton@chromium.org>
Reviewed-by: Antonio Maiorano <amaiorano@google.com>
Reviewed-by: James Price <jrprice@google.com>
diff --git a/src/ast/node.h b/src/ast/node.h
index fe74af8..18eee89 100644
--- a/src/ast/node.h
+++ b/src/ast/node.h
@@ -81,7 +81,7 @@
 
 /// @param node a pointer to an AST node
 /// @returns the ProgramID of the given AST node.
-inline ProgramID ProgramIDOf(ast::Node* node) {
+inline ProgramID ProgramIDOf(const ast::Node* node) {
   return node ? node->program_id() : ProgramID();
 }
 
diff --git a/src/clone_context.h b/src/clone_context.h
index 55af18f..16f8fa1 100644
--- a/src/clone_context.h
+++ b/src/clone_context.h
@@ -32,11 +32,14 @@
 class CloneContext;
 class Program;
 class ProgramBuilder;
-
 namespace ast {
 class FunctionList;
+class Node;
 }  // namespace ast
 
+ProgramID ProgramIDOf(const Program*);
+ProgramID ProgramIDOf(const ast::Node*);
+
 /// Cloneable is the base class for all objects that can be cloned
 class Cloneable : public Castable<Cloneable> {
  public:
@@ -46,6 +49,11 @@
   virtual Cloneable* Clone(CloneContext* ctx) const = 0;
 };
 
+/// @returns an invalid ProgramID
+inline ProgramID ProgramIDOf(const Cloneable*) {
+  return ProgramID();
+}
+
 /// ShareableCloneable is the base class for Cloneable objects which will only
 /// be cloned once when CloneContext::Clone() is called with the same object
 /// pointer.
@@ -95,6 +103,8 @@
       return nullptr;
     }
 
+    TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(src, a);
+
     // Have we cloned this object already, or was Replace() called for this
     // object?
     auto it = cloned_.find(a);
@@ -127,7 +137,11 @@
       cloned_.emplace(a, cloned);
     }
 
-    return CheckedCast<T>(cloned);
+    auto* out = CheckedCast<T>(cloned);
+
+    TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(dst, out);
+
+    return out;
   }
 
   /// Clones the Node or type::Type `a` into the ProgramBuilder #dst if `a` is
@@ -149,6 +163,8 @@
       return nullptr;
     }
 
+    TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(src, a);
+
     // Have we seen this object before? If so, return the previously cloned
     // version instead of making yet another copy.
     auto it = cloned_.find(a);
diff --git a/src/clone_context_test.cc b/src/clone_context_test.cc
index 7337de6..a038f6b 100644
--- a/src/clone_context_test.cc
+++ b/src/clone_context_test.cc
@@ -97,6 +97,23 @@
   }
 };
 
+struct ProgramNode : public Castable<ProgramNode, Cloneable> {
+  ProgramNode(Allocator* alloc, ProgramID id, ProgramID cloned_id)
+      : allocator(alloc), program_id(id), cloned_program_id(cloned_id) {}
+
+  Allocator* const allocator;
+  ProgramID const program_id;
+  ProgramID const cloned_program_id;
+
+  ProgramNode* Clone(CloneContext*) const override {
+    return allocator->Create<ProgramNode>(cloned_program_id, cloned_program_id);
+  }
+};
+
+ProgramID ProgramIDOf(const ProgramNode* node) {
+  return node->program_id;
+}
+
 struct UniqueTypes {
   using Node = UniqueNode;
   using Replaceable = UniqueReplaceable;
@@ -642,6 +659,38 @@
   EXPECT_EQ(cloned.Symbols().NameFor(new_c), "c");
 }
 
+TYPED_TEST(CloneContextTest, ProgramIDs) {
+  ProgramBuilder dst;
+  Program src(ProgramBuilder{});
+  CloneContext ctx(&dst, &src);
+  Allocator allocator;
+  ctx.Clone(allocator.Create<ProgramNode>(src.ID(), dst.ID()));
+}
+
+TYPED_TEST(CloneContextTest, ProgramIDs_ObjectNotOwnedBySrc) {
+  EXPECT_FATAL_FAILURE(
+      {
+        ProgramBuilder dst;
+        Program src(ProgramBuilder{});
+        CloneContext ctx(&dst, &src);
+        Allocator allocator;
+        ctx.Clone(allocator.Create<ProgramNode>(ProgramID::New(), dst.ID()));
+      },
+      R"(internal compiler error: TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(src, a))");
+}
+
+TYPED_TEST(CloneContextTest, ProgramIDs_ObjectNotOwnedByDst) {
+  EXPECT_FATAL_FAILURE(
+      {
+        ProgramBuilder dst;
+        Program src(ProgramBuilder{});
+        CloneContext ctx(&dst, &src);
+        Allocator allocator;
+        ctx.Clone(allocator.Create<ProgramNode>(src.ID(), ProgramID::New()));
+      },
+      R"(internal compiler error: TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(dst, out))");
+}
+
 }  // namespace
 
 TINT_INSTANTIATE_TYPEINFO(UniqueNode);
@@ -651,5 +700,6 @@
 TINT_INSTANTIATE_TYPEINFO(ShareableReplaceable);
 TINT_INSTANTIATE_TYPEINFO(ShareableReplacement);
 TINT_INSTANTIATE_TYPEINFO(NotANode);
+TINT_INSTANTIATE_TYPEINFO(ProgramNode);
 
 }  // namespace tint
diff --git a/src/program.h b/src/program.h
index f2f6eaf..815e751 100644
--- a/src/program.h
+++ b/src/program.h
@@ -171,6 +171,12 @@
   bool moved_ = false;
 };
 
+/// @param program the Program
+/// @returns the ProgramID of the Program
+inline ProgramID ProgramIDOf(const Program* program) {
+  return program->ID();
+}
+
 }  // namespace tint
 
 #endif  // SRC_PROGRAM_H_
diff --git a/src/program_builder.h b/src/program_builder.h
index 5125419..5691c01 100644
--- a/src/program_builder.h
+++ b/src/program_builder.h
@@ -1505,6 +1505,12 @@
 };
 //! @endcond
 
+/// @param builder the ProgramBuilder
+/// @returns the ProgramID of the ProgramBuilder
+inline ProgramID ProgramIDOf(const ProgramBuilder* builder) {
+  return builder->ID();
+}
+
 }  // namespace tint
 
 #endif  // SRC_PROGRAM_BUILDER_H_
diff --git a/src/type/type.h b/src/type/type.h
index 39bc710..95ed371 100644
--- a/src/type/type.h
+++ b/src/type/type.h
@@ -102,6 +102,13 @@
   Type();
 };
 
+/// @returns the ProgramID of the given type.
+inline ProgramID ProgramIDOf(const Type*) {
+  /// TODO(crbug.com/tint/724): Actually implement this once we split the `type`
+  /// namespace into ast::Type and sem::Type.
+  return ProgramID();
+}
+
 }  // namespace type
 }  // namespace tint