From ddca2cf160a757cd2afc7b6c09aa8309ea72b189 Mon Sep 17 00:00:00 2001
From: TB Schardl <neboat@mit.edu>
Date: Tue, 12 Nov 2024 10:50:29 -0500
Subject: [PATCH] [TapirTaskInfo] Add first set of methods to update TaskInfo
 analysis on the fly.

---
 llvm/include/llvm/Analysis/TapirTaskInfo.h | 42 ++++++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/llvm/include/llvm/Analysis/TapirTaskInfo.h b/llvm/include/llvm/Analysis/TapirTaskInfo.h
index fef97f7faabd..6c95c70ee001 100644
--- a/llvm/include/llvm/Analysis/TapirTaskInfo.h
+++ b/llvm/include/llvm/Analysis/TapirTaskInfo.h
@@ -20,6 +20,7 @@
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instruction.h"
@@ -304,6 +305,17 @@ class Spindle {
     DenseBlockSet.insert(&B);
   }
 
+  /// Raw method to remove block B from this spindle.
+  void removeBlock(BasicBlock &B) {
+    DenseBlockSet.erase(&B);
+    for (auto Iter = Blocks.begin(); Iter != Blocks.end(); ++Iter) {
+      if (*Iter == &B) {
+        Blocks.erase(Iter);
+        break;
+      }
+    }
+  }
+
   // Returns true if the basic block B predeces this spindle.
   bool blockPrecedesSpindle(const BasicBlock *B) const {
     for (const BasicBlock *SB : successors(B))
@@ -1456,6 +1468,27 @@ class TaskInfo {
     SpindleMap[S] = T;
   }
 
+  // Move spindle S from its task to that task's parent.
+  //
+  // NOTE: The resulting TaskInfo may not exactly match the state of a freshly
+  // computed TaskInfo.
+  void moveSpindlesToParent(Task *T) {
+    // Currently this method simply adds T's spindles to T's parent task
+    // and removes those spindles from T.
+    // TODO: Update Spindle types and edges.
+    Task *Parent = T->getParentTask();
+    // Add all spindles to parent task.
+    for (Spindle *S : T->Spindles) {
+      Parent->addSpindle(*S);
+      S->setParentTask(Parent);
+      SpindleMap[S] = Parent;
+    }
+
+    // Remove all spindles from this task.
+    while (!T->Spindles.empty())
+      T->Spindles.pop_back();
+  }
+
   // Add spindle S to task T, where S is a shared exception-handling spindle
   // among subtasks of T.
   void addEHSpindleToTask(Spindle *S, Task *T) {
@@ -1472,6 +1505,15 @@ class TaskInfo {
     BBMap[&B] = S;
   }
 
+  // Remove basic block B from its spindle.
+  void removeBlock(BasicBlock &B) {
+    if (!BBMap.count(&B))
+      return;
+    Spindle *S = BBMap[&B];
+    S->removeBlock(B);
+    BBMap.erase(&B);
+  }
+
   // Associate a task T with the spindle TFSpindle that creates its taskframe.
   void AssociateTaskFrameWithUser(Task *T, Spindle *TFSpindle) {
     TFSpindle->TaskFrameUser = T;