Framework for copy-on-write visitors (WIP)

Signed-off-by: Chris Dodd <cdodd@nvidia.com>
p4lang · ChrisDodd · Nov 11, 2024 · Dec 5, 2024 · Dec 12, 2024 · 70531dac432eedfc826e323d65dd623d5d100d04
commit 70531dac432eedfc826e323d65dd623d5d100d04
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -24,8 +24,8 @@ endif()
 
 project (P4C)
 
-# set (CMAKE_CXX_EXTENSIONS OFF) # prefer using -std=c++17 rather than -std=gnu++17
-set (CMAKE_CXX_STANDARD 17)
+# set (CMAKE_CXX_EXTENSIONS OFF) # prefer using -std=c++20 rather than -std=gnu++20
+set (CMAKE_CXX_STANDARD 20)
 set (CMAKE_CXX_STANDARD_REQUIRED ON)
 
 set (CMAKE_MODULE_PATH ${PROJECT_SOURCE_DIR}/cmake)

diff --git a/frontends/p4/removeReturns.cpp b/frontends/p4/removeReturns.cpp
@@ -20,7 +20,7 @@ limitations under the License.
 
 namespace P4 {
 
-bool MoveToElseAfterBranch::preorder(IR::BlockStatement *block) {
+bool MoveToElseAfterBranch::preorder(IR::COWptr<IR::BlockStatement> block) {
     movedToIfBranch = false;
     for (auto it = block->components.begin(); it != block->components.end();) {
         if (movedToIfBranch)
@@ -32,7 +32,7 @@ bool MoveToElseAfterBranch::preorder(IR::BlockStatement *block) {
     return false;
 }
 
-bool MoveToElseAfterBranch::moveFromParentTo(const IR::Statement *&child) {
+template<class T> bool MoveToElseAfterBranch::moveFromParentTo(T &child) {
     auto parent = getParent<IR::BlockStatement>();
     size_t next = getContext()->child_index + 1;
     if (!parent || next >= parent->components.size()) {
@@ -43,9 +43,9 @@ bool MoveToElseAfterBranch::moveFromParentTo(const IR::Statement *&child) {
     IR::BlockStatement *modified = nullptr;
     if (!child)
         modified = new IR::BlockStatement;
-    else if (auto *t = child->to<IR::BlockStatement>())
+    else if (auto *t = child->template to<IR::BlockStatement>())
         modified = t->clone();
-    else if (child->is<IR::EmptyStatement>())
+    else if (child->template is<IR::EmptyStatement>())
         modified = new IR::BlockStatement;
     else
         modified = new IR::BlockStatement({child});
@@ -55,7 +55,7 @@ bool MoveToElseAfterBranch::moveFromParentTo(const IR::Statement *&child) {
     return true;
 }
 
-bool MoveToElseAfterBranch::preorder(IR::IfStatement *ifStmt) {
+bool MoveToElseAfterBranch::preorder(IR::COWptr<IR::IfStatement> ifStmt) {
     hasJumped = false;
     bool movedCode = false;
     visit(ifStmt->ifTrue, "ifTrue", 1);
@@ -74,11 +74,11 @@ bool MoveToElseAfterBranch::preorder(IR::IfStatement *ifStmt) {
     return false;
 }
 
-bool MoveToElseAfterBranch::preorder(IR::SwitchStatement *swch) {
+bool MoveToElseAfterBranch::preorder(IR::COWptr<IR::SwitchStatement> swch) {
     // TBD: if there is exactly one case that falls through (all others end with a branch)
     // then we could move subsequent code into that case, as it done with 'if'
     bool canFallThrough = false;
-    for (auto &c : swch->cases) {
+    for (auto c : swch->cases) {
         hasJumped = false;
         visit(c, "cases");
         canFallThrough |= !hasJumped;
@@ -87,7 +87,7 @@ bool MoveToElseAfterBranch::preorder(IR::SwitchStatement *swch) {
     return false;
 }
 
-void MoveToElseAfterBranch::postorder(IR::LoopStatement *) {
+void MoveToElseAfterBranch::postorder(IR::COWptr<IR::LoopStatement>) {
     // after a loop body is never unreachable
     hasJumped = false;
 }

diff --git a/frontends/p4/removeReturns.h b/frontends/p4/removeReturns.h
@@ -65,7 +65,7 @@ introduce a boolean flag and extra tests to remove those branches.
 precondition:
     switchAddDefault pass has run to ensure switch statements cover all cases
 */
-class MoveToElseAfterBranch : public Modifier {
+class MoveToElseAfterBranch : public COWModifier {
     /* This pass does not use (inherit from) ControlFlowVisitor, even though it is doing
      * control flow analysis, as it turns out to be more efficient to do it directly here
      * by overloading the branching constructs (if/switch/loops) and not cloning the visitor,
@@ -80,22 +80,22 @@ class MoveToElseAfterBranch : public Modifier {
      * indicating that it needs to be removed from the BlockStatment */
     bool movedToIfBranch = false;
 
-    bool preorder(IR::BlockStatement *) override;
-    bool moveFromParentTo(const IR::Statement *&child);
-    bool preorder(IR::IfStatement *) override;
-    bool preorder(IR::SwitchStatement *) override;
-    void postorder(IR::LoopStatement *) override;
+    bool preorder(IR::COWptr<IR::BlockStatement>) override;
+    template<class T> bool moveFromParentTo(T &child);
+    bool preorder(IR::COWptr<IR::IfStatement>) override;
+    bool preorder(IR::COWptr<IR::SwitchStatement>) override;
+    void postorder(IR::COWptr<IR::LoopStatement>) override;
     bool branch() {
         hasJumped = true;
         // no need to visit children
         return false;
     }
-    bool preorder(IR::BreakStatement *) override { return branch(); }
-    bool preorder(IR::ContinueStatement *) override { return branch(); }
-    bool preorder(IR::ExitStatement *) override { return branch(); }
-    bool preorder(IR::ReturnStatement *) override { return branch(); }
+    bool preorder(IR::COWptr<IR::BreakStatement>) override { return branch(); }
+    bool preorder(IR::COWptr<IR::ContinueStatement>) override { return branch(); }
+    bool preorder(IR::COWptr<IR::ExitStatement>) override { return branch(); }
+    bool preorder(IR::COWptr<IR::ReturnStatement>) override { return branch(); }
     // Only visit statements, skip all expressions
-    bool preorder(IR::Expression *) override { return false; }
+    bool preorder(IR::COWptr<IR::Expression>) override { return false; }
 
  public:
     MoveToElseAfterBranch() {}

diff --git a/ir/copy_on_write_inl.h b/ir/copy_on_write_inl.h
@@ -0,0 +1,238 @@
+/*
+Copyright 2024 NVIDIA CORPORATION.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+#ifndef IR_COPY_ON_WRITE_INL_H_
+#define IR_COPY_ON_WRITE_INL_H_
+
+/* template methods declared in "copy_on_write_ptr.h" that can't be defined there
+ * due to order-of-declaration issues. */
+
+namespace P4::IR {
+
+template<class T> const T *COWinfo<T>::get() const {
+    if (clone) return clone->checkedTo<T>();
+    return orig->checkedTo<T>();
+}
+
+template<class T> const T *COWinfo<T>::getOrig() const {
+    return orig->checkedTo<T>();
+}
+
+template<class T> T *COWinfo<T>::mkClone() {
+    if (!clone) clone = orig->clone();
+    return clone->checkedTo<T>();
+}
+
+template<class T> T *COWinfo<T>::getClone() const {
+    BUG_CHECK(clone != nullptr, "Not yet cloned in getClone");
+    return clone->checkedTo<T>();
+}
+
+/* specializations for IR::Node pointers */
+template <class T, typename U, const U *T::*field>
+requires std::derived_from<U, Node>
+struct COWfieldref<T, const U *, field> {
+    COWinfo<T>        *info;
+
+    const U *get() const { return info->get()->*field; }
+    operator const U *() const { return get(); }
+    const U *operator=(const U *val) const {
+        if (!info->isCloned() && info->get()->*field == val) return val;
+        return info->mkClone()->*field = val;
+    }
+    void set(const U *val) const { *this = val; }
+    template<typename V> requires (std::is_base_of_v<V, U> && !std::same_as<V, U>)
+    void set(const V *val) const { *this = val ? val->template checkedTo<U>() : (U *)nullptr; }
+    const U *operator->() const { return get(); }
+    const U &operator*() const { return *get(); }
+};
+
+template<class T, class C, class U, C T::*field>
+struct COW_element_ref {
+    COWinfo<T> *info;
+    bool is_const;
+    union {
+        typename C::const_iterator ci;
+        typename C::iterator ni;
+    };
+    COW_element_ref(COWinfo<T> *inf, typename C::const_iterator i)
+    : info(inf), is_const(true) { ci = i; }
+    COW_element_ref(COWinfo<T> *inf, typename C::iterator i)
+    : info(inf), is_const(false) { ni = i; }
+    void clone_fixup() {
+        if (is_const) {
+            // messy problem -- need to clone (iff not yet cloned) and then find the
+            // corresponding iterator in the clone
+            auto i = (info->mkClone()->*field).begin();
+            auto &orig_vec = info->getOrig()->*field;
+            for (auto oi = orig_vec.begin(); oi != ci; ++oi, ++i)
+                BUG_CHECK(oi != orig_vec.end(), "Invalid iterator in clone_fixup");
+            ni = i;
+            is_const = false;
+        }
+    }
+    U get() {
+        if (is_const && info->isCloned()) clone_fixup();
+        return *ci;
+    }
+    operator U() {
+        if (is_const && info->isCloned()) clone_fixup();
+        return *ci;
+    }
+    U operator=(U val) {
+        clone_fixup();
+        return *ni = val;
+    }
+    void set(U val) {
+        clone_fixup();
+        *ni = val;
+    }
+    void set(const Node *val) requires std::is_pointer_v<U>
+    {
+        set(val ? val->checkedTo<std::remove_pointer<U>::type>() : (U)nullptr);
+    }
+};
+
+template<class T, class C, class U, C T::*field>
+struct COW_iterator {
+    COW_element_ref<T, C, U, field> ref;
+    COW_iterator(COWinfo<T> *inf, typename C::const_iterator i) : ref(inf, i) {}
+    COW_iterator(COWinfo<T> *inf, typename C::iterator i) : ref(inf, i) {}
+    COW_iterator &operator++() { ++ref.ci; return *this; }
+    COW_iterator operator++(int) { COW_iterator<T, C, U, field> rv = *this; ++ref.ci; return rv; }
+    COW_iterator &operator--() { --ref.ci; return *this; }
+    COW_iterator operator--(int) { COW_iterator<T, C, U, field> rv = *this; --ref.ci; return rv; }
+    bool operator==(const COW_iterator<T, C, U, field> &i) const { return ref.ci == i.ref.ci; }
+    bool operator!=(const COW_iterator<T, C, U, field> &i) const { return ref.ci != i.ref.ci; }
+    COW_element_ref<T, C, U, field> &operator *() { return ref; }
+};
+
+
+/* specialization for safe_vector */
+template <class T, typename U, safe_vector<U> T::*field>
+struct COWfieldref<T, safe_vector<U>, field> {
+    COWinfo<T>        *info;
+
+    using iterator = COW_iterator<T, safe_vector<U>, U, field>;
+
+    void visit_children(Visitor &) { BUG("TBD"); }
+
+    const safe_vector<U> &get() const { return info->get()->*field; }
+    safe_vector<U> &modify() const { return info->mkClone()->*field; }
+    operator const safe_vector<U>&() const { return get(); }
+    safe_vector<U> &operator=(const safe_vector<U> &val) const { return modify() = val; }
+    safe_vector<U> &operator=(safe_vector<U> &&val) const { return modify() = std::move(val); }
+    iterator begin() {
+        if (info->isCloned())
+            return iterator(info, (info->getClone()->*field).begin());
+        else
+            return iterator(info, (info->get()->*field).begin());
+    }
+    iterator end() {
+        if (info->isCloned())
+            return iterator(info, (info->getClone()->*field).begin());
+        else
+            return iterator(info, (info->get()->*field).begin());
+    }
+    // FIXME need to add insert/appeand/prepend/emplace_back specializations
+};
+
+
+/* specializations for IR::Vector */
+template <class T, typename U, Vector<U> T::*field>
+struct COWfieldref<T, Vector<U>, field> {
+    COWinfo<T>        *info;
+
+    using iterator = COW_iterator<T, Vector<U>, const U *, field>;
+
+    void visit_children(Visitor &) { BUG("TBD"); }
+
+    const Vector<U> &get() const { return info->get()->*field; }
+    Vector<U> &modify() const { return info->mkClone()->*field; }
+    operator const Vector<U>&() const { return get(); }
+    Vector<U> &operator=(const Vector<U> &val) const { return modify() = val; }
+    Vector<U> &operator=(Vector<U> &&val) const { return modify() = std::move(val); }
+    iterator begin() {
+        if (info->isCloned())
+            return iterator(info, (info->getClone()->*field).begin());
+        else
+            return iterator(info, (info->get()->*field).begin());
+    }
+    iterator end() {
+        if (info->isCloned())
+            return iterator(info, (info->getClone()->*field).begin());
+        else
+            return iterator(info, (info->get()->*field).begin());
+    }
+    iterator erase(iterator i) {
+        i.ref.clone_fixup();
+        Vector<U> &vec = info->getClone()->*field;
+        return iterator(info, vec.erase(i.ref.ni));
+    }
+    // FIXME need to add insert/appeand/prepend/emplace_back specializations
+};
+
+/* specializations for IR::IndexedVector */
+template <class T, typename U, IndexedVector<U> T::*field>
+struct COWfieldref<T, IndexedVector<U>, field> {
+    COWinfo<T>        *info;
+
+    using iterator = COW_iterator<T, IndexedVector<U>, const U *, field>;
+
+    void visit_children(Visitor &) { BUG("TBD"); }
+
+    const IndexedVector<U> &get() const { return info->get()->*field; }
+    IndexedVector<U> &modify() const { return info->mkClone()->*field; }
+    operator const IndexedVector<U>&() const { return get(); }
+    IndexedVector<U> &operator=(const IndexedVector<U> &val) const { return modify() = val; }
+    IndexedVector<U> &operator=(IndexedVector<U> &&val) const { return modify() = std::move(val); }
+    iterator begin() {
+        if (info->isCloned())
+            return iterator(info, (info->getClone()->*field).begin());
+        else
+            return iterator(info, (info->get()->*field).begin());
+    }
+    iterator end() {
+        if (info->isCloned())
+            return iterator(info, (info->getClone()->*field).begin());
+        else
+            return iterator(info, (info->get()->*field).begin());
+    }
+    iterator erase(iterator i) {
+        i.ref.clone_fixup();
+        IndexedVector<U> &vec = info->getClone()->*field;
+        return iterator(info, vec.erase(i.ref.ni));
+    }
+    // FIXME need to add insert/appeand/prepend/emplace_back/removeByName specializations
+};
+
+/* specializations for IR::NameMap */
+template <class T, typename U,
+          template <class K, class V, class COMP, class ALLOC> class MAP,
+          NameMap<U, MAP> T::*field>
+struct COWfieldref<T, NameMap<U, MAP>, field> {
+    COWinfo<T>        *info;
+
+    using iterator = COW_iterator<T, NameMap<U, MAP>, const U *, field>;
+
+    void visit_children(Visitor &) { BUG("TBD"); }
+};
+
+// FIXME -- need NodeMap specializations if any backend ever uses that template
+
+}  // namespace P4::IR
+
+#endif /* IR_COPY_ON_WRITE_INL_H_ */