From 1c4cd3432053ec5ae5d8890ad1fe315ae524644a Mon Sep 17 00:00:00 2001
From: Dan Klishch <danilklishch@gmail.com>
Date: Sat, 2 Sep 2023 14:33:35 -0400
Subject: [PATCH] JSSpecCompiler: Restrict usage of NodeSubtreePointer

This class stores a non-owning raw pointer to a member of `Node`, so
extra care is needed to ensure that referenced `Node`s will be alive
by the time `NodeSubtreePointer` is used. Since we only need to use this
class while traversing AST in `RecursiveASTVisitor`, access to class
methods can be restricted using `Badge<RecursiveASTVisitor>`.
---
 Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/AST/AST.h     | 5 +++--
 .../JSSpecCompiler/Compiler/GenericASTPass.cpp               | 4 ++--
 Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/Forward.h     | 3 +++
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/AST/AST.h b/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/AST/AST.h
index 27841c64ac..2ae3026af5 100644
--- a/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/AST/AST.h
+++ b/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/AST/AST.h
@@ -6,6 +6,7 @@
 
 #pragma once
 
+#include <AK/Badge.h>
 #include <AK/RefCounted.h>
 #include <AK/RefPtr.h>
 #include <AK/Vector.h>
@@ -27,9 +28,9 @@ public:
     {
     }
 
-    Tree& get() { return *m_tree_ptr; }
+    Tree& get(Badge<RecursiveASTVisitor>) { return *m_tree_ptr; }
 
-    void replace_subtree(Tree tree) { *m_tree_ptr = move(tree); }
+    void replace_subtree(Badge<RecursiveASTVisitor>, Tree tree) { *m_tree_ptr = move(tree); }
 
 private:
     Tree* m_tree_ptr;
diff --git a/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/Compiler/GenericASTPass.cpp b/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/Compiler/GenericASTPass.cpp
index aae5168488..f37dd23a6b 100644
--- a/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/Compiler/GenericASTPass.cpp
+++ b/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/Compiler/GenericASTPass.cpp
@@ -18,7 +18,7 @@ void RecursiveASTVisitor::run_in_subtree(Tree& tree)
 
 void RecursiveASTVisitor::replace_current_node_with(Tree tree)
 {
-    m_current_subtree_pointer->replace_subtree(move(tree));
+    m_current_subtree_pointer->replace_subtree({}, move(tree));
 }
 
 RecursionDecision RecursiveASTVisitor::recurse(Tree root, NodeSubtreePointer& pointer)
@@ -30,7 +30,7 @@ RecursionDecision RecursiveASTVisitor::recurse(Tree root, NodeSubtreePointer& po
 
     if (decision == RecursionDecision::Recurse) {
         for (auto& child : root->subtrees()) {
-            if (recurse(child.get(), child) == RecursionDecision::Break)
+            if (recurse(child.get({}), child) == RecursionDecision::Break)
                 return RecursionDecision::Break;
         }
     }
diff --git a/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/Forward.h b/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/Forward.h
index a6cd0f0e6b..7ffdb27aa7 100644
--- a/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/Forward.h
+++ b/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/Forward.h
@@ -37,6 +37,9 @@ class Variable;
 class FunctionPointer;
 using FunctionPointerRef = NonnullRefPtr<FunctionPointer>;
 
+// Compiler/GenericASTPass.h
+class RecursiveASTVisitor;
+
 // Parser/SpecParser.h
 class AlgorithmStep;
 class AlgorithmStepList;