insertExtensionLabel{Before,After}: allow initial node to point at extension

This will be useful in insertCopiesUnder because the extension
is introduced below a given node rather than above a given node.
In a repeated application, the child of the node will point
to the extension node introduced before.

Author: Sven Verdoolaege

tc/core/polyhedral/schedule_transforms.cc

@@ -588,6 +588,10 @@ void insertExtensionLabelBefore(
   size_t pos;
   auto parent = tree->ancestor(root, 1);
   ScheduleTree* seqTree;
+  if (tree->elemAs<detail::ScheduleTreeElemExtension>()) {
+    tree = tree->child({0});
+    parent = tree;
+  }
   if (tree->elemAs<detail::ScheduleTreeElemSequence>()) {
     seqTree = tree;
     pos = 0;
@@ -610,6 +614,10 @@ void insertExtensionLabelAfter(
   size_t pos;
   auto parent = tree->ancestor(root, 1);
   ScheduleTree* seqTree;
+  if (tree->elemAs<detail::ScheduleTreeElemExtension>()) {
+    tree = tree->child({0});
+    parent = tree;
+  }
   if (tree->elemAs<detail::ScheduleTreeElemSequence>()) {
     seqTree = tree;
     pos = tree->numChildren();

tc/core/polyhedral/schedule_transforms.h

@@ -225,7 +225,8 @@ void insertExtensionLabelAt(
 
 // Insert a zero-dimensional extension statement with the given identifier
 // before node "tree".
-// If "tree" is a sequence node or a grandchild of a sequence node,
+// If "tree" is a sequence node, an extension node with a sequence child,
+// or a grandchild of a sequence node,
 // then the new statement is inserted in the right position
 // of that sequence node.
 // Otherwise, a new sequence node is inserted.
@@ -236,7 +237,8 @@ void insertExtensionLabelBefore(
 
 // Insert a zero-dimensional extension statement with the given identifier
 // after node "tree".
-// If "tree" is a sequence node or a grandchild of a sequence node,
+// If "tree" is a sequence node, an extension node with a sequence child,
+// or a grandchild of a sequence node,
 // then the new statement is inserted in the right position
 // of that sequence node.
 // Otherwise, a new sequence node is inserted.